com.google.datastore.v1.RunQueryRequest Java Examples

/**
 * Cloud Datastore system tables with statistics are periodically updated. This method fetches
 * the latest timestamp (in microseconds) of statistics update using the {@code __Stat_Total__}
 * table.
 */
private static long queryLatestStatisticsTimestamp(
    Datastore datastore, @Nullable String namespace) throws DatastoreException {
  Query.Builder query = Query.newBuilder();
  // Note: namespace either being null or empty represents the default namespace, in which
  // case we treat it as not provided by the user.
  if (Strings.isNullOrEmpty(namespace)) {
    query.addKindBuilder().setName("__Stat_Total__");
  } else {
    query.addKindBuilder().setName("__Stat_Ns_Total__");
  }
  query.addOrder(makeOrder("timestamp", DESCENDING));
  query.setLimit(Int32Value.newBuilder().setValue(1));
  RunQueryRequest request = makeRequest(query.build(), namespace);

  RunQueryResponse response = datastore.runQuery(request);
  QueryResultBatch batch = response.getBatch();
  if (batch.getEntityResultsCount() == 0) {
    throw new NoSuchElementException("Datastore total statistics unavailable");
  }
  Entity entity = batch.getEntityResults(0).getEntity();
  return entity.getProperties().get("timestamp").getTimestampValue().getSeconds() * 1000000;
}
 

/**
 * Gets a list of split keys given a desired number of splits.
 *
 * <p>This list will contain multiple split keys for each split. Only a single split key
 * will be chosen as the split point, however providing multiple keys allows for more uniform
 * sharding.
 *
 * @param numSplits the number of desired splits.
 * @param query the user query.
 * @param partition the partition to run the query in.
 * @param datastore the datastore containing the data.
 * @throws DatastoreException if there was an error when executing the datastore query.
 */
private List<Key> getScatterKeys(
    int numSplits, Query query, PartitionId partition, Datastore datastore)
    throws DatastoreException {
  Query.Builder scatterPointQuery = createScatterQuery(query, numSplits);

  List<Key> keySplits = new ArrayList<Key>();

  QueryResultBatch batch;
  do {
    RunQueryRequest scatterRequest =
        RunQueryRequest.newBuilder()
            .setPartitionId(partition)
            .setQuery(scatterPointQuery)
            .build();
    batch = datastore.runQuery(scatterRequest).getBatch();
    for (EntityResult result : batch.getEntityResultsList()) {
      keySplits.add(result.getEntity().getKey());
    }
    scatterPointQuery.setStartCursor(batch.getEndCursor());
    scatterPointQuery.getLimitBuilder().setValue(
        scatterPointQuery.getLimit().getValue() - batch.getEntityResultsCount());
  } while (batch.getMoreResults() == MoreResultsType.NOT_FINISHED);
  Collections.sort(keySplits, DatastoreHelper.getKeyComparator());
  return keySplits;
}
 

/**
 * Run a query on the datastore.
 *
 * @return The entities returned by the query.
 * @throws DatastoreException on error
 */
private List<Entity> runQuery(Query query) throws DatastoreException {
  RunQueryRequest.Builder request = RunQueryRequest.newBuilder();
  request.setQuery(query);
  RunQueryResponse response = datastore.runQuery(request.build());

  if (response.getBatch().getMoreResults() == QueryResultBatch.MoreResultsType.NOT_FINISHED) {
    System.err.println("WARNING: partial results\n");
  }
  List<EntityResult> results = response.getBatch().getEntityResultsList();
  List<Entity> entities = new ArrayList<Entity>(results.size());
  for (EntityResult result : results) {
    entities.add(result.getEntity());
  }
  return entities;
}
 

private RunQueryResponse runQueryWithRetries(RunQueryRequest request) throws Exception {
  Sleeper sleeper = Sleeper.DEFAULT;
  BackOff backoff = RUNQUERY_BACKOFF.backoff();
  while (true) {
    try {
      RunQueryResponse response = datastore.runQuery(request);
      rpcSuccesses.inc();
      return response;
    } catch (DatastoreException exception) {
      rpcErrors.inc();

      if (NON_RETRYABLE_ERRORS.contains(exception.getCode())) {
        throw exception;
      }
      if (!BackOffUtils.next(sleeper, backoff)) {
        LOG.error("Aborting after {} retries.", MAX_RETRIES);
        throw exception;
      }
    }
  }
}
 

