diff --git a/google-cloud-clients/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResults.java b/google-cloud-clients/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResults.java index 27c6f19aebf4..3f73824dcdbc 100644 --- a/google-cloud-clients/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResults.java +++ b/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 extends Iterator { * } */ int getSkippedResults(); + + /** Returns MoreResults state of the query after the current batch. */ + QueryResultBatch.MoreResultsType getMoreResults(); } diff --git a/google-cloud-clients/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResultsImpl.java b/google-cloud-clients/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResultsImpl.java index c6179686707e..3c8bba10a089 100644 --- a/google-cloud-clients/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResultsImpl.java +++ b/google-cloud-clients/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResultsImpl.java @@ -37,6 +37,7 @@ class QueryResultsImpl extends AbstractIterator implements QueryResults private boolean lastBatch; private Iterator entityResultPbIter; private ByteString cursor; + private MoreResultsType moreResults; QueryResultsImpl( DatastoreImpl datastore, com.google.datastore.v1.ReadOptions readOptionsPb, Query 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/google-cloud-clients/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java b/google-cloud-clients/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java index 25d9c9d88c40..2183a4983efd 100644 --- a/google-cloud-clients/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java +++ b/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 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 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 buildResponsesForQueryPagination() { Entity entity4 = Entity.newBuilder(KEY4).set("value", StringValue.of("value")).build(); Entity entity5 = Entity.newBuilder(KEY5).set("value", "value").build();