diff --git a/gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/QueryResults.java b/gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/QueryResults.java index b23c56a7c395..b882131ba939 100644 --- a/gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/QueryResults.java +++ b/gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/QueryResults.java @@ -22,8 +22,9 @@ * The result of a Google Cloud Datastore query submission. * When the result is not typed it is possible to cast it to its appropriate type according to * the {@link #resultClass} value. - * Results are loaded lazily; therefore it is possible to get a {@code DatastoreException} - * upon {@link Iterator#hasNext hasNext} or {@link Iterator#next next} calls. + * Results are loaded lazily in batches, where batch size is set by Cloud Datastore. As a result, it + * is possible to get a {@code DatastoreException} upon {@link Iterator#hasNext hasNext} or + * {@link Iterator#next next} calls. * * @param the type of the results value. */ @@ -35,8 +36,8 @@ public interface QueryResults extends Iterator { Class resultClass(); /** - * Returns the Cursor for point after the value returned in the last {@link #next} call. - * Not currently implemented (depends on v1beta3). + * Returns the Cursor for the point after the value returned in the last {@link #next} call. + * Currently, {@code cursorAfter} returns null in all cases but the last result. */ Cursor cursorAfter(); } diff --git a/gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/QueryResultsImpl.java b/gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/QueryResultsImpl.java index cd3fe9dd776b..3c2e0d177f80 100644 --- a/gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/QueryResultsImpl.java +++ b/gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/QueryResultsImpl.java @@ -21,6 +21,7 @@ import com.google.common.base.Preconditions; import com.google.common.collect.AbstractIterator; import com.google.gcloud.datastore.Query.ResultType; +import com.google.protobuf.ByteString; import java.util.Iterator; import java.util.Objects; @@ -36,7 +37,7 @@ class QueryResultsImpl extends AbstractIterator implements QueryResults private DatastoreV1.QueryResultBatch queryResultBatchPb; private boolean lastBatch; private Iterator entityResultPbIter; - //private ByteString cursor; // only available in v1beta3 + private ByteString cursor; // only available in v1beta3 QueryResultsImpl(DatastoreImpl datastore, DatastoreV1.ReadOptions readOptionsPb, @@ -83,6 +84,7 @@ protected T computeNext() { sendRequest(); } if (!entityResultPbIter.hasNext()) { + cursor = queryResultBatchPb.getEndCursor(); return endOfData(); } DatastoreV1.EntityResult entityResultPb = entityResultPbIter.next(); @@ -99,7 +101,7 @@ public Class resultClass() { @Override public Cursor cursorAfter() { + return cursor == null ? null : new Cursor(cursor); //return new Cursor(cursor); // only available in v1beta3 - return null; } } diff --git a/gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/StructuredQuery.java b/gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/StructuredQuery.java index 293c17cf3c57..bbb73df4a79c 100644 --- a/gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/StructuredQuery.java +++ b/gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/StructuredQuery.java @@ -633,6 +633,20 @@ protected static class BaseBuilder> { this.resultType = resultType; } + BaseBuilder(StructuredQuery query) { + resultType = query.type(); + namespace = query.namespace(); + kind = query.kind; + projection.addAll(query.projection); + filter = query.filter; + groupBy.addAll(query.groupBy); + orderBy.addAll(query.orderBy); + startCursor = query.startCursor; + endCursor = query.endCursor; + offset = query.offset; + limit = query.limit; + } + @SuppressWarnings("unchecked") B self() { return (B) this; @@ -773,6 +787,10 @@ static final class Builder extends BaseBuilder> { Builder(ResultType resultType) { super(resultType); } + + Builder(StructuredQuery query) { + super(query); + } } /** @@ -953,6 +971,10 @@ public Integer limit() { return limit; } + public Builder toBuilder() { + return new Builder<>(this); + } + @Override void populatePb(DatastoreV1.RunQueryRequest.Builder requestPb) { requestPb.setQuery(toPb()); diff --git a/gcloud-java-datastore/src/test/java/com/google/gcloud/datastore/DatastoreTest.java b/gcloud-java-datastore/src/test/java/com/google/gcloud/datastore/DatastoreTest.java index 636fb3708926..e6f84c76ad40 100644 --- a/gcloud-java-datastore/src/test/java/com/google/gcloud/datastore/DatastoreTest.java +++ b/gcloud-java-datastore/src/test/java/com/google/gcloud/datastore/DatastoreTest.java @@ -27,6 +27,9 @@ import com.google.api.services.datastore.DatastoreV1; import com.google.api.services.datastore.DatastoreV1.EntityResult; +import com.google.api.services.datastore.DatastoreV1.QueryResultBatch; +import com.google.api.services.datastore.DatastoreV1.RunQueryRequest; +import com.google.api.services.datastore.DatastoreV1.RunQueryResponse; import com.google.common.collect.Iterators; import com.google.gcloud.RetryParams; import com.google.gcloud.datastore.Query.ResultType; @@ -462,6 +465,89 @@ public void testRunStructuredQuery() { assertFalse(results4.hasNext()); } + @Test + public void testQueryPaginationWithLimit() throws DatastoreRpcException { + DatastoreRpcFactory rpcFactoryMock = EasyMock.createStrictMock(DatastoreRpcFactory.class); + DatastoreRpc rpcMock = EasyMock.createStrictMock(DatastoreRpc.class); + EasyMock.expect(rpcFactoryMock.create(EasyMock.anyObject(DatastoreOptions.class))) + .andReturn(rpcMock); + List responses = buildResponsesForQueryPaginationWithLimit(); + 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 mockDatastore = options.toBuilder() + .retryParams(RetryParams.defaultInstance()) + .serviceRpcFactory(rpcFactoryMock) + .build() + .service(); + int limit = 2; + int totalCount = 0; + StructuredQuery query = Query.entityQueryBuilder().limit(limit).build(); + while (true) { + QueryResults results = mockDatastore.run(query); + int resultCount = 0; + while (results.hasNext()) { + results.next(); + resultCount++; + totalCount++; + } + if (resultCount < limit) { + break; + } + query = query.toBuilder().startCursor(results.cursorAfter()).build(); + } + assertEquals(totalCount, 5); + EasyMock.verify(rpcFactoryMock, rpcMock); + } + + private List buildResponsesForQueryPaginationWithLimit() { + Entity entity4 = Entity.builder(KEY4).set("value", StringValue.of("value")).build(); + Entity entity5 = Entity.builder(KEY5).set("value", "value").build(); + datastore.add(ENTITY3, entity4, entity5); + List responses = new ArrayList<>(); + Query query = Query.entityQueryBuilder().build(); + RunQueryRequest.Builder requestPb = RunQueryRequest.newBuilder(); + query.populatePb(requestPb); + QueryResultBatch queryResultBatchPb = RunQueryResponse.newBuilder() + .mergeFrom(((DatastoreImpl) datastore).runQuery(requestPb.build())) + .getBatch(); + QueryResultBatch queryResultBatchPb1 = QueryResultBatch.newBuilder() + .mergeFrom(queryResultBatchPb) + .setMoreResults(QueryResultBatch.MoreResultsType.NOT_FINISHED) + .clearEntityResult() + .addAllEntityResult(queryResultBatchPb.getEntityResultList().subList(0, 1)) + .setEndCursor(queryResultBatchPb.getEntityResultList().get(0).getCursor()) + .build(); + responses.add(RunQueryResponse.newBuilder().setBatch(queryResultBatchPb1).build()); + QueryResultBatch queryResultBatchPb2 = QueryResultBatch.newBuilder() + .mergeFrom(queryResultBatchPb) + .setMoreResults(QueryResultBatch.MoreResultsType.MORE_RESULTS_AFTER_LIMIT) + .clearEntityResult() + .addAllEntityResult(queryResultBatchPb.getEntityResultList().subList(1, 2)) + .setEndCursor(queryResultBatchPb.getEntityResultList().get(1).getCursor()) + .build(); + responses.add(RunQueryResponse.newBuilder().setBatch(queryResultBatchPb2).build()); + QueryResultBatch queryResultBatchPb3 = QueryResultBatch.newBuilder() + .mergeFrom(queryResultBatchPb) + .setMoreResults(QueryResultBatch.MoreResultsType.MORE_RESULTS_AFTER_LIMIT) + .clearEntityResult() + .addAllEntityResult(queryResultBatchPb.getEntityResultList().subList(2, 4)) + .setEndCursor(queryResultBatchPb.getEntityResultList().get(3).getCursor()) + .build(); + responses.add(RunQueryResponse.newBuilder().setBatch(queryResultBatchPb3).build()); + QueryResultBatch queryResultBatchPb4 = QueryResultBatch.newBuilder() + .mergeFrom(queryResultBatchPb) + .setMoreResults(QueryResultBatch.MoreResultsType.NO_MORE_RESULTS) + .clearEntityResult() + .addAllEntityResult(queryResultBatchPb.getEntityResultList().subList(4, 5)) + .setEndCursor(queryResultBatchPb.getEntityResultList().get(4).getCursor()) + .build(); + responses.add(RunQueryResponse.newBuilder().setBatch(queryResultBatchPb4).build()); + return responses; + } + @Test public void testAllocateId() { KeyFactory keyFactory = datastore.newKeyFactory().kind(KIND1);