@Override
public ListenableFuture<QueryResult> executeAsync(final Query statement, final ListenableFuture<TransactionResult> txn) {
  final ListenableFuture<Response> httpResponse = Futures.transformAsync(txn, result -> {
    final String namespace = config.getNamespace();
    final RunQueryRequest.Builder request = RunQueryRequest.newBuilder()
      .setQuery(statement.getPb(namespace != null ? namespace : ""));
    if (namespace != null) {
      request.setPartitionId(PartitionId.newBuilder().setNamespaceId(namespace));
    }
    final ByteString transaction = result.getTransaction();
    if (transaction != null) {
      request.setReadOptions(ReadOptions.newBuilder().setTransaction(transaction));
    }
    final ProtoHttpContent payload = new ProtoHttpContent(request.build());
    return ListenableFutureAdapter.asGuavaFuture(prepareRequest("runQuery", payload).execute());
  }, MoreExecutors.directExecutor());
  return Futures.transformAsync(httpResponse, response -> {
    if (!isSuccessful(response.getStatusCode())) {
      throw new DatastoreException(response.getStatusCode(), response.getResponseBody());
    }
    final RunQueryResponse query = RunQueryResponse.parseFrom(streamResponse(response));
    return Futures.immediateFuture(QueryResult.build(query));
  }, MoreExecutors.directExecutor());
}
 

/**
 * Tests {@link DatastoreV1.Read#getEstimatedSizeBytes} to fetch and return estimated size for a
 * query.
 */
@Test
public void testEstimatedSizeBytes() throws Exception {
  long entityBytes = 100L;
  // In seconds
  long timestamp = 1234L;

  RunQueryRequest latestTimestampRequest =
      makeRequest(makeLatestTimestampQuery(NAMESPACE), NAMESPACE);
  RunQueryResponse latestTimestampResponse = makeLatestTimestampResponse(timestamp);
  // Per Kind statistics request and response
  RunQueryRequest statRequest = makeRequest(makeStatKindQuery(NAMESPACE, timestamp), NAMESPACE);
  RunQueryResponse statResponse = makeStatKindResponse(entityBytes);

  when(mockDatastore.runQuery(latestTimestampRequest)).thenReturn(latestTimestampResponse);
  when(mockDatastore.runQuery(statRequest)).thenReturn(statResponse);

  assertEquals(entityBytes, getEstimatedSizeBytes(mockDatastore, QUERY, NAMESPACE));
  verify(mockDatastore, times(1)).runQuery(latestTimestampRequest);
  verify(mockDatastore, times(1)).runQuery(statRequest);
}
 

private Iterator<EntityResult> getIteratorAndMoveCursor() throws DatastoreException {
  Query.Builder query = this.query.toBuilder();
  query.setLimit(Int32Value.newBuilder().setValue(QUERY_BATCH_LIMIT));
  if (currentBatch != null && !currentBatch.getEndCursor().isEmpty()) {
    query.setStartCursor(currentBatch.getEndCursor());
  }

  RunQueryRequest request = makeRequest(query.build(), namespace);
  RunQueryResponse response = datastore.runQuery(request);

  currentBatch = response.getBatch();

  int numFetch = currentBatch.getEntityResultsCount();
  // All indications from the API are that there are/may be more results.
  moreResults =
      (numFetch == QUERY_BATCH_LIMIT) || (currentBatch.getMoreResults() == NOT_FINISHED);

  // May receive a batch of 0 results if the number of records is a multiple
  // of the request limit.
  if (numFetch == 0) {
    return null;
  }

  return currentBatch.getEntityResultsList().iterator();
}
 

