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

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