java/src/main/com/google/appengine/api/datastore/QueryResultsSourceCloudDatastoreV1.java

   1 package com.google.appengine.api.datastore;
   2
   3 import com.google.datastore.v1beta3.QueryResultBatch;
   4 import com.google.datastore.v1beta3.RunQueryRequest;
   5 import com.google.datastore.v1beta3.RunQueryResponse;
   6
   7 import java.util.concurrent.Future;
   8
   9 class QueryResultsSourceCloudDatastoreV1 extends
  10     BaseQueryResultsSource<RunQueryResponse, RunQueryRequest, RunQueryResponse> {
  11
  12   private final CloudDatastoreV1Proxy cloudDatastoreV1Proxy;
  13   private final RunQueryRequest initialRequest;
  14   private int remainingLimit;
  15
  16   QueryResultsSourceCloudDatastoreV1(DatastoreCallbacks callbacks, FetchOptions fetchOptions,
  17       Transaction txn, Query query, RunQueryRequest request,
  18       Future<RunQueryResponse> runQueryResponse, CloudDatastoreV1Proxy cloudDatastoreV1Proxy) {
  19     super(callbacks, fetchOptions, txn, query, runQueryResponse);
  20     this.initialRequest = request;
  21     this.cloudDatastoreV1Proxy = cloudDatastoreV1Proxy;
  22     remainingLimit = fetchOptions.getLimit() != null ? fetchOptions.getLimit() : -1;
  23   }
  24
  25   @Override
  26   RunQueryRequest buildNextCallPrototype(RunQueryResponse initialResponse) {
  27     return initialRequest;
  28   }
  29
  30   @Override
  31   Future<RunQueryResponse> makeNextCall(RunQueryRequest prototype, WrappedQueryResult latestResult, Integer fetchCount, Integer offset) {
  32     RunQueryRequest.Builder runQueryRequest = prototype.toBuilder();
  33     com.google.datastore.v1beta3.Query.Builder query = runQueryRequest.getQueryBuilder();
  34     query.setStartCursor(latestResult.getEndCursor().toByteString());
  35     QueryResultBatch latestBatch = ((WrappedQueryResultCloudDatastoreV1) latestResult).getBatch();
  36     if (query.hasLimit()) {
  37       remainingLimit -= latestBatch.getEntityResultsCount();
  38       query.getLimitBuilder().setValue(Math.max(remainingLimit, 0));
  39     }
  40     if (offset != null) {
  41       query.setOffset(offset);
  42     } else {
  43       query.clearOffset();
  44     }
  45     return cloudDatastoreV1Proxy.runQuery(runQueryRequest.build());
  46   }
  47
  48   @Override
  49   WrappedQueryResult wrapInitialResult(RunQueryResponse initialResponse) {
  50     return new WrappedQueryResultCloudDatastoreV1(initialResponse.getBatch());
  51   }
  52
  53   @Override
  54   WrappedQueryResult wrapResult(RunQueryResponse res) {
  55     return new WrappedQueryResultCloudDatastoreV1(res.getBatch());
  56   }
  57 }