/** Tests that {@link ReadFn} retries after an error. */
@Test
public void testReadFnRetriesErrors() throws Exception {
  // An empty query to read entities.
  Query query = Query.newBuilder().setLimit(Int32Value.newBuilder().setValue(1)).build();

  // Use mockResponseForQuery to generate results.
  when(mockDatastore.runQuery(any(RunQueryRequest.class)))
      .thenThrow(new DatastoreException("RunQuery", Code.DEADLINE_EXCEEDED, "", null))
      .thenAnswer(
          invocationOnMock -> {
            Query q = ((RunQueryRequest) invocationOnMock.getArguments()[0]).getQuery();
            return mockResponseForQuery(q);
          });

  ReadFn readFn = new ReadFn(V_1_OPTIONS, mockDatastoreFactory);
  DoFnTester<Query, Entity> doFnTester = DoFnTester.of(readFn);
  doFnTester.setCloningBehavior(CloningBehavior.DO_NOT_CLONE);
  doFnTester.processBundle(query);
}
 

@Test
public void testTranslateGqlQueryWithLimit() throws Exception {
  String gql = "SELECT * from DummyKind LIMIT 10";
  String gqlWithZeroLimit = gql + " LIMIT 0";
  GqlQuery gqlQuery = GqlQuery.newBuilder().setQueryString(gql).setAllowLiterals(true).build();
  GqlQuery gqlQueryWithZeroLimit =
      GqlQuery.newBuilder().setQueryString(gqlWithZeroLimit).setAllowLiterals(true).build();
  RunQueryRequest gqlRequest = makeRequest(gqlQuery, V_1_OPTIONS.getNamespace());
  RunQueryRequest gqlRequestWithZeroLimit =
      makeRequest(gqlQueryWithZeroLimit, V_1_OPTIONS.getNamespace());
  when(mockDatastore.runQuery(gqlRequestWithZeroLimit))
      .thenThrow(
          new DatastoreException(
              "runQuery",
              Code.INVALID_ARGUMENT,
              "invalid query",
              // dummy
              new RuntimeException()));
  when(mockDatastore.runQuery(gqlRequest))
      .thenReturn(RunQueryResponse.newBuilder().setQuery(QUERY).build());
  assertEquals(
      translateGqlQueryWithLimitCheck(gql, mockDatastore, V_1_OPTIONS.getNamespace()), QUERY);
  verify(mockDatastore, times(1)).runQuery(gqlRequest);
  verify(mockDatastore, times(1)).runQuery(gqlRequestWithZeroLimit);
}
 

@Test
public void runQuery() throws Exception {
  RunQueryRequest.Builder request = RunQueryRequest.newBuilder();
  request.getQueryBuilder();
  RunQueryResponse.Builder response = RunQueryResponse.newBuilder();
  response.getBatchBuilder()
      .setEntityResultType(EntityResult.ResultType.FULL)
      .setMoreResults(QueryResultBatch.MoreResultsType.NOT_FINISHED);
  expectRpc("runQuery", request.build(), response.build());
}
 

public RunQueryResponse runQuery(RunQueryRequest request) throws DatastoreException {
  try (InputStream is = remoteRpc.call("runQuery", request)) {
    return RunQueryResponse.parseFrom(is);
  } catch (IOException exception) {
    throw invalidResponseException("runQuery", exception);
  }
}
 

/** Build a datastore query request. */
private static RunQueryRequest makeRequest(Query query, @Nullable String namespace) {
  RunQueryRequest.Builder requestBuilder = RunQueryRequest.newBuilder().setQuery(query);
  if (namespace != null) {
    requestBuilder.getPartitionIdBuilder().setNamespaceId(namespace);
  }
  return requestBuilder.build();
}
 

/** Helper function to run a test reading from a {@link ReadFn}. */
private void readFnTest(int numEntities) throws Exception {
  // An empty query to read entities.
  Query query =
      Query.newBuilder().setLimit(Int32Value.newBuilder().setValue(numEntities)).build();

  // Use mockResponseForQuery to generate results.
  when(mockDatastore.runQuery(any(RunQueryRequest.class)))
      .thenAnswer(
          invocationOnMock -> {
            Query q = ((RunQueryRequest) invocationOnMock.getArguments()[0]).getQuery();
            return mockResponseForQuery(q);
          });

  ReadFn readFn = new ReadFn(V_1_OPTIONS, mockDatastoreFactory);
  DoFnTester<Query, Entity> doFnTester = DoFnTester.of(readFn);
  /**
   * Although Datastore client is marked transient in {@link ReadFn}, when injected through mock
   * factory using a when clause for unit testing purposes, it is not serializable because it
   * doesn't have a no-arg constructor. Thus disabling the cloning to prevent the test object from
   * being serialized.
   */
  doFnTester.setCloningBehavior(CloningBehavior.DO_NOT_CLONE);
  List<Entity> entities = doFnTester.processBundle(query);

  int expectedNumCallsToRunQuery = (int) Math.ceil((double) numEntities / QUERY_BATCH_LIMIT);
  verify(mockDatastore, times(expectedNumCallsToRunQuery)).runQuery(any(RunQueryRequest.class));
  // Validate the number of results.
  assertEquals(numEntities, entities.size());
}
 

