com.google.datastore.v1.QueryResultBatch Java Examples

/**
 * Initialize a QueryState from a Datastore QueryResultBatch.MoreResultsType enum value.
 *
 * @param moreResultsType
 *          original enum value.
 * @return query state enum
 */
public static QueryState forMoreResultsType(QueryResultBatch.MoreResultsType moreResultsType) {
  if (null == moreResultsType) {
    return null;
  }
  switch (moreResultsType) {
    case MORE_RESULTS_TYPE_UNSPECIFIED:
      return MORE_RESULTS_TYPE_UNSPECIFIED;
    case NOT_FINISHED:
      return NOT_FINISHED;
    case MORE_RESULTS_AFTER_LIMIT:
      return MORE_RESULTS_AFTER_LIMIT;
    case MORE_RESULTS_AFTER_CURSOR:
      return MORE_RESULTS_AFTER_CURSOR;
    case NO_MORE_RESULTS:
      return NO_MORE_RESULTS;
    case UNRECOGNIZED:
      return UNRECOGNIZED;
    default:
      throw new IllegalArgumentException(
          "Unknown or unimplemented QueryResultBatch.MoreResultsType: " + moreResultsType);
  }
}
 

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

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

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

@Parameterized.Parameters
public static Iterable<Object[]> testCases() {
  return Arrays.asList(new Object[][] {
    {
      QueryResultBatch.MoreResultsType.MORE_RESULTS_TYPE_UNSPECIFIED,
      QueryResponseMetadata.QueryState.MORE_RESULTS_TYPE_UNSPECIFIED
    },
    {
      QueryResultBatch.MoreResultsType.NOT_FINISHED,
      QueryResponseMetadata.QueryState.NOT_FINISHED,
    },
    {
      QueryResultBatch.MoreResultsType.MORE_RESULTS_AFTER_LIMIT,
      QueryResponseMetadata.QueryState.MORE_RESULTS_AFTER_LIMIT,
    },
    {
      QueryResultBatch.MoreResultsType.MORE_RESULTS_AFTER_CURSOR,
      QueryResponseMetadata.QueryState.MORE_RESULTS_AFTER_CURSOR,
    },
    {
      QueryResultBatch.MoreResultsType.NO_MORE_RESULTS,
      QueryResponseMetadata.QueryState.NO_MORE_RESULTS,
    },
    {
      QueryResultBatch.MoreResultsType.UNRECOGNIZED,
      QueryResponseMetadata.QueryState.UNRECOGNIZED,
    },
    {
      null,
      null,
    },
  });
}
 

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

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

@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 QueryResponseMetadataTest(QueryResultBatch.MoreResultsType input,
  QueryResponseMetadata.QueryState expected) {
  this.input = input;
  this.expected = expected;
}
 

@Override
public QueryResultBatch.MoreResultsType getMoreResults() {
  return results.getMoreResults();
}
 

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