---

Praful Makani · sduskis · commit 90c3da64940a · 2019-02-21T10:39:57.000-05:00
yaml --- r: 13827 b: refs/heads/autosynth-logging c: a1679e4 h: refs/heads/master i: 13825: 2bc4ec4 13823: 0d10c48
diff --git a/[refs] b/[refs]
@@ -105,7 +105,7 @@ refs/tags/v0.61.0: e4b526656bb1bf5eefd0ee578b7405147821225e
 refs/tags/v0.62.0: bbede7385d48ba08f487bdd29ec10668ace96396
 refs/heads/0.60.0-alpha: 10939381ffe0b8da32db4fe3087c86e3aa7f3e55
 refs/heads/autosynth-dlp: 90e4efe3be392e89dc36bce22afab981a905b1bc
-refs/heads/autosynth-logging: 07ab94105146ef83558985e100ac86ef2785a11c
+refs/heads/autosynth-logging: a1679e4551b3bc0f9f78db47cc0e7747121a88cf
 refs/heads/dupes: 3478c5d81fd242d0e985656645a679420a2060c2
 refs/tags/v0.63.0: 94f19b71d40f46b36120e7b9d78a1a3d41bfcbd6
 refs/tags/v0.64.0: 456e8fbd129deced3ca025f239a2d8a82bde1d0a
diff --git a/branches/autosynth-logging/google-cloud-clients/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResults.java b/branches/autosynth-logging/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-logging/google-cloud-clients/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResultsImpl.java b/branches/autosynth-logging/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-logging/google-cloud-clients/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java b/branches/autosynth-logging/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();
