Update query-document mapping from bundles.

packages/firestore/src/core/bundle.ts

@@ -31,6 +31,7 @@ import { MaybeDocument, NoDocument } from '../model/document';
 import { debugAssert } from '../util/assert';
 import {
   applyBundleDocuments,
+  getQueryDocumentMapping,
   LocalStore,
   saveNamedQuery
 } from '../local/local_store';
@@ -218,16 +219,20 @@ export class BundleLoader {
       'Bundled documents ends with a document metadata and missing document.'
     );
 
-    for (const q of this.queries) {
-      await saveNamedQuery(this.localStore, q);
-    }
+    const changedDocuments = await applyBundleDocuments(
+      this.localStore,
+      this.documents
+    );
 
-    const changedDocs = await applyBundleDocuments(
+    const queryDocumentMap = getQueryDocumentMapping(
       this.localStore,
       this.documents
     );
+    for (const q of this.queries) {
+      await saveNamedQuery(this.localStore, q, queryDocumentMap.get(q.name!));
+    }
 
     this.progress.taskState = 'Success';
-    return new BundleLoadResult({ ...this.progress }, changedDocs);
+    return new BundleLoadResult({ ...this.progress }, changedDocuments);
   }
 }

packages/firestore/src/local/local_store.ts

@@ -1249,6 +1249,30 @@ export async function ignoreIfPrimaryLeaseLoss(
   }
 }
 
