Fix overlay bug when offline (#3537)

* Fix: #3528

* Fix test failure

* Simpler fix

* Ugly fix fixes

Author: wu-hui

firebase-firestore/CHANGELOG.md

@@ -2,12 +2,13 @@ Android changes are not released automatically. Ensure that changes are released
 by opting into a release at 
 [go/firebase-android-release](http:go/firebase-android-release) (Googlers only).
 
-# Unreleased
+# 24.1.0
 - [feature] Added experimental support for indexed query execution. Indexes can
   be enabled by invoking `FirebaseFirestore.setIndexConfiguration()` with the
   JSON index definition exported by the Firestore CLI. Queries against the
   cache are executed using an index once the asynchronous operation to generate
   the index entries completes.
+- [fixed] Fixed missing document fields issue with offline overlays (#3528)
 
 # 24.0.2
 - [fixed] Fixed an issue of long grpc reconnection period, when App moves to

firebase-firestore/src/androidTest/java/com/google/firebase/firestore/QueryTest.java

@@ -884,4 +884,40 @@ public void testCollectionGroupQueriesWithWhereFiltersOnArbitraryDocumentIds() {
                 .get());
     assertEquals(asList("cg-doc2"), querySnapshotToIds(querySnapshot));
   }
+
+  // See: https://github.com/firebase/firebase-android-sdk/issues/3528
+  // TODO(Overlay): These two tests should be part of local store tests instead.
+  @Test
+  public void testAddThenUpdatesWhileOffline() {
+    CollectionReference collection = testCollection();
+    collection.getFirestore().disableNetwork();
+
+    collection.add(map("foo", "zzyzx", "bar", "1"));
+
+    QuerySnapshot snapshot1 = waitFor(collection.get(Source.CACHE));
+    assertEquals(asList(map("foo", "zzyzx", "bar", "1")), querySnapshotToValues(snapshot1));
+    DocumentReference doc = snapshot1.getDocuments().get(0).getReference();
+
+    doc.update(map("bar", "2"));
+
+    QuerySnapshot snapshot2 = waitFor(collection.get(Source.CACHE));
+    assertEquals(asList(map("foo", "zzyzx", "bar", "2")), querySnapshotToValues(snapshot2));
+  }
+
+  @Test
+  public void testMultipleUpdatesWhileOffline() {
+    CollectionReference collection = testCollection();
+    collection.getFirestore().disableNetwork();
+
+    DocumentReference doc = collection.document();
+    doc.set(map("foo", "zzyzx", "bar", "1"), SetOptions.mergeFields("foo", "bar"));
+
+    QuerySnapshot snapshot1 = waitFor(collection.get(Source.CACHE));
+    assertEquals(asList(map("foo", "zzyzx", "bar", "1")), querySnapshotToValues(snapshot1));
+
+    doc.update(map("bar", "2"));
+
+    QuerySnapshot snapshot2 = waitFor(collection.get(Source.CACHE));
+    assertEquals(asList(map("foo", "zzyzx", "bar", "2")), querySnapshotToValues(snapshot2));
+  }
 }

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

@@ -122,6 +122,19 @@ ImmutableSortedMap<DocumentKey, Document> getLocalViewOfDocuments(
     return computeViews(docs, overlays, existenceStateChanged);
   }
 
