Datastore: Fix more_results value to query results (#4516)

* add moreResults

* added getter

* updated code

* modified code and added test case

Author: Praful Makani

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();
 }

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;
+  }
 }

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();