Skip to content

Commit 50a5a66

Browse files
author
Ajay Kannan
committed
Query continuation implemented for GqlQuery + tests
1 parent 8245607 commit 50a5a66

5 files changed

Lines changed: 142 additions & 28 deletions

File tree

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

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -379,9 +379,11 @@ 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 StructuredQuery<V> nextQuery(com.google.datastore.v1beta3.RunQueryResponse responsePb) {
383+
return new StructuredQuery.Builder<>(type())
384+
.mergeFrom(responsePb.getQuery())
385+
.prepareNext(responsePb.getBatch())
386+
.build();
385387
}
386388

387389
@Override

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

Lines changed: 1 addition & 1 deletion
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‎

Lines changed: 8 additions & 10 deletions
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‎

Lines changed: 16 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -669,9 +669,23 @@ B mergeFrom(com.google.datastore.v1beta3.Query queryPb) {
669669
return self();
670670
}
671671

672+
B prepareNext(com.google.datastore.v1beta3.QueryResultBatch queryResultBatchPb) {
673+
startCursor(new Cursor(queryResultBatchPb.getEndCursor()));
674+
if (offset > 0 && queryResultBatchPb.getSkippedResults() < offset) {
675+
offset(offset - queryResultBatchPb.getSkippedResults());
676+
} else {
677+
offset(0);
678+
if (limit != null) {
679+
limit(limit - queryResultBatchPb.getEntityResultsCount());
680+
}
681+
}
682+
return self();
683+
}
684+
672685
public StructuredQuery<V> build() {
673686
return new StructuredQuery<>(this);
674687
}
688+
675689
}
676690