+  /**
+   * Similar to {@link #getDocuments}, but creates the local view from the given {@code baseDocs}
+   * without retrieving documents from the local store.
+   *
+   * @param docs The documents to apply local mutations to get the local views.
+   */
+  ImmutableSortedMap<DocumentKey, Document> getLocalViewOfDocuments(
+      Map<DocumentKey, MutableDocument> docs) {
+    Map<DocumentKey, Overlay> overlays = new HashMap<>();
+    populateOverlays(overlays, docs.keySet());
+    return computeViews(docs, overlays, new HashSet<>());
+  }
+
   /*Computes the local view for doc */
   private ImmutableSortedMap<DocumentKey, Document> computeViews(
       Map<DocumentKey, MutableDocument> docs,

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

@@ -244,9 +244,22 @@ public LocalDocumentsResult writeLocally(List<Mutation> mutations) {
     return persistence.runTransaction(
         "Locally write mutations",
         () -> {
+          // Figure out which keys do not have a remote version in the cache, this is needed to
+          // create the right overlay mutation: if no remote version presents, we do not need to
+          // create overlays as patch mutations.
+          // TODO(Overlay): Is there a better way to determine this? Document version does not work
+          // because local mutations set them back to 0.
+          Map<DocumentKey, MutableDocument> remoteDocs = remoteDocuments.getAll(keys);
+          Set<DocumentKey> docsWithoutRemoteVersion = new HashSet<>();
+          for (Map.Entry<DocumentKey, MutableDocument> entry : remoteDocs.entrySet()) {
+            if (!entry.getValue().isValidDocument()) {
+              docsWithoutRemoteVersion.add(entry.getKey());
+            }
+          }
           // Load and apply all existing mutations. This lets us compute the current base state for
           // all non-idempotent transforms before applying any additional user-provided writes.
-          ImmutableSortedMap<DocumentKey, Document> documents = localDocuments.getDocuments(keys);
+          ImmutableSortedMap<DocumentKey, Document> documents =
+              localDocuments.getLocalViewOfDocuments(remoteDocs);
 
           // For non-idempotent mutations (such as `FieldValue.increment()`), we record the base
           // state in a separate patch mutation. This is later used to guarantee consistent values
@@ -270,7 +283,8 @@ public LocalDocumentsResult writeLocally(List<Mutation> mutations) {
 
           MutationBatch batch =
               mutationQueue.addMutationBatch(localWriteTime, baseMutations, mutations);
-          Map<DocumentKey, Mutation> overlays = batch.applyToLocalDocumentSet(documents);
+          Map<DocumentKey, Mutation> overlays =
+              batch.applyToLocalDocumentSet(documents, docsWithoutRemoteVersion);
           documentOverlayCache.saveOverlays(batch.getBatchId(), overlays);
           return new LocalDocumentsResult(batch.getBatchId(), documents);
         });

firebase-firestore/src/main/java/com/google/firebase/firestore/model/mutation/MutationBatch.java

@@ -136,7 +136,8 @@ public FieldMask applyToLocalView(MutableDocument document, @Nullable FieldMask
    * applications.
    */
   public Map<DocumentKey, Mutation> applyToLocalDocumentSet(
-      ImmutableSortedMap<DocumentKey, Document> documentMap) {
+      ImmutableSortedMap<DocumentKey, Document> documentMap,
+      Set<DocumentKey> documentsWithoutRemoteVersion) {
     // TODO(mrschmidt): This implementation is O(n^2). If we iterate through the mutations first
     // (as done in `applyToLocalView(MutableDocument d)`), we can reduce the complexity to
     // O(n).
@@ -146,6 +147,9 @@ public Map<DocumentKey, Mutation> applyToLocalDocumentSet(
       // remove this cast.
       MutableDocument document = (MutableDocument) documentMap.get(key);
       FieldMask mutatedFields = applyToLocalView(document);
+      // Set mutationFields to null if the document is only from local mutations, this creates
+      // a Set(or Delete) mutation, instead of trying to create a patch mutation as the overlay.
+      mutatedFields = documentsWithoutRemoteVersion.contains(key) ? null : mutatedFields;
       Mutation overlay = Mutation.calculateOverlayMutation(document, mutatedFields);
       overlays.put(key, overlay);
       if (!document.isValidDocument()) {

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

@@ -759,7 +759,7 @@ public void testHandlesDeleteMutationThenPatchMutationThenAckThenAck() {
 
     acknowledgeMutation(2); // delete mutation
     assertRemoved("foo/bar");
-    assertContains(deletedDoc("foo/bar", 2).setHasCommittedMutations());
+    assertContains(deletedDoc("foo/bar", 0).setHasLocalMutations());
 
     acknowledgeMutation(3); // patch mutation
     assertChanged(unknownDoc("foo/bar", 3));