@Test
public void testTranslateGqlQueryWithNoLimit() throws Exception {
  String gql = "SELECT * from DummyKind";
  String gqlWithZeroLimit = gql + " LIMIT 0";
  GqlQuery gqlQueryWithZeroLimit =
      GqlQuery.newBuilder().setQueryString(gqlWithZeroLimit).setAllowLiterals(true).build();
  RunQueryRequest gqlRequestWithZeroLimit =
      makeRequest(gqlQueryWithZeroLimit, V_1_OPTIONS.getNamespace());
  when(mockDatastore.runQuery(gqlRequestWithZeroLimit))
      .thenReturn(RunQueryResponse.newBuilder().setQuery(QUERY).build());
  assertEquals(
      translateGqlQueryWithLimitCheck(gql, mockDatastore, V_1_OPTIONS.getNamespace()), QUERY);
  verify(mockDatastore, times(1)).runQuery(gqlRequestWithZeroLimit);
}
 

/** Tests {@link SplitQueryFn} when no query splits is specified. */
@Test
public void testSplitQueryFnWithoutNumSplits() throws Exception {
  // Force SplitQueryFn to compute the number of query splits
  int numSplits = 0;
  int expectedNumSplits = 20;
  long entityBytes = expectedNumSplits * DEFAULT_BUNDLE_SIZE_BYTES;
  // In seconds
  long timestamp = 1234L;

  RunQueryRequest latestTimestampRequest =
      makeRequest(makeLatestTimestampQuery(NAMESPACE), NAMESPACE);
  RunQueryResponse latestTimestampResponse = makeLatestTimestampResponse(timestamp);

  // Per Kind statistics request and response
  RunQueryRequest statRequest = makeRequest(makeStatKindQuery(NAMESPACE, timestamp), NAMESPACE);
  RunQueryResponse statResponse = makeStatKindResponse(entityBytes);

  when(mockDatastore.runQuery(latestTimestampRequest)).thenReturn(latestTimestampResponse);
  when(mockDatastore.runQuery(statRequest)).thenReturn(statResponse);
  when(mockQuerySplitter.getSplits(
          eq(QUERY), any(PartitionId.class), eq(expectedNumSplits), any(Datastore.class)))
      .thenReturn(splitQuery(QUERY, expectedNumSplits));

  SplitQueryFn splitQueryFn = new SplitQueryFn(V_1_OPTIONS, numSplits, mockDatastoreFactory);
  DoFnTester<Query, Query> doFnTester = DoFnTester.of(splitQueryFn);
  doFnTester.setCloningBehavior(CloningBehavior.DO_NOT_CLONE);
  List<Query> queries = doFnTester.processBundle(QUERY);

  assertEquals(expectedNumSplits, queries.size());
  verify(mockQuerySplitter, times(1))
      .getSplits(eq(QUERY), any(PartitionId.class), eq(expectedNumSplits), any(Datastore.class));
  verify(mockDatastore, times(1)).runQuery(latestTimestampRequest);
  verify(mockDatastore, times(1)).runQuery(statRequest);
}
 

/** Translates a gql query string to {@link Query}. */
private static Query translateGqlQuery(String gql, Datastore datastore, String namespace)
    throws DatastoreException {
  GqlQuery gqlQuery = GqlQuery.newBuilder().setQueryString(gql).setAllowLiterals(true).build();
  RunQueryRequest req = makeRequest(gqlQuery, namespace);
  return datastore.runQuery(req).getQuery();
}
 

