Firestore overlays

packages/firestore/src/local/document_overlay_cache.ts

@@ -43,6 +43,15 @@ export interface DocumentOverlayCache {
     key: DocumentKey
   ): PersistencePromise<Overlay | null>;
 
+  /**
+   * Gets the saved overlay mutation for the given document keys. Skips keys for
+   * which there are no overlays.
+   */
+  getOverlays(
+    transaction: PersistenceTransaction,
+    keys: DocumentKeySet
+  ): PersistencePromise<OverlayMap>;
+
   /**
    * Saves the given document mutation map to persistence as overlays.
    * All overlays will have their largest batch id set to `largestBatchId`.

packages/firestore/src/local/indexeddb_document_overlay_cache.ts

@@ -81,6 +81,24 @@ export class IndexedDbDocumentOverlayCache implements DocumentOverlayCache {
       });
   }
 
+  getOverlays(
+    transaction: PersistenceTransaction,
+    keys: DocumentKeySet
+  ): PersistencePromise<OverlayMap> {
+    const result = newOverlayMap();
+    const promises: Array<PersistencePromise<void>> = [];
+    keys.forEach(key => {
+      promises.push(
+        this.getOverlay(transaction, key).next(overlay => {
+          if (overlay !== null) {
+            result.set(key, overlay);
+          }
+        })
+      );
+    });
+    return PersistencePromise.waitFor(promises).next(() => result);
+  }
+
   saveOverlays(
     transaction: PersistenceTransaction,
     largestBatchId: number,

packages/firestore/src/local/indexeddb_schema.ts

@@ -31,7 +31,7 @@ import { DbTimestampKey } from './indexeddb_sentinels';
 // TODO(indexing): Remove this constant
 export const INDEXING_ENABLED = false;
 
-export const INDEXING_SCHEMA_VERSION = 14;
+export const INDEXING_SCHEMA_VERSION = 15;
 
 /**
  * Schema Version for the Web client:
@@ -54,10 +54,11 @@ export const INDEXING_SCHEMA_VERSION = 14;
  * 12. Add document overlays.
  * 13. Rewrite the keys of the remote document cache to allow for efficient
  *     document lookup via `getAll()`.
- * 14. Add indexing support.
+ * 14. Add overlays.
+ * 15. Add indexing support.
  */
 
-export const SCHEMA_VERSION = INDEXING_ENABLED ? INDEXING_SCHEMA_VERSION : 13;
+export const SCHEMA_VERSION = INDEXING_ENABLED ? INDEXING_SCHEMA_VERSION : 14;
 
 /**
  * Wrapper class to store timestamps (seconds and nanos) in IndexedDb objects.

packages/firestore/src/local/indexeddb_schema_converter.ts

@@ -15,8 +15,16 @@
  * limitations under the License.
  */
 
+import { User } from '../auth/user';
+import { ListenSequence } from '../core/listen_sequence';
 import { SnapshotVersion } from '../core/snapshot_version';
+import {
+  DocumentKeyMap,
+  DocumentKeySet,
+  documentKeySet
+} from '../model/collections';
 import { DocumentKey } from '../model/document_key';
+import { FieldMask } from '../model/field_mask';
 import { ResourcePath } from '../model/path';
 import { debugAssert, fail, hardAssert } from '../util/assert';
 import { BATCHID_UNKNOWN } from '../util/types';
@@ -25,10 +33,13 @@ import {
   decodeResourcePath,
   encodeResourcePath
 } from './encoded_resource_path';
+import { IndexedDbDocumentOverlayCache } from './indexeddb_document_overlay_cache';
 import {
   dbDocumentSize,
   removeMutationBatch
 } from './indexeddb_mutation_batch_impl';
