Add getOverlays(CollectionGroup) (#3330)

Author: schmidt-sebastian

firebase-firestore/src/main/java/com/google/firebase/firestore/local/DocumentOverlayCache.java

@@ -56,4 +56,18 @@ public interface DocumentOverlayCache {
    * @return Mapping of each document key in the collection to its overlay.
    */
   Map<DocumentKey, Overlay> getOverlays(ResourcePath collection, int sinceBatchId);
+
+  /**
+   * Returns {@code count} overlays with a batch ID higher than {@code sinceBatchId} for the
+   * provided collection group, processed by ascending batch ID. The method always returns all
+   * overlays for a batch even if the last batch contains more documents than the remaining limit.
+   *
+   * @param collectionGroup The collection group to get the overlays for.
+   * @param sinceBatchId The minimum batch ID to filter by (exclusive). Only overlays that contain a
+   *     change past `sinceBatchId` are returned.
+   * @param count The number of overlays to return. Can be exceeded if the last batch contains more
+   *     entries.
+   * @return Mapping of each document key in the collection group to its overlay.
+   */
+  Map<DocumentKey, Overlay> getOverlays(String collectionGroup, int sinceBatchId, int count);
 }

firebase-firestore/src/main/java/com/google/firebase/firestore/local/MemoryDocumentOverlayCache.java

@@ -23,6 +23,7 @@
 import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
+import java.util.SortedMap;
 import java.util.TreeMap;
 
 public class MemoryDocumentOverlayCache implements DocumentOverlayCache {
@@ -100,4 +101,35 @@ public Map<DocumentKey, Overlay> getOverlays(ResourcePath collection, int sinceB
 
     return result;
   }
+
+  @Override
+  public Map<DocumentKey, Overlay> getOverlays(
+      String collectionGroup, int sinceBatchId, int count) {
+    SortedMap<Integer, Map<DocumentKey, Overlay>> batchIdToOverlays = new TreeMap<>();
+
+    for (Overlay overlay : overlays.values()) {
+      DocumentKey key = overlay.getKey();
+      if (!key.getCollectionGroup().equals(collectionGroup)) {
+        continue;
+      }
+      if (overlay.getLargestBatchId() > sinceBatchId) {
+        Map<DocumentKey, Overlay> overlays = batchIdToOverlays.get(overlay.getLargestBatchId());
+        if (overlays == null) {
+          overlays = new HashMap<>();
+          batchIdToOverlays.put(overlay.getLargestBatchId(), overlays);
+        }
+        overlays.put(overlay.getKey(), overlay);
+      }
+    }
+
+    Map<DocumentKey, Overlay> result = new HashMap<>();
+    for (Map<DocumentKey, Overlay> overlays : batchIdToOverlays.values()) {
+      result.putAll(overlays);
+      if (result.size() >= count) {
+        break;
+      }
+    }
+
+    return result;
+  }
 }

firebase-firestore/src/main/java/com/google/firebase/firestore/local/SQLiteDocumentOverlayCache.java

@@ -99,6 +99,51 @@ public Map<DocumentKey, Overlay> getOverlays(ResourcePath collection, int sinceB
     return result;
   }
 
+  @Override
+  public Map<DocumentKey, Overlay> getOverlays(
+      String collectionGroup, int sinceBatchId, int count) {
+    Map<DocumentKey, Overlay> result = new HashMap<>();
+    Overlay[] lastOverlay = new Overlay[] {null};
+
+    db.query(
+            "SELECT overlay_mutation, largest_batch_id FROM document_overlays "
+                + "WHERE uid = ? AND collection_group = ? AND largest_batch_id > ? "
+                + "ORDER BY largest_batch_id, collection_path, document_id LIMIT ?")
+        .binding(uid, collectionGroup, sinceBatchId, count)
+        .forEach(
+            row -> {
+              lastOverlay[0] = decodeOverlay(row);
+              result.put(lastOverlay[0].getKey(), lastOverlay[0]);
+            });
+
+    if (lastOverlay[0] == null) {
+      return result;
+    }
+
+    // Finish batch
+    DocumentKey key = lastOverlay[0].getKey();
+    String encodedCollectionPath = EncodedPath.encode(key.getCollectionPath());
+    db.query(
+            "SELECT overlay_mutation, largest_batch_id FROM document_overlays "
+                + "WHERE uid = ? AND collection_group = ? "
+                + "AND (collection_path > ? OR (collection_path = ? AND document_id > ?)) "
+                + "AND largest_batch_id = ?")
+        .binding(
+            uid,
+            collectionGroup,
+            encodedCollectionPath,
+            encodedCollectionPath,
+            key.getDocumentId(),
+            lastOverlay[0].getLargestBatchId())
+        .forEach(
+            row -> {
+              Overlay overlay = decodeOverlay(row);
+              result.put(overlay.getKey(), overlay);
+            });
+
+    return result;
+  }
+
   private Overlay decodeOverlay(android.database.Cursor row) {
     try {
       Write write = Write.parseFrom(row.getBlob(0));

firebase-firestore/src/test/java/com/google/firebase/firestore/local/DocumentOverlayCacheTestCase.java

@@ -153,6 +153,44 @@ public void testGetAllOverlaysSinceBatchId() {
     verifyOverlayContains(overlays, "coll/doc3", "coll/doc4");
   }
 
+  @Test
+  public void testGetAllOverlaysFromCollectionGroupEnforcesCollectionGroup() {
+    saveOverlays(2, "coll1/doc1", "coll2/doc1");
+    saveOverlays(3, "coll1/doc2");
+    saveOverlays(4, "coll2/doc2");
+
+    Map<DocumentKey, Overlay> overlays = cache.getOverlays("coll1", -1, 50);
+    verifyOverlayContains(overlays, "coll1/doc1", "coll1/doc2");
+  }
+
+  @Test
+  public void testGetAllOverlaysFromCollectionGroupEnforcesBatchId() {
+    saveOverlays(2, "coll/doc1");
+    saveOverlays(3, "coll/doc2");
+
+    Map<DocumentKey, Overlay> overlays = cache.getOverlays("coll", 2, 50);
+    verifyOverlayContains(overlays, "coll/doc2");
+  }
+
+  @Test
+  public void testGetAllOverlaysFromCollectionGroupEnforcesLimit() {
+    saveOverlays(1, "coll/doc1");
+    saveOverlays(2, "coll/doc2");
+    saveOverlays(3, "coll/doc3");
+
+    Map<DocumentKey, Overlay> overlays = cache.getOverlays("coll", -1, 2);
+    verifyOverlayContains(overlays, "coll/doc1", "coll/doc2");
+  }
+
+  @Test
+  public void testGetAllOverlaysFromCollectionGroupWithLimitIncludesFullBatches() {
+    saveOverlays(1, "coll/doc1");
+    saveOverlays(2, "coll/doc2", "coll/doc3");
+
+    Map<DocumentKey, Overlay> overlays = cache.getOverlays("coll", -1, 2);
+    verifyOverlayContains(overlays, "coll/doc1", "coll/doc2", "coll/doc3");
+  }
+
   void verifyOverlayContains(Map<DocumentKey, Overlay> overlays, String... keys) {
     Set<DocumentKey> expected = Arrays.stream(keys).map(TestUtil::key).collect(Collectors.toSet());
     assertThat(overlays.keySet()).containsExactlyElementsIn(expected);