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 dsApiProxy;
  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 dsApiProxy) {
  19     super(callbacks, fetchOptions, txn, query, runQueryResponse);
  20     this.initialRequest = request;
  21     this.dsApiProxy = dsApiProxy;
  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     QueryResultBatch latestBatch = ((WrappedQueryResultCloudDatastoreV1) latestResult).getBatch();
  35     if (latestBatch.getEndCursor().isEmpty()) {
  36       throw new IllegalArgumentException();
  37     }
  38     query.setStartCursor(latestBatch.getEndCursor());
  39     if (query.hasLimit()) {
  40       remainingLimit -= latestBatch.getEntityResultsCount();
  41       query.getLimitBuilder().setValue(Math.max(remainingLimit, 0));
  42     }
  43     if (offset != null) {
  44       query.setOffset(offset);
  45     } else {
  46       query.clearOffset();
  47     }
  48     return dsApiProxy.runQuery(runQueryRequest.build());
  49   }
  50
  51   @Override
  52   WrappedQueryResult wrapInitialResult(RunQueryResponse initialResponse) {
  53     return new WrappedQueryResultCloudDatastoreV1(initialResponse.getBatch());
  54   }
  55
  56   @Override
  57   WrappedQueryResult wrapResult(RunQueryResponse res) {
  58     return new WrappedQueryResultCloudDatastoreV1(res.getBatch());
  59   }
  60 }