+import { IndexedDbMutationQueue } from './indexeddb_mutation_queue';
+import { newIndexedDbRemoteDocumentCache } from './indexeddb_remote_document_cache';
 import {
   DbCollectionParent,
   DbDocumentMutation,
@@ -107,13 +118,16 @@ import {
   DbTargetQueryTargetsKeyPath,
   DbTargetStore
 } from './indexeddb_sentinels';
+import { IndexedDbTransaction } from './indexeddb_transaction';
+import { LocalDocumentsView } from './local_documents_view';
 import {
   fromDbMutationBatch,
   fromDbTarget,
   LocalSerializer,
   toDbTarget
 } from './local_serializer';
 import { MemoryCollectionParentIndex } from './memory_index_manager';
+import { MemoryEagerDelegate, MemoryPersistence } from './memory_persistence';
 import { PersistencePromise } from './persistence_promise';
 import { SimpleDbSchemaConverter, SimpleDbTransaction } from './simple_db';
 
@@ -240,9 +254,11 @@ export class SchemaConverter implements SimpleDbSchemaConverter {
     }
 
     if (fromVersion < 14 && toVersion >= 14) {
-      p = p.next(() => {
-        createFieldIndex(db);
-      });
+      p = p.next(() => this.runOverlayMigration(db, simpleDbTransaction));
+    }
+
+    if (fromVersion < 15 && toVersion >= 15) {
+      p = p.next(() => createFieldIndex(db));
     }
 
     return p;
@@ -455,6 +471,74 @@ export class SchemaConverter implements SimpleDbSchemaConverter {
       })
       .next(() => PersistencePromise.waitFor(writes));
   }
+
+  private runOverlayMigration(
+    db: IDBDatabase,
+    transaction: SimpleDbTransaction
+  ): PersistencePromise<void> {
+    const mutationsStore = transaction.store<
+      DbMutationBatchKey,
+      DbMutationBatch
+    >(DbMutationBatchStore);
+
+    const remoteDocumentCache = newIndexedDbRemoteDocumentCache(
+      this.serializer
+    );
+    const memoryPersistence = new MemoryPersistence(
+      MemoryEagerDelegate.factory,
+      this.serializer.remoteSerializer
+    );
+
+    const promises: Array<
+      PersistencePromise<DocumentKeyMap<FieldMask | null>>
+    > = [];
+    const userToDocumentSet = new Map<string, DocumentKeySet>();
+
+    return mutationsStore
+      .loadAll()
+      .next(dbBatches => {
+        dbBatches.forEach(dbBatch => {
+          let documentSet =
+            userToDocumentSet.get(dbBatch.userId) ?? documentKeySet();
+          const batch = fromDbMutationBatch(this.serializer, dbBatch);
+          batch.keys().forEach(key => (documentSet = documentSet.add(key)));
+          userToDocumentSet.set(dbBatch.userId, documentSet);
+        });
+      })
+      .next(() => {
+        userToDocumentSet.forEach((allDocumentKeysForUser, userId) => {
+          const user = new User(userId);
+          const documentOverlayCache = IndexedDbDocumentOverlayCache.forUser(
+            this.serializer,
+            user
+          );
+          // NOTE: The index manager and the reference delegate are
+          // irrelevant for the purpose of recalculating and saving
+          // overlays. We can therefore simply use the memory
+          // implementation.
+          const indexManager = memoryPersistence.getIndexManager(user);
+          const mutationQueue = IndexedDbMutationQueue.forUser(
+            user,
+            this.serializer,
+            indexManager,
+            memoryPersistence.referenceDelegate
+          );
+          const localDocumentsView = new LocalDocumentsView(
+            remoteDocumentCache,
+            mutationQueue,
+            documentOverlayCache,
+            indexManager
+          );
+          promises.push(
+            localDocumentsView.recalculateAndSaveOverlaysForDocumentKeys(
+              new IndexedDbTransaction(transaction, ListenSequence.INVALID),
+              allDocumentKeysForUser
+            )
+          );
+        });
+      })
+      .next(() => PersistencePromise.waitFor(promises));
+  }
 }
 
 function sentinelKey(path: ResourcePath): DbTargetDocumentKey {

packages/firestore/src/local/indexeddb_sentinels.ts

@@ -407,8 +407,9 @@ export const V13_STORES = [
   DbNamedQueryStore,
   DbDocumentOverlayStore
 ];
-export const V14_STORES = [
-  ...V13_STORES,
+export const V14_STORES = V13_STORES;
+export const V15_STORES = [
+  ...V14_STORES,
   DbIndexConfigurationStore,
   DbIndexStateStore,
   DbIndexEntryStore
@@ -423,7 +424,9 @@ export const ALL_STORES = V12_STORES;
 
 /** Returns the object stores for the provided schema. */
 export function getObjectStores(schemaVersion: number): string[] {
-  if (schemaVersion === 14) {
+  if (schemaVersion === 15) {
+    return V15_STORES;
+  } else if (schemaVersion === 14) {
     return V14_STORES;
   } else if (schemaVersion === 13) {
     return V13_STORES;