---

Praful Makani · sduskis · commit b925bc20ca43 · 2019-02-21T10:39:57.000-05:00
yaml --- r: 17079 b: refs/heads/autosynth-redis c: a1679e4 h: refs/heads/master i: 17077: b880d72 17075: 9e6417f 17071: 1ad7b82
diff --git a/[refs] b/[refs]
@@ -137,7 +137,7 @@ refs/heads/autosynth-iot: f16fc4ba85bbb37eacecb64a2a7336f685813f6c
 refs/heads/autosynth-kms: 4fff8f3cf6c830f7706ac4ef272f036bb87fbaf8
 refs/heads/autosynth-language: 6262e2eb76944f01972c887b3e684aaf65ec999a
 refs/heads/autosynth-os-login: a88a337797996a205873040a63abe1d3116f5789
-refs/heads/autosynth-redis: 07ab94105146ef83558985e100ac86ef2785a11c
+refs/heads/autosynth-redis: a1679e4551b3bc0f9f78db47cc0e7747121a88cf
 refs/heads/autosynth-scheduler: 2f0fe714a8541dc72f88b45bddd8748e5f21cf29
 refs/heads/autosynth-securitycenter: 8e95447ccb176c3648880ee0cb926cd1f4065156
 refs/heads/autosynth-spanner: 563dc6c78c6a2579b681f441f0a9b59005212394
diff --git a/branches/autosynth-redis/google-cloud-clients/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResults.java b/branches/autosynth-redis/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-redis/google-cloud-clients/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResultsImpl.java b/branches/autosynth-redis/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-redis/google-cloud-clients/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java b/branches/autosynth-redis/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();
