com.google.datastore.v1.RunQueryResponse 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;
}
 

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;
      }
    }
  }
}
 

/**
 * 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);
}
 

@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);
}
 

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();
}
 

@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());
}
 

/**
 * 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;
}
 

/** 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();
}
 

/** 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);
}
 

@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);
}
 

/**
 * A helper function that creates mock {@link Entity} results in response to a query. Always
 * indicates that more results are available, unless the batch is limited to fewer than {@link
 * DatastoreV1.Read#QUERY_BATCH_LIMIT} results.
 */
private static RunQueryResponse mockResponseForQuery(Query q) {
  // Every query DatastoreV1 sends should have a limit.
  assertTrue(q.hasLimit());

  // The limit should be in the range [1, QUERY_BATCH_LIMIT]
  int limit = q.getLimit().getValue();
  assertThat(limit, greaterThanOrEqualTo(1));
  assertThat(limit, lessThanOrEqualTo(QUERY_BATCH_LIMIT));

  // Create the requested number of entities.
  List<EntityResult> entities = new ArrayList<>(limit);
  for (int i = 0; i < limit; ++i) {
    entities.add(
        EntityResult.newBuilder()
            .setEntity(Entity.newBuilder().setKey(makeKey("key" + i, i + 1)))
            .build());
  }

  // Fill out the other parameters on the returned result batch.
  RunQueryResponse.Builder ret = RunQueryResponse.newBuilder();
  ret.getBatchBuilder()
      .addAllEntityResults(entities)
      .setEntityResultType(EntityResult.ResultType.FULL)
      .setMoreResults(
          limit == QUERY_BATCH_LIMIT
              ? QueryResultBatch.MoreResultsType.NOT_FINISHED
              : QueryResultBatch.MoreResultsType.NO_MORE_RESULTS);

  return ret.build();
}
 

/** Builds a per-kind statistics response with the given entity size. */
private static RunQueryResponse makeStatKindResponse(long entitySizeInBytes) {
  RunQueryResponse.Builder statKindResponse = RunQueryResponse.newBuilder();
  Entity.Builder entity = Entity.newBuilder();
  entity.setKey(makeKey("dummyKind", "dummyId"));
  entity.putProperties("entity_bytes", makeValue(entitySizeInBytes).build());
  EntityResult.Builder entityResult = EntityResult.newBuilder();
  entityResult.setEntity(entity);
  QueryResultBatch.Builder batch = QueryResultBatch.newBuilder();
  batch.addEntityResults(entityResult);
  statKindResponse.setBatch(batch);
  return statKindResponse.build();
}
 

/** Builds a response of the given timestamp. */
private static RunQueryResponse makeLatestTimestampResponse(long timestamp) {
  RunQueryResponse.Builder timestampResponse = RunQueryResponse.newBuilder();
  Entity.Builder entity = Entity.newBuilder();
  entity.setKey(makeKey("dummyKind", "dummyId"));
  entity.putProperties("timestamp", makeValue(new Date(timestamp * 1000)).build());
  EntityResult.Builder entityResult = EntityResult.newBuilder();
  entityResult.setEntity(entity);
  QueryResultBatch.Builder batch = QueryResultBatch.newBuilder();
  batch.addEntityResults(entityResult);
  timestampResponse.setBatch(batch);
  return timestampResponse.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);
  }
}
 

@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());
}
 

/** 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));
  }
}