Populate cursorAfter in datastore v1beta2

gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/QueryResults.java

@@ -22,8 +22,9 @@
  * The result of a Google Cloud Datastore query submission.
  * When the result is not typed it is possible to cast it to its appropriate type according to
  * the {@link #resultClass} value.
- * Results are loaded lazily; therefore it is possible to get a {@code DatastoreException}
- * upon {@link Iterator#hasNext hasNext} or {@link Iterator#next next} calls.
+ * Results are loaded lazily in batches, where batch size is set by Cloud Datastore. As a result, it
+ * is possible to get a {@code DatastoreException} upon {@link Iterator#hasNext hasNext} or
+ * {@link Iterator#next next} calls.
  *
  * @param <V> the type of the results value.
  */
@@ -35,8 +36,8 @@ public interface QueryResults<V> extends Iterator<V> {
   Class<?> resultClass();
 
   /**
-   * Returns the Cursor for point after the value returned in the last {@link #next} call.
-   * Not currently implemented (depends on v1beta3).
+   * Returns the Cursor for the point after the value returned in the last {@link #next} call.
+   * Currently, {@code cursorAfter} returns null in all cases but the last result.
    */
   Cursor cursorAfter();
 }

gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/QueryResultsImpl.java

@@ -21,6 +21,7 @@
 import com.google.common.base.Preconditions;
 import com.google.common.collect.AbstractIterator;
 import com.google.gcloud.datastore.Query.ResultType;
+import com.google.protobuf.ByteString;
 
 import java.util.Iterator;
 import java.util.Objects;
@@ -36,7 +37,7 @@ class QueryResultsImpl<T> extends AbstractIterator<T> implements QueryResults<T>
   private DatastoreV1.QueryResultBatch queryResultBatchPb;
   private boolean lastBatch;
   private Iterator<DatastoreV1.EntityResult> entityResultPbIter;
-  //private ByteString cursor; // only available in v1beta3
+  private ByteString cursor; // only available in v1beta3
 
 
   QueryResultsImpl(DatastoreImpl datastore, DatastoreV1.ReadOptions readOptionsPb,
@@ -83,6 +84,7 @@ protected T computeNext() {
       sendRequest();
     }
     if (!entityResultPbIter.hasNext()) {
+      cursor = queryResultBatchPb.getEndCursor();
       return endOfData();
     }
     DatastoreV1.EntityResult entityResultPb = entityResultPbIter.next();
@@ -99,7 +101,7 @@ public Class<?> resultClass() {
 
   @Override
   public Cursor cursorAfter() {
+    return cursor == null ? null : new Cursor(cursor);
     //return new Cursor(cursor); // only available in v1beta3
-    return null;
   }
 }

gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/StructuredQuery.java

@@ -633,6 +633,20 @@ protected static class BaseBuilder<V, B extends BaseBuilder<V, B>> {
       this.resultType = resultType;
     }
 
+    BaseBuilder(StructuredQuery<V> query) {
+      resultType = query.type();
+      namespace = query.namespace();
+      kind = query.kind;
+      projection.addAll(query.projection);
+      filter = query.filter;
+      groupBy.addAll(query.groupBy);
+      orderBy.addAll(query.orderBy);
+      startCursor = query.startCursor;
+      endCursor = query.endCursor;
+      offset = query.offset;
+      limit = query.limit;
+    }
+
     @SuppressWarnings("unchecked")
     B self() {
       return (B) this;
@@ -773,6 +787,10 @@ static final class Builder<V> extends BaseBuilder<V, Builder<V>> {
     Builder(ResultType<V> resultType) {
       super(resultType);
     }
+
+    Builder(StructuredQuery<V> query) {
+      super(query);
+    }
   }
 
   /**
@@ -953,6 +971,10 @@ public Integer limit() {
     return limit;
   }
 
+  public Builder<V> toBuilder() {
+    return new Builder<>(this);
+  }
+
   @Override
   void populatePb(DatastoreV1.RunQueryRequest.Builder requestPb) {
     requestPb.setQuery(toPb());

gcloud-java-datastore/src/test/java/com/google/gcloud/datastore/DatastoreTest.java

@@ -27,6 +27,9 @@
 
 import com.google.api.services.datastore.DatastoreV1;
 import com.google.api.services.datastore.DatastoreV1.EntityResult;
+import com.google.api.services.datastore.DatastoreV1.QueryResultBatch;
+import com.google.api.services.datastore.DatastoreV1.RunQueryRequest;
+import com.google.api.services.datastore.DatastoreV1.RunQueryResponse;
 import com.google.common.collect.Iterators;
 import com.google.gcloud.RetryParams;
 import com.google.gcloud.datastore.Query.ResultType;
@@ -462,6 +465,89 @@ public void testRunStructuredQuery() {
     assertFalse(results4.hasNext());
   }
 
+  @Test
+  public void testQueryPaginationWithLimit() throws DatastoreRpcException {
+    DatastoreRpcFactory rpcFactoryMock = EasyMock.createStrictMock(DatastoreRpcFactory.class);
+    DatastoreRpc rpcMock = EasyMock.createStrictMock(DatastoreRpc.class);
+    EasyMock.expect(rpcFactoryMock.create(EasyMock.anyObject(DatastoreOptions.class)))
+        .andReturn(rpcMock);
+    List<RunQueryResponse> responses = buildResponsesForQueryPaginationWithLimit();
+    for (int i = 0; i < responses.size(); i++) {
+      EasyMock.expect(rpcMock.runQuery(EasyMock.anyObject(RunQueryRequest.class)))
+          .andReturn(responses.get(i));
+    }
+    EasyMock.replay(rpcFactoryMock, rpcMock);
+    Datastore mockDatastore = options.toBuilder()
+        .retryParams(RetryParams.defaultInstance())
+        .serviceRpcFactory(rpcFactoryMock)
+        .build()
+        .service();
+    int limit = 2;
+    int totalCount = 0;
+    StructuredQuery<Entity> query = Query.entityQueryBuilder().limit(limit).build();
+    while (true) {
+      QueryResults<Entity> results = mockDatastore.run(query);
+      int resultCount = 0;
+      while (results.hasNext()) {
+        results.next();
+        resultCount++;
+        totalCount++;
+      }
+      if (resultCount < limit) {
+        break;
+      }
+      query = query.toBuilder().startCursor(results.cursorAfter()).build();
+    }
+    assertEquals(totalCount, 5);
+    EasyMock.verify(rpcFactoryMock, rpcMock);
+  }
+
+  private List<RunQueryResponse> buildResponsesForQueryPaginationWithLimit() {
+    Entity entity4 = Entity.builder(KEY4).set("value", StringValue.of("value")).build();
+    Entity entity5 = Entity.builder(KEY5).set("value", "value").build();
+    datastore.add(ENTITY3, entity4, entity5);
+    List<RunQueryResponse> responses = new ArrayList<>();
+    Query<Entity> query = Query.entityQueryBuilder().build();
+    RunQueryRequest.Builder requestPb = RunQueryRequest.newBuilder();
+    query.populatePb(requestPb);
+    QueryResultBatch queryResultBatchPb = RunQueryResponse.newBuilder()
+        .mergeFrom(((DatastoreImpl) datastore).runQuery(requestPb.build()))
+        .getBatch();
+    QueryResultBatch queryResultBatchPb1 = QueryResultBatch.newBuilder()
+        .mergeFrom(queryResultBatchPb)
+        .setMoreResults(QueryResultBatch.MoreResultsType.NOT_FINISHED)
+        .clearEntityResult()
+        .addAllEntityResult(queryResultBatchPb.getEntityResultList().subList(0, 1))
+        .setEndCursor(queryResultBatchPb.getEntityResultList().get(0).getCursor())
+        .build();
+    responses.add(RunQueryResponse.newBuilder().setBatch(queryResultBatchPb1).build());
+    QueryResultBatch queryResultBatchPb2 = QueryResultBatch.newBuilder()
+        .mergeFrom(queryResultBatchPb)
+        .setMoreResults(QueryResultBatch.MoreResultsType.MORE_RESULTS_AFTER_LIMIT)
+        .clearEntityResult()
+        .addAllEntityResult(queryResultBatchPb.getEntityResultList().subList(1, 2))
+        .setEndCursor(queryResultBatchPb.getEntityResultList().get(1).getCursor())
+        .build();
+    responses.add(RunQueryResponse.newBuilder().setBatch(queryResultBatchPb2).build());
+    QueryResultBatch queryResultBatchPb3 = QueryResultBatch.newBuilder()
+        .mergeFrom(queryResultBatchPb)
+        .setMoreResults(QueryResultBatch.MoreResultsType.MORE_RESULTS_AFTER_LIMIT)
+        .clearEntityResult()
+        .addAllEntityResult(queryResultBatchPb.getEntityResultList().subList(2, 4))
+        .setEndCursor(queryResultBatchPb.getEntityResultList().get(3).getCursor())
+        .build();
+    responses.add(RunQueryResponse.newBuilder().setBatch(queryResultBatchPb3).build());
+    QueryResultBatch queryResultBatchPb4 = QueryResultBatch.newBuilder()
+        .mergeFrom(queryResultBatchPb)
+        .setMoreResults(QueryResultBatch.MoreResultsType.NO_MORE_RESULTS)
+        .clearEntityResult()
+        .addAllEntityResult(queryResultBatchPb.getEntityResultList().subList(4, 5))
+        .setEndCursor(queryResultBatchPb.getEntityResultList().get(4).getCursor())
+        .build();
+    responses.add(RunQueryResponse.newBuilder().setBatch(queryResultBatchPb4).build());
+    return responses;
+  }
+
   @Test
   public void testAllocateId() {
     KeyFactory keyFactory = datastore.newKeyFactory().kind(KIND1);