@VisibleForTesting
/** Builds a {@link RunQueryRequest} from the {@code GqlQuery} and {@code namespace}. */
static RunQueryRequest makeRequest(GqlQuery gqlQuery, @Nullable String namespace) {
  return RunQueryRequest.newBuilder()
      .setGqlQuery(gqlQuery)
      .setPartitionId(forNamespace(namespace))
      .build();
}
 

/** Builds a {@link RunQueryRequest} from the {@code query} and {@code namespace}. */
static RunQueryRequest makeRequest(Query query, @Nullable String namespace) {
  return RunQueryRequest.newBuilder()
      .setQuery(query)
      .setPartitionId(forNamespace(namespace))
      .build();
}
 

/** Retrieve latest table statistics for a given kind, namespace, and datastore. */
private static Entity getLatestTableStats(
    String ourKind, @Nullable String namespace, Datastore datastore) throws DatastoreException {
  long latestTimestamp = queryLatestStatisticsTimestamp(datastore, namespace);
  LOG.info("Latest stats timestamp for kind {} is {}", ourKind, latestTimestamp);

  Query.Builder queryBuilder = Query.newBuilder();
  if (Strings.isNullOrEmpty(namespace)) {
    queryBuilder.addKindBuilder().setName("__Stat_Kind__");
  } else {
    queryBuilder.addKindBuilder().setName("__Stat_Ns_Kind__");
  }

  queryBuilder.setFilter(
      makeAndFilter(
          makeFilter("kind_name", EQUAL, makeValue(ourKind).build()).build(),
          makeFilter("timestamp", EQUAL, makeValue(latestTimestamp).build()).build()));

  RunQueryRequest request = makeRequest(queryBuilder.build(), namespace);

  long now = System.currentTimeMillis();
  RunQueryResponse response = datastore.runQuery(request);
  LOG.debug("Query for per-kind statistics took {}ms", System.currentTimeMillis() - now);

  QueryResultBatch batch = response.getBatch();
  if (batch.getEntityResultsCount() == 0) {
    throw new NoSuchElementException(
        "Datastore statistics for kind " + ourKind + " unavailable");
  }
  return batch.getEntityResults(0).getEntity();
}
 

/** Read and output entities for the given query. */
@ProcessElement
public void processElement(ProcessContext context) throws Exception {
  Query query = context.element();
  String namespace = options.getNamespace();
  int userLimit = query.hasLimit() ? query.getLimit().getValue() : Integer.MAX_VALUE;

  boolean moreResults = true;
  QueryResultBatch currentBatch = null;

  while (moreResults) {
    Query.Builder queryBuilder = query.toBuilder();
    queryBuilder.setLimit(
        Int32Value.newBuilder().setValue(Math.min(userLimit, QUERY_BATCH_LIMIT)));

    if (currentBatch != null && !currentBatch.getEndCursor().isEmpty()) {
      queryBuilder.setStartCursor(currentBatch.getEndCursor());
    }

    RunQueryRequest request = makeRequest(queryBuilder.build(), namespace);
    RunQueryResponse response = runQueryWithRetries(request);

    currentBatch = response.getBatch();

    // MORE_RESULTS_AFTER_LIMIT is not implemented yet:
    // https://groups.google.com/forum/#!topic/gcd-discuss/iNs6M1jA2Vw, so
    // use result count to determine if more results might exist.
    int numFetch = currentBatch.getEntityResultsCount();
    if (query.hasLimit()) {
      verify(
          userLimit >= numFetch,
          "Expected userLimit %s >= numFetch %s, because query limit %s must be <= userLimit",
          userLimit,
          numFetch,
          query.getLimit());
      userLimit -= numFetch;
    }

    // output all the entities from the current batch.
    for (EntityResult entityResult : currentBatch.getEntityResultsList()) {
      context.output(entityResult.getEntity());
    }

    // Check if we have more entities to be read.
    moreResults =
        // User-limit does not exist (so userLimit == MAX_VALUE) and/or has not been satisfied
        (userLimit > 0)
            // All indications from the API are that there are/may be more results.
            && ((numFetch == QUERY_BATCH_LIMIT)
                || (currentBatch.getMoreResults() == NOT_FINISHED));
  }
}