fix: Fix limitToLast queries with cursors (#1072)

googleapis · Dec 11, 2023 · 5cb10f5 · 5cb10f5
1 parent 8cd533c
commit 5cb10f5
Show file tree

Hide file tree

Showing 2 changed files with 113 additions and 4 deletions.
diff --git a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/Query.java b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/Query.java
@@ -1289,10 +1289,12 @@ public Query endAt(@Nonnull DocumentSnapshot snapshot) {
   StructuredQuery.Builder buildQuery() {
     StructuredQuery.Builder structuredQuery = buildWithoutClientTranslation();
     if (options.getLimitType().equals(LimitType.Last)) {
-      // Apply client translation for limitToLast.
+      structuredQuery.clearOrderBy();
+      structuredQuery.clearStartAt();
+      structuredQuery.clearEndAt();
 
+      // Apply client translation for limitToLast.
       if (!options.getFieldOrders().isEmpty()) {
-        structuredQuery.clearOrderBy();
         for (FieldOrder order : options.getFieldOrders()) {
           // Flip the orderBy directions since we want the last results
           order =
@@ -1306,7 +1308,6 @@ StructuredQuery.Builder buildQuery() {
       }
 
       if (options.getStartCursor() != null) {
-        structuredQuery.clearEndAt();
         // Swap the cursors to match the flipped query ordering.
         Cursor cursor =
             options
@@ -1318,7 +1319,6 @@ StructuredQuery.Builder buildQuery() {
       }
 
       if (options.getEndCursor() != null) {
-        structuredQuery.clearStartAt();
         // Swap the cursors to match the flipped query ordering.
         Cursor cursor =
             options

diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITQueryWatchTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITQueryWatchTest.java
@@ -17,6 +17,7 @@
 package com.google.cloud.firestore.it;
 
 import static com.google.cloud.firestore.LocalFirestoreHelper.map;
+import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth.assertWithMessage;
 import static java.util.Arrays.asList;
@@ -26,6 +27,7 @@
 import com.google.cloud.firestore.CollectionReference;
 import com.google.cloud.firestore.DocumentChange;
 import com.google.cloud.firestore.DocumentReference;
+import com.google.cloud.firestore.DocumentSnapshot;
 import com.google.cloud.firestore.EventListener;
 import com.google.cloud.firestore.Firestore;
 import com.google.cloud.firestore.FirestoreException;
@@ -39,6 +41,7 @@
 import com.google.common.base.Joiner;
 import com.google.common.base.Joiner.MapJoiner;
 import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Range;
 import com.google.common.truth.Truth;
 import java.util.ArrayList;
@@ -381,6 +384,112 @@ public void limitToLast() throws Exception {
     listenerAssertions.addedIdsIsAnyOf(emptyList(), asList("doc2", "doc3"));
   }
 
+  /**
+   * Verifies that QuerySnapshot for limitToLast() queries work with startAt when the full limit is
+   * used in the result set.
+   */
+  @Test
+  public void limitToLastWithStartAtFullLimit() throws Exception {
+    for (int i = 0; i < 10; i++) {
+      setDocument("doc" + i, Collections.singletonMap("counter", i));
+    }
+    Query query = randomColl.orderBy("counter").startAt(5).limitToLast(3);
+    assertQueryResultContainsDocsInOrder(query, "doc7", "doc8", "doc9");
+  }
+
+  /**
+   * Verifies that QuerySnapshot for limitToLast() queries work with startAt when the partial limit
+   * is used in the result set.
+   */
+  @Test
+  public void limitToLastWithStartAtPartialLimit() throws Exception {
+    for (int i = 0; i < 10; i++) {
+      setDocument("doc" + i, Collections.singletonMap("counter", i));
+    }
+    Query query = randomColl.orderBy("counter").startAt(8).limitToLast(3);
+    assertQueryResultContainsDocsInOrder(query, "doc8", "doc9");
+  }
+
+  /**
+   * Verifies that QuerySnapshot for limitToLast() queries work with startAfter when the full limit
+   * is used in the result set.
+   */
+  @Test
+  public void limitToLastWithStartAfterFullLimit() throws Exception {
+    for (int i = 0; i < 10; i++) {
+      setDocument("doc" + i, Collections.singletonMap("counter", i));
+    }
+    Query query = randomColl.orderBy("counter").startAfter(5).limitToLast(3);
+    assertQueryResultContainsDocsInOrder(query, "doc7", "doc8", "doc9");
+  }
+
+  /**
+   * Verifies that QuerySnapshot for limitToLast() queries work with startAfter when the partial
+   * limit is used in the result set.
+   */
+  @Test
+  public void limitToLastWithStartAfterPartialLimit() throws Exception {
+    for (int i = 0; i < 10; i++) {
+      setDocument("doc" + i, Collections.singletonMap("counter", i));
+    }
+    Query query = randomColl.orderBy("counter").startAfter(7).limitToLast(3);
+    assertQueryResultContainsDocsInOrder(query, "doc8", "doc9");
+  }
+
+  /** Verifies that QuerySnapshot for limitToLast() queries work with endAt. */
+  @Test
+  public void limitToLastWithEndAt() throws Exception {
+    for (int i = 0; i < 10; i++) {
+      setDocument("doc" + i, Collections.singletonMap("counter", i));
+    }
+    Query query = randomColl.orderBy("counter").endAt(5).limitToLast(3);
+    assertQueryResultContainsDocsInOrder(query, "doc3", "doc4", "doc5");
+  }
+
+  /** Verifies that QuerySnapshot for limitToLast() queries work with endBefore. */
+  @Test
+  public void limitToLastWithEndBefore() throws Exception {
+    for (int i = 0; i < 10; i++) {
+      setDocument("doc" + i, Collections.singletonMap("counter", i));
+    }
+    Query query = randomColl.orderBy("counter").endBefore(5).limitToLast(3);
+    assertQueryResultContainsDocsInOrder(query, "doc2", "doc3", "doc4");
+  }
+
+  /**
+   * Verifies that QuerySnapshot for limitToLast() queries work with both startAt and endAt when the
+   * full limit is used in the result set.
+   */
+  @Test
+  public void limitToLastWithStartAtAndEndAtFullLimit() throws Exception {
+    for (int i = 0; i < 10; i++) {
+      setDocument("doc" + i, Collections.singletonMap("counter", i));
+    }
+    Query query = randomColl.orderBy("counter").startAt(3).endAt(6).limitToLast(3);
+    assertQueryResultContainsDocsInOrder(query, "doc4", "doc5", "doc6");
+  }
+
+  /**
+   * Verifies that QuerySnapshot for limitToLast() queries work with both startAt and endAt when the
+   * partial limit is used in the result set.
+   */
+  @Test
+  public void limitToLastWithStartAtAndEndAtPartialLimit() throws Exception {
+    for (int i = 0; i < 10; i++) {
+      setDocument("doc" + i, Collections.singletonMap("counter", i));
+    }
+    Query query = randomColl.orderBy("counter").startAt(5).endAt(6).limitToLast(3);
+    assertQueryResultContainsDocsInOrder(query, "doc5", "doc6");
+  }
+
+  private static void assertQueryResultContainsDocsInOrder(Query query, String... docIds)
+      throws ExecutionException, InterruptedException {
+    QuerySnapshot snapshot = query.get().get();
+    ImmutableList<String> actualDocIds =
+        snapshot.getDocuments().stream().map(DocumentSnapshot::getId).collect(toImmutableList());
+    assertThat(actualDocIds).containsExactlyElementsIn(docIds).inOrder();
+  }
+
   @Test
   public void shutdownNowTerminatesActiveListener() throws Exception {
     Query query = randomColl.whereEqualTo("foo", "bar");