Fix: sort document reference by long type id (#1945)

milaGGL · web-flow · commit c38e3adc08e5 · 2024-12-17T11:11:52.000-05:00
diff --git a/README.md b/README.md
@@ -41,7 +41,7 @@ If you are using Maven without the BOM, add this to your dependencies:
 <dependency>
   <groupId>com.google.cloud</groupId>
   <artifactId>google-cloud-firestore</artifactId>
-  <version>3.30.1</version>
+  <version>3.30.2</version>
 </dependency>
 
 ```
diff --git a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/BasePath.java b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/BasePath.java
@@ -114,7 +114,11 @@ boolean isPrefixOf(BasePath<B> path) {
   }
 
   /**
-   * Compare the current path lexicographically against another Path object.
+   * Compare the current path against another Path object.
+   *
+   * <p>Compare the current path against another Path object. Paths are compared segment by segment,
+   * prioritizing numeric IDs (e.g., "__id123__") in numeric ascending order, followed by string
+   * segments in lexicographical order.
    *
    * @param other The path to compare to.
    * @return -1 if current is less than other, 1 if current greater than other, 0 if equal
@@ -123,16 +127,41 @@ boolean isPrefixOf(BasePath<B> path) {
   public int compareTo(@Nonnull B other) {
     List<String> thisSegments = this.getSegments();
     List<String> otherSegments = other.getSegments();
+
     int length = Math.min(thisSegments.size(), otherSegments.size());
     for (int i = 0; i < length; i++) {
-      int cmp = thisSegments.get(i).compareTo(otherSegments.get(i));
+      int cmp = compareSegments(thisSegments.get(i), otherSegments.get(i));
       if (cmp != 0) {
         return cmp;
       }
     }
     return Integer.compare(thisSegments.size(), otherSegments.size());
   }
 
+  private int compareSegments(String lhs, String rhs) {
+    boolean isLhsNumeric = isNumericId(lhs);
+    boolean isRhsNumeric = isNumericId(rhs);
+
+    if (isLhsNumeric && !isRhsNumeric) { // Only lhs is numeric
+      return -1;
+    } else if (!isLhsNumeric && isRhsNumeric) { // Only rhs is numeric
+      return 1;
+    } else if (isLhsNumeric && isRhsNumeric) { // both numeric
+      return Long.compare(extractNumericId(lhs), extractNumericId(rhs));
+    } else { // both string
+      return lhs.compareTo(rhs);
+    }
+  }
+
+  /** Checks if a segment is a numeric ID (starts with "__id" and ends with "__"). */
+  private boolean isNumericId(String segment) {
+    return segment.startsWith("__id") && segment.endsWith("__");
+  }
+
+  private long extractNumericId(String segment) {
+    return Long.parseLong(segment.substring(4, segment.length() - 2));
+  }
+
   abstract String[] splitChildPath(String path);
 
   abstract B createPathWithSegments(ImmutableList<String> segments);
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
@@ -24,17 +24,20 @@
 import static java.util.Collections.emptyList;
 import static java.util.Collections.singletonList;
 import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertEquals;
 
 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.FieldPath;
 import com.google.cloud.firestore.FieldValue;
 import com.google.cloud.firestore.FirestoreException;
 import com.google.cloud.firestore.ListenerRegistration;
 import com.google.cloud.firestore.LocalFirestoreHelper;
 import com.google.cloud.firestore.Query;
+import com.google.cloud.firestore.Query.Direction;
 import com.google.cloud.firestore.QueryDocumentSnapshot;
 import com.google.cloud.firestore.QuerySnapshot;
 import com.google.cloud.firestore.it.ITQueryWatchTest.QuerySnapshotEventListener.ListenerAssertions;
@@ -644,6 +647,118 @@ public void shutdownNowPreventsAddingNewListener() throws Exception {
     listenerAssertions.hasError();
   }
 
+  @Test
+  public void snapshotListenerSortsQueryByDocumentIdInTheSameOrderAsServer() throws Exception {
+    CollectionReference col = randomColl;
+
+    firestore
+        .batch()
+        .set(col.document("A"), Collections.singletonMap("a", 1))
+        .set(col.document("a"), Collections.singletonMap("a", 1))
+        .set(col.document("Aa"), Collections.singletonMap("a", 1))
+        .set(col.document("7"), Collections.singletonMap("a", 1))
+        .set(col.document("12"), Collections.singletonMap("a", 1))
+        .set(col.document("__id7__"), Collections.singletonMap("a", 1))
+        .set(col.document("__id12__"), Collections.singletonMap("a", 1))
+        .set(col.document("__id-2__"), Collections.singletonMap("a", 1))
+        .set(col.document("_id1__"), Collections.singletonMap("a", 1))
+        .set(col.document("__id1_"), Collections.singletonMap("a", 1))
+        .set(col.document("__id"), Collections.singletonMap("a", 1))
+        .commit()
+        .get();
+
+    Query query = col.orderBy("__name__", Direction.ASCENDING);
+    List<String> expectedOrder =
+        Arrays.asList(
+            "__id-2__",
+            "__id7__",
+            "__id12__",
+            "12",
+            "7",
+            "A",
+            "Aa",
+            "__id",
+            "__id1_",
+            "_id1__",
+            "a");
+
+    QuerySnapshot snapshot = query.get().get();
+    List<String> queryOrder =
+        snapshot.getDocuments().stream().map(doc -> doc.getId()).collect(Collectors.toList());
+    assertEquals(expectedOrder, queryOrder); // Assert order from backend
+
+    CountDownLatch latch = new CountDownLatch(1);
+    List<String> listenerOrder = new ArrayList<>();
+
+    ListenerRegistration registration =
+        query.addSnapshotListener(
+            (value, error) -> {
+              listenerOrder.addAll(
+                  value.getDocuments().stream()
+                      .map(doc -> doc.getId())
+                      .collect(Collectors.toList()));
+
+              latch.countDown();
+            });
+
+    latch.await();
+    registration.remove();
+
+    assertEquals(expectedOrder, listenerOrder); // Assert order in the SDK
+  }
+
+  @Test
+  public void snapshotListenerSortsFilteredQueryByDocumentIdInTheSameOrderAsServer()
+      throws Exception {
+    CollectionReference col = randomColl;
+
+    firestore
+        .batch()
+        .set(col.document("A"), Collections.singletonMap("a", 1))
+        .set(col.document("a"), Collections.singletonMap("a", 1))
+        .set(col.document("Aa"), Collections.singletonMap("a", 1))
+        .set(col.document("7"), Collections.singletonMap("a", 1))
+        .set(col.document("12"), Collections.singletonMap("a", 1))
+        .set(col.document("__id7__"), Collections.singletonMap("a", 1))
+        .set(col.document("__id12__"), Collections.singletonMap("a", 1))
+        .set(col.document("__id-2__"), Collections.singletonMap("a", 1))
+        .set(col.document("_id1__"), Collections.singletonMap("a", 1))
+        .set(col.document("__id1_"), Collections.singletonMap("a", 1))
+        .set(col.document("__id"), Collections.singletonMap("a", 1))
+        .commit()
+        .get();
+
+    Query query =
+        col.whereGreaterThan(FieldPath.documentId(), "__id7__")
+            .whereLessThanOrEqualTo(FieldPath.documentId(), "A")
+            .orderBy("__name__", Direction.ASCENDING);
+    List<String> expectedOrder = Arrays.asList("__id12__", "12", "7", "A");
+
+    QuerySnapshot snapshot = query.get().get();
+    List<String> queryOrder =
+        snapshot.getDocuments().stream().map(doc -> doc.getId()).collect(Collectors.toList());
+    assertEquals(expectedOrder, queryOrder); // Assert order from backend
+
+    CountDownLatch latch = new CountDownLatch(1);
+    List<String> listenerOrder = new ArrayList<>();
+
+    ListenerRegistration registration =
+        query.addSnapshotListener(
+            (value, error) -> {
+              listenerOrder.addAll(
+                  value.getDocuments().stream()
+                      .map(doc -> doc.getId())
+                      .collect(Collectors.toList()));
+
+              latch.countDown();
+            });
+
+    latch.await();
+    registration.remove();
+
+    assertEquals(expectedOrder, listenerOrder); // Assert order in the SDK
+  }
+
   /**
    * A tuple class used by {@code #queryWatch}. This class represents an event delivered to the
    * registered query listener.
