Fix missing query results when it is a match after a local patch mutation. (#1957)

When a local patch mutation is not acked by server yet, but the mutation makes a document matching a query, the document might not show up in the compensated query results. This PR fixes the bug.

See firebase/firebase-android-sdk#155

Author: Hui-Wu

packages/firestore/src/local/local_documents_view.ts

@@ -19,6 +19,7 @@ import { Query } from '../core/query';
 import { SnapshotVersion } from '../core/snapshot_version';
 import {
   DocumentKeySet,
+  documentKeySet,
   DocumentMap,
   documentMap,
   MaybeDocumentMap,
@@ -34,6 +35,7 @@ import { ResourcePath } from '../model/path';
 import { assert } from '../util/assert';
 import { IndexManager } from './index_manager';
 import { MutationQueue } from './mutation_queue';
+import { PatchMutation } from '../model/mutation';
 import { PersistenceTransaction } from './persistence';
 import { PersistencePromise } from './persistence_promise';
 import { RemoteDocumentCache } from './remote_document_cache';
@@ -49,7 +51,7 @@ export class LocalDocumentsView {
     private remoteDocumentCache: RemoteDocumentCache,
     private mutationQueue: MutationQueue,
     private indexManager: IndexManager
-  ) {}
+  ) { }
 
   /**
    * Get the local view of the document identified by `key`.
@@ -208,6 +210,7 @@ export class LocalDocumentsView {
   ): PersistencePromise<DocumentMap> {
     // Query the remote documents and overlay mutations.
     let results: DocumentMap;
+    let mutationBatches: MutationBatch[];
     return this.remoteDocumentCache
       .getDocumentsMatchingQuery(transaction, query)
       .next(queryResults => {
@@ -218,27 +221,41 @@ export class LocalDocumentsView {
         );
       })
       .next(matchingMutationBatches => {
-        for (const batch of matchingMutationBatches) {
-          for (const mutation of batch.mutations) {
-            const key = mutation.key;
-            // Only process documents belonging to the collection.
-            if (!query.path.isImmediateParentOf(key.path)) {
-              continue;
+        mutationBatches = matchingMutationBatches;
+        // It is possible that a PatchMutation can make a document match a query, even if
+        // the version in the RemoteDocumentCache is not a match yet (waiting for server
+        // to ack). To handle this, we find all document keys affected by the PatchMutations
+        // that are not in `result` yet, and back fill them via `remoteDocumentCache.getEntries`,
+        // otherwise those `PatchMutations` will be ignored because no base document can be found,
+        // and lead to missing result for the query.
+        return this.getMissingBaseDocuments(
+          transaction,
+          mutationBatches,
+          results
+        ).next(missingBaseDocuments => {
+          missingBaseDocuments.forEach((key, doc) => {
+            if (doc !== null && doc instanceof Document) {
+              results = results.insert(key, doc);
             }
+          });
 
-            const baseDoc = results.get(key);
-            const mutatedDoc = mutation.applyToLocalView(
-              baseDoc,
-              baseDoc,
-              batch.localWriteTime
-            );
-            if (mutatedDoc instanceof Document) {
-              results = results.insert(key, mutatedDoc);
-            } else {
-              results = results.remove(key);
+          for (const batch of mutationBatches) {
+            for (const mutation of batch.mutations) {
+              const key = mutation.key;
+              const baseDoc = results.get(key);
+              const mutatedDoc = mutation.applyToLocalView(
+                baseDoc,
+                baseDoc,
+                batch.localWriteTime
+              );
+              if (mutatedDoc instanceof Document) {
+                results = results.insert(key, mutatedDoc);
+              } else {
+                results = results.remove(key);
+              }
             }
           }
-        }
+        });
       })
       .next(() => {
         // Finally, filter out any documents that don't actually match
@@ -252,4 +269,28 @@ export class LocalDocumentsView {
         return results;
       });
   }
+
+  private getMissingBaseDocuments(
+    transaction: PersistenceTransaction,
+    matchingMutationBatches: MutationBatch[],
+    existingDocuments: DocumentMap
+  ): PersistencePromise<NullableMaybeDocumentMap> {
+    let missingBaseDocEntriesForPatching = documentKeySet();
+    for (const batch of matchingMutationBatches) {
+      for (const mutation of batch.mutations) {
+        if (
+          mutation instanceof PatchMutation &&
+          existingDocuments.get(mutation.key) === null
+        ) {
+          missingBaseDocEntriesForPatching = missingBaseDocEntriesForPatching.add(
+            mutation.key
+          );
+        }
+      }
+    }
+    return this.remoteDocumentCache.getEntries(
+      transaction,
+      missingBaseDocEntriesForPatching
+    );
+  }
 }

packages/firestore/test/unit/specs/query_spec.test.ts

@@ -16,7 +16,7 @@
  */
 
 import { Query } from '../../../src/core/query';
-import { doc, filter } from '../../util/helpers';
+import { doc, filter, path } from '../../util/helpers';
 
 import { Document } from '../../../src/model/document';
 import { ResourcePath } from '../../../src/model/path';
@@ -98,4 +98,41 @@ describeSpec('Queries:', [], () => {
         hasPendingWrites: true
       });
   });
+
+  specTest(
+    'Latency-compensated updates are included in query results',
+    [],
+    () => {
+      const fullQuery = Query.atPath(path('collection'));
+      const filteredQuery = fullQuery.addFilter(filter('match', '==', true));
+      const docA = doc('collection/a', 1000, { match: false });
+      const docAv2Local = doc(
+        'collection/a',
+        1000,
+        { match: true },
+        { hasLocalMutations: true }
+      );
+
+      return (
+        spec()
+          .userListens(fullQuery)
+          .watchAcksFull(fullQuery, 1000, docA)
+          .expectEvents(fullQuery, { added: [docA] })
+
+          // patch docA so that it will now match the filtered query.
+          .userPatches('collection/a', { match: true })
+          .expectEvents(fullQuery, {
+            modified: [docAv2Local],
+            hasPendingWrites: true
+          })
+          // Make sure docA shows up in filtered query.
+          .userListens(filteredQuery)
+          .expectEvents(filteredQuery, {
+            added: [docAv2Local],
+            fromCache: true,
+            hasPendingWrites: true
+          })
+      );
+    }
+  );
 });