Add API to filter overlays by batch ID (#3145)

Author: schmidt-sebastian

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

@@ -46,6 +46,12 @@ public interface DocumentOverlayCache {
   /** Removes the overlay whose largest-batch-id equals to the given Id. */
   void removeOverlaysForBatchId(int batchId);
 
-  /** Returns all saved overlay for the given collection. */
-  Map<DocumentKey, Mutation> getOverlays(ResourcePath collection);
+  /**
+   * Returns all saved overlays for the given collection.
+   *
+   * @param collection The collection path 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.
+   */
+  Map<DocumentKey, Mutation> getOverlays(ResourcePath collection, int sinceBatchId);
 }

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

@@ -323,7 +323,7 @@ private ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingCollection
           Query query, SnapshotVersion sinceReadTime) {
     ImmutableSortedMap<DocumentKey, MutableDocument> remoteDocuments =
         remoteDocumentCache.getAllDocumentsMatchingQuery(query, sinceReadTime);
-    Map<DocumentKey, Mutation> overlays = documentOverlayCache.getOverlays(query.getPath());
+    Map<DocumentKey, Mutation> overlays = documentOverlayCache.getOverlays(query.getPath(), -1);
 
     // As documents might match the query because of their overlay we need to include all documents
     // in the result.

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

@@ -28,16 +28,16 @@
 public class MemoryDocumentOverlayCache implements DocumentOverlayCache {
   // A map sorted by DocumentKey, whose value is a pair of the largest batch id for the overlay
   // and the overlay itself.
-  private TreeMap<DocumentKey, Pair<Integer, Mutation>> overlays = new TreeMap<>();
-  Map<Integer, Set<DocumentKey>> overlayByBatchId = new HashMap<>();
+  private final TreeMap<DocumentKey, Pair<Integer, Mutation>> overlays = new TreeMap<>();
+  private final Map<Integer, Set<DocumentKey>> overlayByBatchId = new HashMap<>();
 
   @Nullable
   @Override
   public Mutation getOverlay(DocumentKey key) {
-    if (overlays.get(key) != null) {
-      return overlays.get(key).second;
+    Pair<Integer, Mutation> overlay = overlays.get(key);
+    if (overlay != null) {
+      return overlay.second;
     }
-
     return null;
   }
 
@@ -80,7 +80,7 @@ public void removeOverlaysForBatchId(int batchId) {
   }
 
   @Override
-  public Map<DocumentKey, Mutation> getOverlays(ResourcePath collection) {
+  public Map<DocumentKey, Mutation> getOverlays(ResourcePath collection, int sinceBatchId) {
     Map<DocumentKey, Mutation> result = new HashMap<>();
 
     int immediateChildrenPathLength = collection.length() + 1;
@@ -96,7 +96,11 @@ public Map<DocumentKey, Mutation> getOverlays(ResourcePath collection) {
       if (key.getPath().length() != immediateChildrenPathLength) {
         continue;
       }
-      result.put(entry.getKey(), entry.getValue().second);
+
+      Pair<Integer, Mutation> batchIdToOverlay = entry.getValue();
+      if (batchIdToOverlay.first > sinceBatchId) {
+        result.put(entry.getKey(), batchIdToOverlay.second);
+      }
     }
 
     return result;

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

@@ -82,7 +82,7 @@ public void removeOverlaysForBatchId(int batchId) {
   }
 
   @Override
-  public Map<DocumentKey, Mutation> getOverlays(ResourcePath collection) {
+  public Map<DocumentKey, Mutation> getOverlays(ResourcePath collection, int sinceBatchId) {
     int immediateChildrenPathLength = collection.length() + 1;
 
     String prefixPath = EncodedPath.encode(collection);
@@ -91,8 +91,9 @@ public Map<DocumentKey, Mutation> getOverlays(ResourcePath collection) {
     Map<DocumentKey, Mutation> result = new HashMap<>();
 
     db.query(
-            "SELECT path, overlay_mutation FROM document_overlays WHERE uid = ? AND path >= ? AND path < ?")
-        .binding(uid, prefixPath, prefixSuccessorPath)
+            "SELECT path, overlay_mutation FROM document_overlays "
+                + "WHERE uid = ? AND path >= ? AND path < ? AND largest_batch_id > ?")
+        .binding(uid, prefixPath, prefixSuccessorPath, sinceBatchId)
         .forEach(
             row -> {
               try {

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

@@ -152,6 +152,31 @@ public void testGetAllOverlaysForCollection() {
 
     m.remove(key("coll/doc1/sub/sub_doc"));
     m.remove(key("other/doc1"));
-    assertEquals(m, overlays.getOverlays(path("coll")));
+    assertEquals(m, overlays.getOverlays(path("coll"), -1));
+  }
+
+  @Test
+  public void testGetAllOverlaysSinceBatchId() {
+    Mutation m1 = patchMutation("coll/doc1", map("foo", "bar"));
+    Mutation m2 = setMutation("coll/doc2", map("foo", "bar"));
+    Map<DocumentKey, Mutation> m = new HashMap<>();
+    m.put(key("coll/doc1"), m1);
+    m.put(key("coll/doc2"), m2);
+    overlays.saveOverlays(2, m);
+
+    Mutation m3 = deleteMutation("coll/doc3");
+    m = new HashMap<>();
+    m.put(key("coll/doc3"), m3);
+    overlays.saveOverlays(3, m);
+
+    Mutation m4 = deleteMutation("coll/doc4");
+    m = new HashMap<>();
+    m.put(key("coll/doc4"), m4);
+    overlays.saveOverlays(4, m);
+
+    Map<DocumentKey, Mutation> expected = new HashMap<>();
+    expected.put(key("coll/doc3"), m3);
+    expected.put(key("coll/doc4"), m4);
+    assertEquals(expected, overlays.getOverlays(path("coll"), 2));
   }
 }