---

Praful Makani · sduskis · commit 0120c9345b6c · 2019-02-21T10:39:57.000-05:00
yaml --- r: 16119 b: refs/heads/autosynth-errorreporting c: a1679e4 h: refs/heads/master i: 16117: dd3b804 16115: d533bdb 16111: 8187162
diff --git a/[refs] b/[refs]
@@ -131,7 +131,7 @@ refs/heads/autosynth-bigtable-admin: 6379a2bc712f2736c83de0e009b4d26da4fa82ca
 refs/heads/autosynth-containeranalysis: 039ca5b8db725c76c16a965ff26b2774322b8ef8
 refs/heads/autosynth-datastore: 9acd400b484d6691a080c9152a331d88d24fefc1
 refs/heads/autosynth-dialogflow: d7477419376eac9b6dcc7dbcede581152527351d
-refs/heads/autosynth-errorreporting: 07ab94105146ef83558985e100ac86ef2785a11c
+refs/heads/autosynth-errorreporting: a1679e4551b3bc0f9f78db47cc0e7747121a88cf
 refs/heads/autosynth-firestore: 92b27fbc8855c9902168695abb0a8f1f433b750b
 refs/heads/autosynth-iot: 9d732be07d99843d8cb53d34ec0837328a807fce
 refs/heads/autosynth-kms: dcc6e15d68759010c8735cc868135bd7e6c1cc5f
diff --git a/branches/autosynth-errorreporting/google-cloud-clients/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResults.java b/branches/autosynth-errorreporting/google-cloud-clients/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResults.java
@@ -16,6 +16,7 @@
 
 package com.google.cloud.datastore;
 
+import com.google.datastore.v1.QueryResultBatch;
 import java.util.Iterator;
 
 /**
@@ -66,4 +67,7 @@ public interface QueryResults<V> extends Iterator<V> {
    * }</pre>
    */
   int getSkippedResults();
+
+  /** Returns MoreResults state of the query after the current batch. */
+  QueryResultBatch.MoreResultsType getMoreResults();
 }
diff --git a/branches/autosynth-errorreporting/google-cloud-clients/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResultsImpl.java b/branches/autosynth-errorreporting/google-cloud-clients/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResultsImpl.java
@@ -37,6 +37,7 @@ class QueryResultsImpl<T> extends AbstractIterator<T> implements QueryResults<T>
   private boolean lastBatch;
   private Iterator<com.google.datastore.v1.EntityResult> entityResultPbIter;
   private ByteString cursor;
+  private MoreResultsType moreResults;
 
   QueryResultsImpl(
       DatastoreImpl datastore, com.google.datastore.v1.ReadOptions readOptionsPb, Query<T> query) {
@@ -74,7 +75,8 @@ private void sendRequest() {
     if (mostRecentQueryPb == null) {
       mostRecentQueryPb = requestPb.getQuery();
     }
-    lastBatch = runQueryResponsePb.getBatch().getMoreResults() != MoreResultsType.NOT_FINISHED;
+    moreResults = runQueryResponsePb.getBatch().getMoreResults();
+    lastBatch = moreResults != MoreResultsType.NOT_FINISHED;
     entityResultPbIter = runQueryResponsePb.getBatch().getEntityResultsList().iterator();
     actualResultType = ResultType.fromPb(runQueryResponsePb.getBatch().getEntityResultType());
     if (Objects.equals(queryResultType, ResultType.PROJECTION_ENTITY)) {
@@ -117,4 +119,9 @@ public Cursor getCursorAfter() {
   public int getSkippedResults() {
     return runQueryResponsePb.getBatch().getSkippedResults();
   }
+
+  @Override
+  public MoreResultsType getMoreResults() {
+    return moreResults;
+  }
 }
diff --git a/branches/autosynth-errorreporting/google-cloud-clients/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java b/branches/autosynth-errorreporting/google-cloud-clients/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java
@@ -585,6 +585,26 @@ public void testStructuredQueryPagination() throws DatastoreException {
     EasyMock.verify(rpcFactoryMock, rpcMock);
   }
 
+  @Test
+  public void testStructuredQueryPaginationWithMoreResults() throws DatastoreException {
+    List<RunQueryResponse> responses = buildResponsesForQueryPagination();
+    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 datastore = rpcMockOptions.getService();
+    QueryResults<Key> results = datastore.run(Query.newKeyQueryBuilder().build());
+    int count = 0;
+    while (results.hasNext()) {
+      count += 1;
+      results.next();
+    }
+    assertEquals(count, 5);
+    assertEquals(QueryResultBatch.MoreResultsType.NO_MORE_RESULTS, results.getMoreResults());
+    EasyMock.verify(rpcFactoryMock, rpcMock);
+  }
+
   private List<RunQueryResponse> buildResponsesForQueryPagination() {
     Entity entity4 = Entity.newBuilder(KEY4).set("value", StringValue.of("value")).build();
     Entity entity5 = Entity.newBuilder(KEY5).set("value", "value").build();
