fix: query result start with startCursor if specified (#207)

googleapis · Aug 27, 2020 · 023229a · 023229a
1 parent 2cacc10
commit 023229a
Show file tree

Hide file tree

Showing 3 changed files with 44 additions and 4 deletions.
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResultsImpl.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResultsImpl.java
@@ -71,10 +71,7 @@ private void sendRequest() {
     requestPb.setPartitionId(partitionIdPb);
     query.populatePb(requestPb);
     runQueryResponsePb = datastore.runQuery(requestPb.build());
-    mostRecentQueryPb = runQueryResponsePb.getQuery();
-    if (mostRecentQueryPb == null) {
-      mostRecentQueryPb = requestPb.getQuery();
-    }
+    mostRecentQueryPb = requestPb.getQuery();
     moreResults = runQueryResponsePb.getBatch().getMoreResults();
     lastBatch = moreResults != MoreResultsType.NOT_FINISHED;
     entityResultPbIter = runQueryResponsePb.getBatch().getEntityResultsList().iterator();

diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/ITDatastoreTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/ITDatastoreTest.java
@@ -1242,4 +1242,25 @@ public void testGqlQueryWithNullBinding() {
     }
     assertEquals(1, count);
   }
+
+  @Test
+  public void testQueryWithStartCursor() {
+    Entity entity1 =
+        Entity.newBuilder(Key.newBuilder(PROJECT_ID, KIND1, "name-01").build()).build();
+    Entity entity2 =
+        Entity.newBuilder(Key.newBuilder(PROJECT_ID, KIND1, "name-02").build()).build();
+    Entity entity3 =
+        Entity.newBuilder(Key.newBuilder(PROJECT_ID, KIND1, "name-03").build()).build();
+    datastore.put(entity1, entity2, entity3);
+    QueryResults<Entity> run1 = datastore.run(Query.newEntityQueryBuilder().setKind(KIND1).build());
+    run1.next();
+    Cursor cursor1 = run1.getCursorAfter();
+    assertNotNull(cursor1);
+    QueryResults<Entity> run2 =
+        datastore.run(Query.newEntityQueryBuilder().setKind(KIND1).setStartCursor(cursor1).build());
+    Cursor cursor2 = run2.getCursorAfter();
+    assertNotNull(cursor2);
+    assertEquals(cursor2, cursor1);
+    datastore.delete(entity1.getKey(), entity2.getKey(), entity3.getKey());
+  }
 }
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java
@@ -28,6 +28,7 @@
 import com.google.cloud.Timestamp;
 import com.google.cloud.datastore.Batch;
 import com.google.cloud.datastore.BooleanValue;
+import com.google.cloud.datastore.Cursor;
 import com.google.cloud.datastore.Datastore;
 import com.google.cloud.datastore.DatastoreException;
 import com.google.cloud.datastore.DatastoreOptions;
@@ -898,4 +899,25 @@ public void testGqlQueryWithNullBinding() {
     assertEquals(ENTITY1, results.next());
     assertFalse(results.hasNext());
   }
+
+  @Test
+  public void testQueryWithStartCursor() {
+    Entity entity1 =
+        Entity.newBuilder(Key.newBuilder(PROJECT_ID, KIND1, "name-01").build()).build();
+    Entity entity2 =
+        Entity.newBuilder(Key.newBuilder(PROJECT_ID, KIND1, "name-02").build()).build();
+    Entity entity3 =
+        Entity.newBuilder(Key.newBuilder(PROJECT_ID, KIND1, "name-03").build()).build();
+    DATASTORE.put(entity1, entity2, entity3);
+    QueryResults<Entity> run1 = DATASTORE.run(Query.newEntityQueryBuilder().setKind(KIND1).build());
+    run1.next();
+    Cursor cursor1 = run1.getCursorAfter();
+    assertNotNull(cursor1);
+    QueryResults<Entity> run2 =
+        DATASTORE.run(Query.newEntityQueryBuilder().setKind(KIND1).setStartCursor(cursor1).build());
+    Cursor cursor2 = run2.getCursorAfter();
+    assertNotNull(cursor2);
+    assertEquals(cursor2, cursor1);
+    DATASTORE.delete(entity1.getKey(), entity2.getKey(), entity3.getKey());
+  }
 }