Skip to content

Commit b54c501

Browse files
aozarovaozarov
committed
Merge pull request #273 from ajkannan/gql-pagination
Query continuation implemented for GqlQuery + tests
2 parents 9e4692d + 6ffe927 commit b54c501

File tree

5 files changed

+141
-22
lines changed

5 files changed

+141
-22
lines changed

‎gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/GqlQuery.java

+3-3
Original file line numberDiff line numberDiff line change
@@ -379,9 +379,9 @@ protected void populatePb(com.google.datastore.v1beta3.RunQueryRequest.Builder r
379379
}
380380

381381
@Override
382-
protected GqlQuery<V> nextQuery(com.google.datastore.v1beta3.QueryResultBatch responsePb) {
383-
// See issue #17
384-
throw new UnsupportedOperationException("paging for this query is not implemented yet");
382+
protected Query<V> nextQuery(com.google.datastore.v1beta3.RunQueryResponse responsePb) {
383+
return StructuredQuery.<V>fromPb(type(), namespace(), responsePb.getQuery())
384+
.nextQuery(responsePb);
385385
}
386386

387387
@Override

‎gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/Query.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -192,7 +192,7 @@ protected abstract Object fromPb(ResultType<V> resultType, String namespace, byt
192192
protected abstract void populatePb(
193193
com.google.datastore.v1beta3.RunQueryRequest.Builder requestPb);
194194

195-
protected abstract Query<V> nextQuery(com.google.datastore.v1beta3.QueryResultBatch responsePb);
195+
protected abstract Query<V> nextQuery(com.google.datastore.v1beta3.RunQueryResponse responsePb);
196196

197197
/**
198198
* Returns a new {@link GqlQuery} builder.

‎gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/QueryResultsImpl.java

+8-10
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ class QueryResultsImpl<T> extends AbstractIterator<T> implements QueryResults<T>
3333
private final ResultType<T> queryResultType;
3434
private Query<T> query;
3535
private ResultType<?> actualResultType;
36-
private com.google.datastore.v1beta3.QueryResultBatch queryResultBatchPb;
36+
private com.google.datastore.v1beta3.RunQueryResponse runQueryResponsePb;
3737
private com.google.datastore.v1beta3.Query mostRecentQueryPb;
3838
private boolean lastBatch;
3939
private Iterator<com.google.datastore.v1beta3.EntityResult> entityResultPbIter;
@@ -56,8 +56,8 @@ class QueryResultsImpl<T> extends AbstractIterator<T> implements QueryResults<T>
5656
}
5757
partitionIdPb = pbBuilder.build();
5858
sendRequest();
59-
if (queryResultBatchPb.getSkippedResults() > 0) {
60-
cursor = queryResultBatchPb.getSkippedCursor();
59+
if (runQueryResponsePb.getBatch().getSkippedResults() > 0) {
60+
cursor = runQueryResponsePb.getBatch().getSkippedCursor();
6161
} else {
6262
cursor = mostRecentQueryPb.getStartCursor();
6363
}
@@ -71,16 +71,14 @@ private void sendRequest() {
7171
}
7272
requestPb.setPartitionId(partitionIdPb);
7373
query.populatePb(requestPb);
74-
com.google.datastore.v1beta3.RunQueryResponse runQueryResponsePb =
75-
datastore.runQuery(requestPb.build());
76-
queryResultBatchPb = runQueryResponsePb.getBatch();
74+
runQueryResponsePb = datastore.runQuery(requestPb.build());
7775
mostRecentQueryPb = runQueryResponsePb.getQuery();
7876
if (mostRecentQueryPb == null) {
7977
mostRecentQueryPb = requestPb.getQuery();
8078
}
81-
lastBatch = queryResultBatchPb.getMoreResults() != MoreResultsType.NOT_FINISHED;
82-
entityResultPbIter = queryResultBatchPb.getEntityResultsList().iterator();
83-
actualResultType = ResultType.fromPb(queryResultBatchPb.getEntityResultType());
79+
lastBatch = runQueryResponsePb.getBatch().getMoreResults() != MoreResultsType.NOT_FINISHED;
80+
entityResultPbIter = runQueryResponsePb.getBatch().getEntityResultsList().iterator();
81+
actualResultType = ResultType.fromPb(runQueryResponsePb.getBatch().getEntityResultType());
8482
if (Objects.equals(queryResultType, ResultType.PROJECTION_ENTITY)) {
8583
// projection entity can represent all type of results
8684
actualResultType = ResultType.PROJECTION_ENTITY;
@@ -92,7 +90,7 @@ private void sendRequest() {
9290
@Override
9391
protected T computeNext() {
9492
while (!entityResultPbIter.hasNext() && !lastBatch) {
95-
query = query.nextQuery(queryResultBatchPb);
93+
query = query.nextQuery(runQueryResponsePb);
9694
sendRequest();
9795
}
9896
if (!entityResultPbIter.hasNext()) {

‎gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/StructuredQuery.java

+9-7
Original file line numberDiff line numberDiff line change
@@ -844,16 +844,16 @@ protected void populatePb(com.google.datastore.v1beta3.RunQueryRequest.Builder r
844844
}
845845

846846
@Override
847-
protected StructuredQuery<V> nextQuery(com.google.datastore.v1beta3.QueryResultBatch responsePb) {
847+
protected Query<V> nextQuery(com.google.datastore.v1beta3.RunQueryResponse responsePb) {
848848
Builder<V> builder = new Builder<>(type());
849849
builder.mergeFrom(toPb());
850-
builder.startCursor(new Cursor(responsePb.getEndCursor()));
851-
if (offset > 0 && responsePb.getSkippedResults() < offset) {
852-
builder.offset(offset - responsePb.getSkippedResults());
850+
builder.startCursor(new Cursor(responsePb.getBatch().getEndCursor()));
851+
if (offset > 0 && responsePb.getBatch().getSkippedResults() < offset) {
852+
builder.offset(offset - responsePb.getBatch().getSkippedResults());
853853
} else {
854854
builder.offset(0);
855855
if (limit != null) {
856-
builder.limit(limit - responsePb.getEntityResultsCount());
856+
builder.limit(limit - responsePb.getBatch().getEntityResultsCount());
857857
}
858858
}
859859
return builder.build();
@@ -904,7 +904,9 @@ protected Object fromPb(ResultType<V> resultType, String namespace, byte[] bytes
904904
return fromPb(resultType, namespace, com.google.datastore.v1beta3.Query.parseFrom(bytesPb));
905905
}
906906

907-
private static StructuredQuery<?> fromPb(ResultType<?> resultType, String namespace,
907+
@SuppressWarnings("unchecked")
908+
static <V> StructuredQuery<V> fromPb(
909+
ResultType<?> resultType, String namespace,
908910
com.google.datastore.v1beta3.Query queryPb) {
909911
BaseBuilder<?, ?> builder;
910912
if (resultType.equals(ResultType.ENTITY)) {
@@ -914,6 +916,6 @@ private static StructuredQuery<?> fromPb(ResultType<?> resultType, String namesp
914916
} else {
915917
builder = new ProjectionEntityQueryBuilder();
916918
}
917-
return builder.namespace(namespace).mergeFrom(queryPb).build();
919+
return (StructuredQuery<V>) builder.namespace(namespace).mergeFrom(queryPb).build();
918920
}
919921
}

‎gcloud-java-datastore/src/test/java/com/google/gcloud/datastore/DatastoreTest.java

+120-1
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
import com.google.gcloud.datastore.StructuredQuery.OrderBy;
3232
import com.google.gcloud.datastore.StructuredQuery.PropertyFilter;
3333
import com.google.gcloud.spi.DatastoreRpc;
34+
import com.google.gcloud.spi.DatastoreRpc.DatastoreRpcException;
3435
import com.google.gcloud.spi.DatastoreRpc.DatastoreRpcException.Reason;
3536
import com.google.gcloud.spi.DatastoreRpcFactory;
3637

@@ -43,6 +44,7 @@
4344
import org.junit.runners.JUnit4;
4445

4546
import java.io.IOException;
47+
import java.util.ArrayList;
4648
import java.util.Collections;
4749
import java.util.Iterator;
4850
import java.util.List;
@@ -407,6 +409,38 @@ public void testRunGqlQueryWithCasting() {
407409
assertFalse(results3.hasNext());
408410
}
409411

412+
@Test
413+
public void testGqlQueryPagination() throws DatastoreRpcException {
414+
DatastoreRpcFactory rpcFactoryMock = EasyMock.createStrictMock(DatastoreRpcFactory.class);
415+
DatastoreRpc rpcMock = EasyMock.createStrictMock(DatastoreRpc.class);
416+
EasyMock.expect(rpcFactoryMock.create(EasyMock.anyObject(DatastoreOptions.class)))
417+
.andReturn(rpcMock);
418+
List<com.google.datastore.v1beta3.RunQueryResponse> responses =
419+
buildResponsesForQueryPagination();
420+
for (int i = 0; i < responses.size(); i++) {
421+
EasyMock
422+
.expect(rpcMock.runQuery(
423+
EasyMock.anyObject(com.google.datastore.v1beta3.RunQueryRequest.class)))
424+
.andReturn(responses.get(i));
425+
}
426+
EasyMock.replay(rpcFactoryMock, rpcMock);
427+
DatastoreOptions options =
428+
this.options.toBuilder()
429+
.retryParams(RetryParams.getDefaultInstance())
430+
.serviceRpcFactory(rpcFactoryMock)
431+
.build();
432+
Datastore mockDatastore = DatastoreFactory.instance().get(options);
433+
QueryResults<Key> results =
434+
mockDatastore.run(Query.gqlQueryBuilder(ResultType.KEY, "select __key__ from *").build());
435+
int count = 0;
436+
while (results.hasNext()) {
437+
count += 1;
438+
results.next();
439+
}
440+
assertEquals(count, 5);
441+
EasyMock.verify(rpcFactoryMock, rpcMock);
442+
}
443+
410444
@Test
411445
public void testRunStructuredQuery() {
412446
Query<Entity> query =
@@ -449,7 +483,92 @@ public void testRunStructuredQuery() {
449483
assertEquals(20, entity.getLong("age"));
450484
assertEquals(1, entity.properties().size());
451485
assertFalse(results4.hasNext());
452-
// TODO(ozarov): construct a test to verify nextQuery/pagination
486+
}
487+
488+
@Test
489+
public void testStructuredQueryPagination() throws DatastoreRpcException {
490+
DatastoreRpcFactory rpcFactoryMock = EasyMock.createStrictMock(DatastoreRpcFactory.class);
491+
DatastoreRpc rpcMock = EasyMock.createStrictMock(DatastoreRpc.class);
492+
EasyMock.expect(rpcFactoryMock.create(EasyMock.anyObject(DatastoreOptions.class)))
493+
.andReturn(rpcMock);
494+
List<com.google.datastore.v1beta3.RunQueryResponse> responses =
495+
buildResponsesForQueryPagination();
496+
for (int i = 0; i < responses.size(); i++) {
497+
EasyMock
498+
.expect(rpcMock.runQuery(
499+
EasyMock.anyObject(com.google.datastore.v1beta3.RunQueryRequest.class)))
500+
.andReturn(responses.get(i));
501+
}
502+
EasyMock.replay(rpcFactoryMock, rpcMock);
503+
DatastoreOptions options =
504+
this.options.toBuilder()
505+
.retryParams(RetryParams.getDefaultInstance())
506+
.serviceRpcFactory(rpcFactoryMock)
507+
.build();
508+
Datastore mockDatastore = DatastoreFactory.instance().get(options);
509+
QueryResults<Key> results = mockDatastore.run(Query.keyQueryBuilder().build());
510+
int count = 0;
511+
while (results.hasNext()) {
512+
count += 1;
513+
results.next();
514+
}
515+
assertEquals(count, 5);
516+
EasyMock.verify(rpcFactoryMock, rpcMock);
517+
}
518+
519+
private List<com.google.datastore.v1beta3.RunQueryResponse> buildResponsesForQueryPagination() {
520+
Entity entity4 = Entity.builder(KEY4).set("value", StringValue.of("value")).build();
521+
Entity entity5 = Entity.builder(KEY5).set("value", "value").build();
522+
datastore.add(ENTITY3, entity4, entity5);
523+
List<com.google.datastore.v1beta3.RunQueryResponse> responses = new ArrayList<>();
524+
Query<Key> query = Query.keyQueryBuilder().build();
525+
com.google.datastore.v1beta3.RunQueryRequest.Builder requestPb =
526+
com.google.datastore.v1beta3.RunQueryRequest.newBuilder();
527+
query.populatePb(requestPb);
528+
com.google.datastore.v1beta3.QueryResultBatch queryResultBatchPb =
529+
com.google.datastore.v1beta3.RunQueryResponse.newBuilder()
530+
.mergeFrom(((DatastoreImpl) datastore).runQuery(requestPb.build()))
531+
.getBatch();
532+
com.google.datastore.v1beta3.QueryResultBatch queryResultBatchPb1 =
533+
com.google.datastore.v1beta3.QueryResultBatch.newBuilder()
534+
.mergeFrom(queryResultBatchPb)
535+
.setMoreResults(
536+
com.google.datastore.v1beta3.QueryResultBatch.MoreResultsType.NOT_FINISHED)
537+
.clearEntityResults()
538+
.addAllEntityResults(queryResultBatchPb.getEntityResultsList().subList(0, 1))
539+
.setEndCursor(queryResultBatchPb.getEntityResultsList().get(0).getCursor())
540+
.build();
541+
responses.add(
542+
com.google.datastore.v1beta3.RunQueryResponse.newBuilder()
543+
.setBatch(queryResultBatchPb1)
544+
.build());
545+
com.google.datastore.v1beta3.QueryResultBatch queryResultBatchPb2 =
546+
com.google.datastore.v1beta3.QueryResultBatch.newBuilder()
547+
.mergeFrom(queryResultBatchPb)
548+
.setMoreResults(
549+
com.google.datastore.v1beta3.QueryResultBatch.MoreResultsType.NOT_FINISHED)
550+
.clearEntityResults()
551+
.addAllEntityResults(queryResultBatchPb.getEntityResultsList().subList(1, 3))
552+
.setEndCursor(queryResultBatchPb.getEntityResultsList().get(2).getCursor())
553+
.build();
554+
responses.add(
555+
com.google.datastore.v1beta3.RunQueryResponse.newBuilder()
556+
.setBatch(queryResultBatchPb2)
557+
.build());
558+
com.google.datastore.v1beta3.QueryResultBatch queryResultBatchPb3 =
559+
com.google.datastore.v1beta3.QueryResultBatch.newBuilder()
560+
.mergeFrom(queryResultBatchPb)
561+
.setMoreResults(
562+
com.google.datastore.v1beta3.QueryResultBatch.MoreResultsType.NO_MORE_RESULTS)
563+
.clearEntityResults()
564+
.addAllEntityResults(queryResultBatchPb.getEntityResultsList().subList(3, 5))
565+
.setEndCursor(queryResultBatchPb.getEntityResultsList().get(4).getCursor())
566+
.build();
567+
responses.add(
568+
com.google.datastore.v1beta3.RunQueryResponse.newBuilder()
569+
.setBatch(queryResultBatchPb3)
570+
.build());
571+
return responses;
453572
}
454573

455574
@Test

0 commit comments

Comments
 (0)