Add a type parameter to Persistence (#1047)

* Cherry pick sequence number starting point

* Working on typed transactions

* Start plumbing in sequence number

* Back out sequence number changes

* [AUTOMATED]: Prettier Code Styling

* Fix tests

* [AUTOMATED]: Prettier Code Styling

* Fix lint

* [AUTOMATED]: Prettier Code Styling

* Uncomment line

* MemoryPersistenceTransaction -> MemoryTransaction

* [AUTOMATED]: Prettier Code Styling

* Review updates

* Style

* Lint and style

* Review feedback

* [AUTOMATED]: Prettier Code Styling

* Revert some unintentional import churn

* Line 44 should definitely be empty

* Checkpoint before adding helper function for stores

* Use a helper for casting PersistenceTransaction to IndexedDbTransaction

* [AUTOMATED]: Prettier Code Styling

* Remove errant generic type

* Lint

* Fix typo

Author: Greg Soltis

packages/firestore/src/local/indexeddb_mutation_queue.ts

@@ -40,7 +40,8 @@ import { LocalSerializer } from './local_serializer';
 import { MutationQueue } from './mutation_queue';
 import { PersistenceTransaction } from './persistence';
 import { PersistencePromise } from './persistence_promise';
-import { SimpleDb, SimpleDbStore } from './simple_db';
+import { SimpleDbStore } from './simple_db';
+import { IndexedDbPersistence } from './indexeddb_persistence';
 
 /** A mutation queue for a specific user, backed by IndexedDB. */
 export class IndexedDbMutationQueue implements MutationQueue {
@@ -563,7 +564,7 @@ function convertStreamToken(token: ProtoByteString): string {
 function mutationsStore(
   txn: PersistenceTransaction
 ): SimpleDbStore<DbMutationBatchKey, DbMutationBatch> {
-  return SimpleDb.getStore<DbMutationBatchKey, DbMutationBatch>(
+  return IndexedDbPersistence.getStore<DbMutationBatchKey, DbMutationBatch>(
     txn,
     DbMutationBatch.store
   );
@@ -575,10 +576,10 @@ function mutationsStore(
 function documentMutationsStore(
   txn: PersistenceTransaction
 ): SimpleDbStore<DbDocumentMutationKey, DbDocumentMutation> {
-  return SimpleDb.getStore<DbDocumentMutationKey, DbDocumentMutation>(
-    txn,
-    DbDocumentMutation.store
-  );
+  return IndexedDbPersistence.getStore<
+    DbDocumentMutationKey,
+    DbDocumentMutation
+  >(txn, DbDocumentMutation.store);
 }
 
 /**
@@ -587,7 +588,7 @@ function documentMutationsStore(
 function mutationQueuesStore(
   txn: PersistenceTransaction
 ): SimpleDbStore<DbMutationQueueKey, DbMutationQueue> {
-  return SimpleDb.getStore<DbMutationQueueKey, DbMutationQueue>(
+  return IndexedDbPersistence.getStore<DbMutationQueueKey, DbMutationQueue>(
     txn,
     DbMutationQueue.store
   );

packages/firestore/src/local/indexeddb_persistence.ts

@@ -17,7 +17,7 @@
 import { User } from '../auth/user';
 import { DatabaseInfo } from '../core/database_info';
 import { JsonProtoSerializer } from '../remote/serializer';
-import { assert } from '../util/assert';
+import { assert, fail } from '../util/assert';
 import { Code, FirestoreError } from '../util/error';
 import * as log from '../util/log';
 import { AutoId } from '../util/misc';
@@ -34,11 +34,11 @@ import {
 } from './indexeddb_schema';
 import { LocalSerializer } from './local_serializer';
 import { MutationQueue } from './mutation_queue';
-import { Persistence } from './persistence';
+import { Persistence, PersistenceTransaction } from './persistence';
 import { PersistencePromise } from './persistence_promise';
 import { QueryCache } from './query_cache';
 import { RemoteDocumentCache } from './remote_document_cache';
-import { SimpleDb, SimpleDbTransaction } from './simple_db';
+import { SimpleDb, SimpleDbStore, SimpleDbTransaction } from './simple_db';
 
 const LOG_TAG = 'IndexedDbPersistence';
 
@@ -59,6 +59,12 @@ const UNSUPPORTED_PLATFORM_ERROR_MSG =
   ' IndexedDB or is known to have an incomplete implementation. Offline' +
   ' persistence has been disabled.';
 
+export class IndexedDbTransaction extends PersistenceTransaction {
+  constructor(readonly simpleDbTransaction: SimpleDbTransaction) {
+    super();
+  }
+}
+
 /**
  * An IndexedDB-backed instance of Persistence. Data is stored persistently
  * across sessions.
@@ -89,6 +95,17 @@ const UNSUPPORTED_PLATFORM_ERROR_MSG =
  * owner lease immediately regardless of the current lease timestamp.
  */
 export class IndexedDbPersistence implements Persistence {
+  static getStore<Key extends IDBValidKey, Value>(
+    txn: PersistenceTransaction,
+    store: string
+  ): SimpleDbStore<Key, Value> {
+    if (txn instanceof IndexedDbTransaction) {
+      return SimpleDb.getStore<Key, Value>(txn.simpleDbTransaction, store);
+    } else {
+      fail('IndexedDbPersistence must use instances of IndexedDbTransaction');
+    }
+  }
+
   /**
    * The name of the main (and currently only) IndexedDB database. this name is
    * appended to the prefix provided to the IndexedDbPersistence constructor.
@@ -176,7 +193,7 @@ export class IndexedDbPersistence implements Persistence {
 
   runTransaction<T>(
     action: string,
-    operation: (transaction: SimpleDbTransaction) => PersistencePromise<T>
+    operation: (transaction: IndexedDbTransaction) => PersistencePromise<T>
   ): Promise<T> {
     if (this.persistenceError) {
       return Promise.reject(this.persistenceError);
@@ -186,10 +203,16 @@ export class IndexedDbPersistence implements Persistence {
 
     // Do all transactions as readwrite against all object stores, since we
     // are the only reader/writer.
-    return this.simpleDb.runTransaction('readwrite', ALL_STORES, txn => {
-      // Verify that we still have the owner lease as part of every transaction.
-      return this.ensureOwnerLease(txn).next(() => operation(txn));
-    });
+    return this.simpleDb.runTransaction(
+      'readwrite',
+      ALL_STORES,
+      simpleDbTxn => {
+        // Verify that we still have the owner lease as part of every transaction.
+        return this.ensureOwnerLease(simpleDbTxn).next(() =>
+          operation(new IndexedDbTransaction(simpleDbTxn))
+        );
+      }
+    );
   }
 
   static isAvailable(): boolean {
@@ -326,12 +349,16 @@ export class IndexedDbPersistence implements Persistence {
     // would increase the chances of us not refreshing on time if the queue is
     // backed up for some reason.
     this.ownerLeaseRefreshHandle = setInterval(() => {
-      const txResult = this.runTransaction('Refresh owner timestamp', txn => {
-        // NOTE: We don't need to validate the current owner contents, since
-        // runTransaction does that automatically.
-        const store = txn.store<DbOwnerKey, DbOwner>(DbOwner.store);
-        return store.put('owner', new DbOwner(this.ownerId, Date.now()));
-      });
+      const txResult = this.simpleDb.runTransaction(
+        'readwrite',
+        ALL_STORES,
+        txn => {
+          // NOTE: We don't need to validate the current owner contents, since
+          // runTransaction does that automatically.
+          const store = txn.store<DbOwnerKey, DbOwner>(DbOwner.store);
+          return store.put('owner', new DbOwner(this.ownerId, Date.now()));
+        }
+      );
 
       txResult.catch(reason => {
         // Probably means we lost the lease. Report the error and stop trying to

packages/firestore/src/local/indexeddb_query_cache.ts

@@ -38,7 +38,8 @@ import { PersistenceTransaction } from './persistence';
 import { PersistencePromise } from './persistence_promise';
 import { QueryCache } from './query_cache';
 import { QueryData } from './query_data';
-import { SimpleDb, SimpleDbStore } from './simple_db';
+import { SimpleDbStore } from './simple_db';
+import { IndexedDbPersistence } from './indexeddb_persistence';
 
 export class IndexedDbQueryCache implements QueryCache {
   constructor(private serializer: LocalSerializer) {}
@@ -52,7 +53,7 @@ export class IndexedDbQueryCache implements QueryCache {
   /**
    * A cached copy of the metadata for the query cache.
    */
-  private metadata = null;
+  private metadata: DbTargetGlobal = null;
 
   /** The garbage collector to notify about potential garbage keys. */
   private garbageCollector: GarbageCollector | null = null;
@@ -297,6 +298,8 @@ export class IndexedDbQueryCache implements QueryCache {
     this.garbageCollector = gc;
   }
 
+  // TODO(gsoltis): we can let the compiler assert that txn !== null if we
+  // drop null from the type bounds on txn.
   containsKey(
     txn: PersistenceTransaction | null,
     key: DocumentKey
@@ -335,7 +338,10 @@ export class IndexedDbQueryCache implements QueryCache {
 function targetsStore(
   txn: PersistenceTransaction
 ): SimpleDbStore<DbTargetKey, DbTarget> {
-  return SimpleDb.getStore<DbTargetKey, DbTarget>(txn, DbTarget.store);
+  return IndexedDbPersistence.getStore<DbTargetKey, DbTarget>(
+    txn,
+    DbTarget.store
+  );
 }
 
 /**
@@ -344,7 +350,7 @@ function targetsStore(
 function globalTargetStore(
   txn: PersistenceTransaction
 ): SimpleDbStore<DbTargetGlobalKey, DbTargetGlobal> {
-  return SimpleDb.getStore<DbTargetGlobalKey, DbTargetGlobal>(
+  return IndexedDbPersistence.getStore<DbTargetGlobalKey, DbTargetGlobal>(
     txn,
     DbTargetGlobal.store
   );
@@ -356,7 +362,7 @@ function globalTargetStore(
 function documentTargetStore(
   txn: PersistenceTransaction
 ): SimpleDbStore<DbTargetDocumentKey, DbTargetDocument> {
-  return SimpleDb.getStore<DbTargetDocumentKey, DbTargetDocument>(
+  return IndexedDbPersistence.getStore<DbTargetDocumentKey, DbTargetDocument>(
     txn,
     DbTargetDocument.store
   );

packages/firestore/src/local/indexeddb_remote_document_cache.ts

@@ -20,11 +20,12 @@ import { Document, MaybeDocument } from '../model/document';
 import { DocumentKey } from '../model/document_key';
 
 import { DbRemoteDocument, DbRemoteDocumentKey } from './indexeddb_schema';
+import { IndexedDbPersistence } from './indexeddb_persistence';
 import { LocalSerializer } from './local_serializer';
 import { PersistenceTransaction } from './persistence';
 import { PersistencePromise } from './persistence_promise';
 import { RemoteDocumentCache } from './remote_document_cache';
-import { SimpleDb, SimpleDbStore } from './simple_db';
+import { SimpleDbStore } from './simple_db';
 
 export class IndexedDbRemoteDocumentCache implements RemoteDocumentCache {
   constructor(private serializer: LocalSerializer) {}
@@ -88,7 +89,7 @@ export class IndexedDbRemoteDocumentCache implements RemoteDocumentCache {
 function remoteDocumentsStore(
   txn: PersistenceTransaction
 ): SimpleDbStore<DbRemoteDocumentKey, DbRemoteDocument> {
-  return SimpleDb.getStore<DbRemoteDocumentKey, DbRemoteDocument>(
+  return IndexedDbPersistence.getStore<DbRemoteDocumentKey, DbRemoteDocument>(
     txn,
     DbRemoteDocument.store
   );

packages/firestore/src/local/memory_persistence.ts

@@ -85,9 +85,12 @@ export class MemoryPersistence implements Persistence {
     operation: (transaction: PersistenceTransaction) => PersistencePromise<T>
   ): Promise<T> {
     debug(LOG_TAG, 'Starting transaction:', action);
-    return operation(new MemoryPersistenceTransaction()).toPromise();
+    return operation(new MemoryTransaction()).toPromise();
   }
 }
 
-/** Dummy class since memory persistence doesn't actually use transactions. */
-class MemoryPersistenceTransaction implements PersistenceTransaction {}
+/**
+ * Memory persistence is not actually transactional, but future implementations
+ * may have transaction-scoped state.
+ */
+export class MemoryTransaction implements PersistenceTransaction {}

packages/firestore/src/local/persistence.ts

@@ -17,7 +17,6 @@
 import { User } from '../auth/user';
 
 import { MutationQueue } from './mutation_queue';
-import { PersistenceTransaction } from './persistence';
 import { PersistencePromise } from './persistence_promise';
 import { QueryCache } from './query_cache';
 import { RemoteDocumentCache } from './remote_document_cache';
@@ -29,7 +28,7 @@ import { RemoteDocumentCache } from './remote_document_cache';
  * pass it to your callback. You then pass it to any method that operates
  * on persistence.
  */
-export interface PersistenceTransaction {}
+export abstract class PersistenceTransaction {}
 
 /**
  * Persistence is the lowest-level shared interface to persistent storage in

packages/firestore/src/local/simple_db.ts

@@ -14,14 +14,12 @@
  * limitations under the License.
  */
 
-import { assert, fail } from '../util/assert';
+import { assert } from '../util/assert';
 import { debug } from '../util/log';
 import { AnyDuringMigration } from '../util/misc';
-
 import { PersistencePromise } from './persistence_promise';
 import { SCHEMA_VERSION } from './indexeddb_schema';
 import { Deferred } from '../util/promise';
-import { PersistenceTransaction } from './persistence';
 import { Code, FirestoreError } from '../util/error';
 
 const LOG_TAG = 'SimpleDb';
@@ -150,14 +148,10 @@ export class SimpleDb {
 
   /** Helper to get a typed SimpleDbStore from a transaction. */
   static getStore<KeyType extends IDBValidKey, ValueType>(
-    txn: PersistenceTransaction,
+    txn: SimpleDbTransaction,
     store: string
   ): SimpleDbStore<KeyType, ValueType> {
-    if (txn instanceof SimpleDbTransaction) {
-      return txn.store<KeyType, ValueType>(store);
-    } else {
-      return fail('Invalid transaction object provided!');
-    }
+    return txn.store<KeyType, ValueType>(store);
   }
 
   constructor(private db: IDBDatabase) {}
@@ -534,25 +528,19 @@ export class SimpleDbStore<KeyType extends IDBValidKey, ValueType> {
   }
 
   private cursor(options: IterateOptions): IDBRequest {
-    let direction = 'next';
+    let direction: IDBCursorDirection = 'next';
     if (options.reverse) {
       direction = 'prev';
     }
     if (options.index) {
       const index = this.store.index(options.index);
       if (options.keysOnly) {
-        return index.openKeyCursor(
-          options.range,
-          direction as AnyDuringMigration
-        );
+        return index.openKeyCursor(options.range, direction);
       } else {
-        return index.openCursor(options.range, direction as AnyDuringMigration);
+        return index.openCursor(options.range, direction);
       }
     } else {
-      return this.store.openCursor(
-        options.range,
-        direction as AnyDuringMigration
-      );
+      return this.store.openCursor(options.range, direction);
     }
   }
 }

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

@@ -65,7 +65,7 @@ describe('EagerGarbageCollector', () => {
         expect(referenceSet.isEmpty()).to.equal(false);
 
         referenceSet.removeReferencesForId(2);
-        return gc.collectGarbage(true).toPromise();
+        return gc.collectGarbage(null).toPromise();
       })
       .then(garbage => {
         expectSetToEqual(garbage, [key3]);