+export function getQueryDocumentMapping(
+  localStore: LocalStore,
+  documents: BundledDocuments
+): Map<string, DocumentKeySet> {
+  const localStoreImpl = debugCast(localStore, LocalStoreImpl);
+  const queryDocumentMap = new Map<string, DocumentKeySet>();
+  const bundleConverter = new BundleConverter(localStoreImpl.serializer);
+  for (const bundleDoc of documents) {
+    if (bundleDoc.metadata.queries) {
+      const documentKey = bundleConverter.toDocumentKey(
+        bundleDoc.metadata.name!
+      );
+      for (const queryName of bundleDoc.metadata.queries) {
+        const documentKeys = (
+          queryDocumentMap.get(queryName) || documentKeySet()
+        ).add(documentKey);
+        queryDocumentMap.set(queryName, documentKeys);
+      }
+    }
+  }
+
+  return queryDocumentMap;
+}
+
 /**
  * Applies the documents from a bundle to the "ground-state" (remote)
  * documents.
@@ -1300,6 +1324,9 @@ export function applyBundleDocuments(
             txn,
             changedDocs
           );
+        })
+        .next(changedDocuments => {
+          return PersistencePromise.resolve(changedDocuments);
         });
     }
   );
@@ -1373,7 +1400,8 @@ export function getNamedQuery(
  */
 export async function saveNamedQuery(
   localStore: LocalStore,
-  query: bundleProto.NamedQuery
+  query: bundleProto.NamedQuery,
+  documents: DocumentKeySet = documentKeySet()
 ): Promise<void> {
   // Allocate a target for the named query such that it can be resumed
   // from associated read time if users use it to listen.
@@ -1389,21 +1417,41 @@ export async function saveNamedQuery(
     'readwrite',
     transaction => {
       // Update allocated target's read time, if the bundle's read time is newer.
-      let updateReadTime = PersistencePromise.resolve();
+      let updateReadTime = PersistencePromise.resolve(false);
       const readTime = fromVersion(query.readTime!);
       if (allocated.snapshotVersion.compareTo(readTime) < 0) {
         const newTargetData = allocated.withResumeToken(
           ByteString.EMPTY_BYTE_STRING,
           readTime
         );
-        updateReadTime = localStoreImpl.targetCache.updateTargetData(
-          transaction,
+        updateReadTime = localStoreImpl.targetCache
+          .updateTargetData(transaction, newTargetData)
+          .next(
+            () => true,
+            () => false
+          );
+        localStoreImpl.targetDataByTarget = localStoreImpl.targetDataByTarget.insert(
+          newTargetData.targetId,
           newTargetData
         );
       }
-      return updateReadTime.next(() =>
-        localStoreImpl.bundleCache.saveNamedQuery(transaction, query)
-      );
+      return updateReadTime
+        .next(updated => {
+          if (updated) {
+            return localStoreImpl.targetCache
+              .removeMatchingKeysForTargetId(transaction, allocated.targetId)
+              .next(() =>
+                localStoreImpl.targetCache.addMatchingKeys(
+                  transaction,
+                  documents,
+                  allocated.targetId
+                )
+              );
+          }
+        })
+        .next(() =>
+          localStoreImpl.bundleCache.saveNamedQuery(transaction, query)
+        );
     }
   );
 }

packages/firestore/src/protos/firestore/bundle.proto

@@ -79,6 +79,9 @@ message BundledDocumentMetadata {
 
   // Whether the document exists.
   bool exists = 3;
+
+  // The names of the queries in this bundle that this document matches to.
+  repeated string queries = 4;
 }
 
 // Metadata describing the bundle file/stream.

packages/firestore/src/protos/firestore_bundle_proto.ts

@@ -54,6 +54,9 @@ export interface BundledDocumentMetadata {
 
   /** BundledDocumentMetadata exists */
   exists?: boolean | null;
+
+  /** The names of the queries in this bundle that this document matches to. */
+  queries?: string[];
 }
 
 /** Properties of a BundleMetadata. */

packages/firestore/test/unit/local/local_store.test.ts

@@ -48,6 +48,7 @@ import { LocalViewChanges } from '../../../src/local/local_view_changes';
 import { Persistence } from '../../../src/local/persistence';
 import { SimpleQueryEngine } from '../../../src/local/simple_query_engine';
 import {
+  DocumentKeySet,
   documentKeySet,
   MaybeDocumentMap
 } from '../../../src/model/collections';
@@ -91,6 +92,7 @@ import {
   query,
   setMutation,
   TestBundledDocuments,
+  TestNamedQuery,
   TestSnapshotVersion,
   transformMutation,
   unknownDoc,
@@ -137,7 +139,7 @@ class LocalStoreTester {
       | RemoteEvent
       | LocalViewChanges
       | TestBundledDocuments
-      | bundleProto.NamedQuery
+      | TestNamedQuery
   ): LocalStoreTester {
     if (op instanceof Mutation) {
       return this.afterMutations([op]);
@@ -188,17 +190,21 @@ class LocalStoreTester {
 
     this.promiseChain = this.promiseChain
       .then(() => applyBundleDocuments(this.localStore, documents))
-      .then((result: MaybeDocumentMap) => {
+      .then(result => {
         this.lastChanges = result;
       });
     return this;
   }
 
-  afterNamedQuery(namedQuery: bundleProto.NamedQuery): LocalStoreTester {
+  afterNamedQuery(testQuery: TestNamedQuery): LocalStoreTester {
     this.prepareNextStep();
 
     this.promiseChain = this.promiseChain.then(() =>
-      saveNamedQuery(this.localStore, namedQuery)
+      saveNamedQuery(
+        this.localStore,
+        testQuery.namedQuery,
+        testQuery.matchingDocuments
+      )
     );
     return this;
   }
@@ -432,6 +438,29 @@ class LocalStoreTester {
     return this;
   }
 
+  toHaveQueryDocumentMapping(
+    persistence: Persistence,
+    targetId: TargetId,
+    expectedKeys: DocumentKeySet
+  ): LocalStoreTester {
+    this.promiseChain = this.promiseChain.then(() => {
+      return persistence.runTransaction(
+        'toHaveQueryDocumentMapping',
+        'readonly',
+        transaction => {
+          return persistence
+            .getTargetCache()
+            .getMatchingKeysForTargetId(transaction, targetId)
+            .next(matchedKeys => {
+              expect(matchedKeys.isEqual(expectedKeys)).to.be.true;
+            });
+        }
+      );
+    });
+
+    return this;
+  }
+
   toHaveNewerBundle(
     metadata: bundleProto.BundleMetadata,
     expected: boolean
@@ -1706,6 +1735,68 @@ function genericLocalStoreTests(
       .finish();
   });
 
+  it('loading named queries allocates targets and updates target document mapping', async () => {
+    const expectedQueryDocumentMap = new Map([
+      ['query-1', documentKeySet(key('foo1/bar'))],
+      ['query-2', documentKeySet(key('foo2/bar'))]
+    ]);
+    const version1 = SnapshotVersion.fromTimestamp(Timestamp.fromMillis(10000));
+    const version2 = SnapshotVersion.fromTimestamp(Timestamp.fromMillis(20000));
+
+    return expectLocalStore()
+      .after(
+        bundledDocuments(
+          [doc('foo1/bar', 1, { sum: 1337 }), doc('foo2/bar', 2, { sum: 42 })],
+          [['query-1'], ['query-2']]
+        )
+      )
+      .toReturnChanged(
+        doc('foo1/bar', 1, { sum: 1337 }),
+        doc('foo2/bar', 2, { sum: 42 })
+      )
+      .toContain(doc('foo1/bar', 1, { sum: 1337 }))
+      .toContain(doc('foo2/bar', 2, { sum: 42 }))
+      .after(
+        namedQuery(
+          'query-1',
+          query('foo1'),
+          /* limitType */ 'FIRST',
+          version1,
+          expectedQueryDocumentMap.get('query-1')
+        )
+      )
+      .toHaveNamedQuery({
+        name: 'query-1',
+        query: query('foo1'),
+        readTime: version1
+      })
+      .toHaveQueryDocumentMapping(
+        persistence,
+        /*targetId*/ 2,
+        /*expectedKeys*/ documentKeySet(key('foo1/bar'))
+      )
+      .after(
+        namedQuery(
+          'query-2',
+          query('foo2'),
+          /* limitType */ 'FIRST',
+          version2,
+          expectedQueryDocumentMap.get('query-2')
+        )
+      )
+      .toHaveNamedQuery({
+        name: 'query-2',
+        query: query('foo2'),
+        readTime: version2
+      })
+      .toHaveQueryDocumentMapping(
+        persistence,
+        /*targetId*/ 4,
+        /*expectedKeys*/ documentKeySet(key('foo2/bar'))
+      )
+      .finish();
+  });
+
   it('handles saving and loading limit to last queries', async () => {
     const now = Timestamp.now();
     return expectLocalStore()

packages/firestore/test/util/helpers.ts

@@ -445,14 +445,16 @@ export class TestBundledDocuments {
 }
 
 export function bundledDocuments(
-  documents: MaybeDocument[]
+  documents: MaybeDocument[],
+  queryNames?: string[][]
 ): TestBundledDocuments {
-  const result = documents.map(d => {
+  const result = documents.map((d, index) => {
     return {
       metadata: {
         name: toName(JSON_SERIALIZER, d.key),
         readTime: toVersion(JSON_SERIALIZER, d.version),
-        exists: d instanceof Document
+        exists: d instanceof Document,
+        queries: queryNames ? queryNames[index] : undefined
       },
       document:
         d instanceof Document ? toDocument(JSON_SERIALIZER, d) : undefined
@@ -462,21 +464,32 @@ export function bundledDocuments(
   return new TestBundledDocuments(result);
 }
 
+export class TestNamedQuery {
+  constructor(
+    public namedQuery: bundleProto.NamedQuery,
+    public matchingDocuments: DocumentKeySet
+  ) {}
+}
+
 export function namedQuery(
   name: string,
   query: Query,
   limitType: bundleProto.LimitType,
-  readTime: SnapshotVersion
-): bundleProto.NamedQuery {
+  readTime: SnapshotVersion,
+  matchingDocuments: DocumentKeySet = documentKeySet()
+): TestNamedQuery {
   return {
-    name,
-    readTime: toTimestamp(JSON_SERIALIZER, readTime.toTimestamp()),
-    bundledQuery: {
-      parent: toQueryTarget(JSON_SERIALIZER, queryToTarget(query)).parent,
-      limitType,
-      structuredQuery: toQueryTarget(JSON_SERIALIZER, queryToTarget(query))
-        .structuredQuery
-    }
+    namedQuery: {
+      name,
+      readTime: toTimestamp(JSON_SERIALIZER, readTime.toTimestamp()),
+      bundledQuery: {
+        parent: toQueryTarget(JSON_SERIALIZER, queryToTarget(query)).parent,
+        limitType,
+        structuredQuery: toQueryTarget(JSON_SERIALIZER, queryToTarget(query))
+          .structuredQuery
+      }
+    },
+    matchingDocuments
   };
 }