Update query-document mapping from bundles.

packages/firestore/src/core/bundle.ts

@@ -218,16 +218,17 @@ export class BundleLoader {
       'Bundled documents ends with a document metadata and missing document.'
     );
 
+    const result = await applyBundleDocuments(this.localStore, this.documents);
+
     for (const q of this.queries) {
-      await saveNamedQuery(this.localStore, q);
+      await saveNamedQuery(
+        this.localStore,
+        q,
+        result.queryDocumentMap.get(q.name!)
+      );
     }
 
-    const changedDocs = await applyBundleDocuments(
-      this.localStore,
-      this.documents
-    );
-
     this.progress.taskState = 'Success';
-    return new BundleLoadResult({ ...this.progress }, changedDocs);
+    return new BundleLoadResult({ ...this.progress }, result.changedDocuments);
   }
 }

packages/firestore/src/local/local_store.ts

@@ -1249,6 +1249,11 @@ export async function ignoreIfPrimaryLeaseLoss(
   }
 }
 
+export interface ApplyBundleDocumentsResult {
+  changedDocuments: MaybeDocumentMap;
+  queryDocumentMap: Map<string, DocumentKeySet>;
+}
+
 /**
  * Applies the documents from a bundle to the "ground-state" (remote)
  * documents.
@@ -1259,11 +1264,12 @@ export async function ignoreIfPrimaryLeaseLoss(
 export function applyBundleDocuments(
   localStore: LocalStore,
   documents: BundledDocuments
-): Promise<MaybeDocumentMap> {
+): Promise<ApplyBundleDocumentsResult> {
   const localStoreImpl = debugCast(localStore, LocalStoreImpl);
   const bundleConverter = new BundleConverter(localStoreImpl.serializer);
   let documentMap = maybeDocumentMap();
   let versionMap = documentVersionMap();
+  const queryDocumentMap: Map<string, DocumentKeySet> = new Map();
   for (const bundleDoc of documents) {
     const documentKey = bundleConverter.toDocumentKey(bundleDoc.metadata.name!);
     documentMap = documentMap.insert(
@@ -1274,6 +1280,14 @@ export function applyBundleDocuments(
       documentKey,
       bundleConverter.toSnapshotVersion(bundleDoc.metadata.readTime!)
     );
+    if (bundleDoc.metadata.queries) {
+      for (const queryName of bundleDoc.metadata.queries) {
+        const documentKeys = (
+          queryDocumentMap.get(queryName) || documentKeySet()
+        ).add(documentKey);
+        queryDocumentMap.set(queryName, documentKeys);
+      }
+    }
   }
 
   const documentBuffer = localStoreImpl.remoteDocuments.newChangeBuffer({
@@ -1300,6 +1314,12 @@ export function applyBundleDocuments(
             txn,
             changedDocs
           );
+        })
+        .next(changedDocuments => {
+          return PersistencePromise.resolve({
+            changedDocuments,
+            queryDocumentMap
+          });
         });
     }
   );
@@ -1373,7 +1393,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.
@@ -1400,10 +1421,22 @@ export async function saveNamedQuery(
           transaction,
           newTargetData
         );
+        localStoreImpl.targetDataByTarget = localStoreImpl.targetDataByTarget.insert(
+          newTargetData.targetId,
+          newTargetData
+        );
       }
-      return updateReadTime.next(() =>
-        localStoreImpl.bundleCache.saveNamedQuery(transaction, query)
-      );
+      return updateReadTime
+        .next(() =>
+          localStoreImpl.bundleCache.saveNamedQuery(transaction, query)
+        )
+        .next(() =>
+          localStoreImpl.targetCache.addMatchingKeys(
+            transaction,
+            documents,
+            allocated.targetId
+          )
+        );
     }
   );
 }

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[] | null;
 }
 
 /** Properties of a BundleMetadata. */

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

@@ -35,6 +35,7 @@ import { IndexFreeQueryEngine } from '../../../src/local/index_free_query_engine
 import { IndexedDbPersistence } from '../../../src/local/indexeddb_persistence';
 import {
   applyBundleDocuments,
+  ApplyBundleDocumentsResult,
   getNamedQuery,
   hasNewerBundle,
   LocalStore,
@@ -48,6 +49,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 +93,7 @@ import {
   query,
   setMutation,
   TestBundledDocuments,
+  TestNamedQuery,
   TestSnapshotVersion,
   transformMutation,
   unknownDoc,
@@ -113,6 +116,7 @@ export interface LocalStoreComponents {
 class LocalStoreTester {
   private promiseChain: Promise<void> = Promise.resolve();
   private lastChanges: MaybeDocumentMap | null = null;
+  private lastBundleQueryDocumentMap: Map<string, DocumentKeySet> | null = null;
   private lastTargetId: TargetId | null = null;
   private batches: MutationBatch[] = [];
 
@@ -125,6 +129,7 @@ class LocalStoreTester {
   private prepareNextStep(): void {
     this.promiseChain = this.promiseChain.then(() => {
       this.lastChanges = null;
+      this.lastBundleQueryDocumentMap = null;
       this.lastTargetId = null;
       this.queryEngine.resetCounts();
     });
@@ -137,7 +142,7 @@ class LocalStoreTester {
       | RemoteEvent
       | LocalViewChanges
       | TestBundledDocuments
-      | bundleProto.NamedQuery
+      | TestNamedQuery
   ): LocalStoreTester {
     if (op instanceof Mutation) {
       return this.afterMutations([op]);
@@ -188,17 +193,22 @@ class LocalStoreTester {
 
     this.promiseChain = this.promiseChain
       .then(() => applyBundleDocuments(this.localStore, documents))
-      .then((result: MaybeDocumentMap) => {
-        this.lastChanges = result;
+      .then((result: ApplyBundleDocumentsResult) => {
+        this.lastChanges = result.changedDocuments;
+        this.lastBundleQueryDocumentMap = result.queryDocumentMap;
       });
     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;
   }
@@ -372,6 +382,25 @@ class LocalStoreTester {
     return this;
   }
 
+  toReturnQueryDocumentMap(
+    expected: Map<string, DocumentKeySet>
+  ): LocalStoreTester {
+    this.promiseChain = this.promiseChain.then(() => {
+      debugAssert(
+        !!this.lastBundleQueryDocumentMap,
+        'Expecting query document map to be returned from applying bundled documents'
+      );
+      for (const k of Array.from(this.lastBundleQueryDocumentMap.keys())) {
+        expect(expected.has(k)).to.be.true;
+        expect(
+          this.lastBundleQueryDocumentMap.get(k)!.isEqual(expected.get(k)!)
+        ).to.be.true;
+      }
+      this.lastBundleQueryDocumentMap = null;
+    });
+    return this;
+  }
+
   toReturnRemoved(...keyStrings: string[]): LocalStoreTester {
     this.promiseChain = this.promiseChain.then(() => {
       debugAssert(
@@ -432,6 +461,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 +1758,69 @@ function genericLocalStoreTests(
       .finish();
   });
 
+  it('handles 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 })
+      )
+      .toReturnQueryDocumentMap(expectedQueryDocumentMap)
+      .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
   };
 }