677691
static final class Builder<V> extends BaseBuilder<V, Builder<V>> {
@@ -844,19 +858,8 @@ protected void populatePb(com.google.datastore.v1beta3.RunQueryRequest.Builder r
844858
}
845859

846860
@Override
847-
protected StructuredQuery<V> nextQuery(com.google.datastore.v1beta3.QueryResultBatch responsePb) {
848-
Builder<V> builder = new Builder<>(type());
849-
builder.mergeFrom(toPb());
850-
builder.startCursor(new Cursor(responsePb.getEndCursor()));
851-
if (offset > 0 && responsePb.getSkippedResults() < offset) {
852-
builder.offset(offset - responsePb.getSkippedResults());
853-
} else {
854-
builder.offset(0);
855-
if (limit != null) {
856-
builder.limit(limit - responsePb.getEntityResultsCount());
857-
}
858-
}
859-
return builder.build();
861+
protected StructuredQuery<V> nextQuery(com.google.datastore.v1beta3.RunQueryResponse responsePb) {
862+
return new Builder<>(type()).mergeFrom(toPb()).prepareNext(responsePb.getBatch()).build();
860863
}
861864

862865
@Override

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

Lines changed: 112 additions & 1 deletion
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;
@@ -403,6 +405,42 @@ public void testRunGqlQueryWithCasting() {
403405
assertFalse(results3.hasNext());
404406
}
405407

408+
@Test
409+
public void testGqlQueryPagination() {
410+
DatastoreRpcFactory rpcFactoryMock = EasyMock.createStrictMock(DatastoreRpcFactory.class);
411+
DatastoreRpc rpcMock = EasyMock.createStrictMock(DatastoreRpc.class);
412+
EasyMock.expect(rpcFactoryMock.create(EasyMock.anyObject(DatastoreOptions.class)))
413+
.andReturn(rpcMock);
414+
List<com.google.datastore.v1beta3.RunQueryResponse> responses =
415+
buildResponsesForQueryPagination();
416+
for (int i = 0; i < responses.size(); i++) {
417+
try {
418+
EasyMock
419+
.expect(rpcMock.runQuery(
420+
EasyMock.anyObject(com.google.datastore.v1beta3.RunQueryRequest.class)))
421+
.andReturn(responses.get(i));
422+
} catch (DatastoreRpcException e) {
423+
fail("Unexpected DatastoreRpcException");
424+
}
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, responses.size());
441+
EasyMock.verify(rpcFactoryMock, rpcMock);
442+
}
443+
406444
@Test
407445
public void testRunStructuredQuery() {
408446
Query<Entity> query =
@@ -445,7 +483,80 @@ public void testRunStructuredQuery() {
445483
assertEquals(20, entity.getLong("age"));
446484
assertEquals(1, entity.properties().size());
447485
assertFalse(results4.hasNext());
448-
// TODO(ozarov): construct a test to verify nextQuery/pagination
486+
}
487+
488+
@Test
489+
public void testStructuredQueryPagination() {
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+
try {
498+
EasyMock
499+
.expect(rpcMock.runQuery(
500+
EasyMock.anyObject(com.google.datastore.v1beta3.RunQueryRequest.class)))
501+
.andReturn(responses.get(i));
502+
} catch (DatastoreRpcException e) {
503+
fail("Unexpected DatastoreRpcException");
504+
}
505+
}
506+
EasyMock.replay(rpcFactoryMock, rpcMock);
507+
DatastoreOptions options =
508+
this.options.toBuilder()
509+
.retryParams(RetryParams.getDefaultInstance())
510+
.serviceRpcFactory(rpcFactoryMock)
511+
.build();
512+
Datastore mockDatastore = DatastoreFactory.instance().get(options);
513+
QueryResults<Key> results = mockDatastore.run(Query.keyQueryBuilder().build());
514+
int count = 0;
515+
while (results.hasNext()) {
516+
count += 1;
517+
results.next();
518+
}
519+
assertEquals(count, responses.size());
520+
EasyMock.verify(rpcFactoryMock, rpcMock);
521+
}
522+
523+
private List<com.google.datastore.v1beta3.RunQueryResponse> buildResponsesForQueryPagination() {
524+
List<com.google.datastore.v1beta3.RunQueryResponse> responses = new ArrayList<>();
525+
Query<Key> query = Query.keyQueryBuilder().build();
526+
com.google.datastore.v1beta3.RunQueryRequest.Builder requestPb =
527+
com.google.datastore.v1beta3.RunQueryRequest.newBuilder();
528+
query.populatePb(requestPb);
529+
com.google.datastore.v1beta3.QueryResultBatch queryResultBatchPb =
530+
com.google.datastore.v1beta3.RunQueryResponse.newBuilder()
531+
.mergeFrom(((DatastoreImpl) datastore).runQuery(requestPb.build()))
532+
.getBatch();
533+
com.google.datastore.v1beta3.QueryResultBatch queryResultBatchPb1 =
534+
com.google.datastore.v1beta3.QueryResultBatch.newBuilder()
535+
.mergeFrom(queryResultBatchPb)
536+
.setMoreResults(
537+
com.google.datastore.v1beta3.QueryResultBatch.MoreResultsType.NOT_FINISHED)
538+
.clearEntityResults()
539+
.addAllEntityResults(queryResultBatchPb.getEntityResultsList().subList(0, 1))
540+
.setEndCursor(queryResultBatchPb.getEntityResultsList().get(0).getCursor())
541+
.build();
542+
responses.add(
543+
com.google.datastore.v1beta3.RunQueryResponse.newBuilder()
544+
.setBatch(queryResultBatchPb1)
545+
.build());
546+
com.google.datastore.v1beta3.QueryResultBatch queryResultBatchPb2 =
547+
com.google.datastore.v1beta3.QueryResultBatch.newBuilder()
548+
.mergeFrom(queryResultBatchPb)
549+
.setMoreResults(
550+
com.google.datastore.v1beta3.QueryResultBatch.MoreResultsType.NO_MORE_RESULTS)
551+
.clearEntityResults()
552+
.addAllEntityResults(queryResultBatchPb.getEntityResultsList().subList(1, 2))
553+
.setEndCursor(queryResultBatchPb.getEntityResultsList().get(1).getCursor())
554+
.build();
555+
responses.add(
556+
com.google.datastore.v1beta3.RunQueryResponse.newBuilder()
557+
.setBatch(queryResultBatchPb2)
558+
.build());
559+
return responses;
449560
}
450561

451562
@Test

0 commit comments

Comments
 (0)