Add a type parameter to Persistence

packages/firestore/src/core/firestore_client.ts

@@ -30,7 +30,7 @@ import { IndexedDbPersistence } from '../local/indexeddb_persistence';
 import { LocalStore } from '../local/local_store';
 import { MemoryPersistence } from '../local/memory_persistence';
 import { NoOpGarbageCollector } from '../local/no_op_garbage_collector';
-import { Persistence } from '../local/persistence';
+import { Persistence, PersistenceTransaction } from '../local/persistence';
 import {
   DocumentKeySet,
   documentKeySet,
@@ -76,8 +76,10 @@ export class FirestoreClient {
   // undefined checks.
   private eventMgr: EventManager;
   private garbageCollector: GarbageCollector;
-  private persistence: Persistence;
-  private localStore: LocalStore;
+  // Note that localStore will end up with the same type parameter as
+  // persistence because the LocalStore constructor enforces this.
+  private persistence: Persistence<PersistenceTransaction>;
+  private localStore: LocalStore<PersistenceTransaction>;
   private remoteStore: RemoteStore;
   private syncEngine: SyncEngine;
 

packages/firestore/src/core/sync_engine.ts

@@ -41,7 +41,7 @@ import { ObjectMap } from '../util/obj_map';
 import { Deferred } from '../util/promise';
 import { SortedMap } from '../util/sorted_map';
 import { isNullOrUndefined } from '../util/types';
-
+import { PersistenceTransaction } from '../local/persistence';
 import { Query } from './query';
 import { SnapshotVersion } from './snapshot_version';
 import { TargetIdGenerator } from './target_id_generator';
@@ -143,7 +143,7 @@ export class SyncEngine implements RemoteSyncer {
   private targetIdGenerator = TargetIdGenerator.forSyncEngine();
 
   constructor(
-    private localStore: LocalStore,
+    private localStore: LocalStore<PersistenceTransaction>,
     private remoteStore: RemoteStore,
     private currentUser: User
   ) {

packages/firestore/src/local/indexeddb_mutation_queue.ts

@@ -38,12 +38,13 @@ import {
 } from './indexeddb_schema';
 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 { IndexedDbTransaction } from './indexeddb_persistence';
 
 /** A mutation queue for a specific user, backed by IndexedDB. */
-export class IndexedDbMutationQueue implements MutationQueue {
+export class IndexedDbMutationQueue
+  implements MutationQueue<IndexedDbTransaction> {
   /**
    * Next value to use when assigning sequential IDs to each mutation batch.
    *
@@ -87,7 +88,7 @@ export class IndexedDbMutationQueue implements MutationQueue {
     return new IndexedDbMutationQueue(userId, serializer);
   }
 
-  start(transaction: PersistenceTransaction): PersistencePromise<void> {
+  start(transaction: IndexedDbTransaction): PersistencePromise<void> {
     return IndexedDbMutationQueue.loadNextBatchIdFromDb(transaction)
       .next(nextBatchId => {
         this.nextBatchId = nextBatchId;
@@ -129,7 +130,7 @@ export class IndexedDbMutationQueue implements MutationQueue {
    * are no mutations returns 0. Note that batch IDs are global.
    */
   static loadNextBatchIdFromDb(
-    txn: PersistenceTransaction
+    txn: IndexedDbTransaction
   ): PersistencePromise<BatchId> {
     let maxBatchId = BATCHID_UNKNOWN;
     return mutationsStore(txn)
@@ -152,7 +153,7 @@ export class IndexedDbMutationQueue implements MutationQueue {
       .next(() => maxBatchId + 1);
   }
 
-  checkEmpty(transaction: PersistenceTransaction): PersistencePromise<boolean> {
+  checkEmpty(transaction: IndexedDbTransaction): PersistencePromise<boolean> {
     let empty = true;
     const range = IDBKeyRange.bound(
       this.keyForBatchId(Number.NEGATIVE_INFINITY),
@@ -167,19 +168,19 @@ export class IndexedDbMutationQueue implements MutationQueue {
   }
 
   getNextBatchId(
-    transaction: PersistenceTransaction
+    transaction: IndexedDbTransaction
   ): PersistencePromise<BatchId> {
     return PersistencePromise.resolve(this.nextBatchId);
   }
 
   getHighestAcknowledgedBatchId(
-    transaction: PersistenceTransaction
+    transaction: IndexedDbTransaction
   ): PersistencePromise<BatchId> {
     return PersistencePromise.resolve(this.metadata.lastAcknowledgedBatchId);
   }
 
   acknowledgeBatch(
-    transaction: PersistenceTransaction,
+    transaction: IndexedDbTransaction,
     batch: MutationBatch,
     streamToken: ProtoByteString
   ): PersistencePromise<void> {
@@ -196,21 +197,21 @@ export class IndexedDbMutationQueue implements MutationQueue {
   }
 
   getLastStreamToken(
-    transaction: PersistenceTransaction
+    transaction: IndexedDbTransaction
   ): PersistencePromise<ProtoByteString> {
     return PersistencePromise.resolve(this.metadata.lastStreamToken);
   }
 
   setLastStreamToken(
-    transaction: PersistenceTransaction,
+    transaction: IndexedDbTransaction,
     streamToken: ProtoByteString
   ): PersistencePromise<void> {
     this.metadata.lastStreamToken = convertStreamToken(streamToken);
     return mutationQueuesStore(transaction).put(this.metadata);
   }
 
   addMutationBatch(
-    transaction: PersistenceTransaction,
+    transaction: IndexedDbTransaction,
     localWriteTime: Timestamp,
     mutations: Mutation[]
   ): PersistencePromise<MutationBatch> {
@@ -245,7 +246,7 @@ export class IndexedDbMutationQueue implements MutationQueue {
   }
 
   lookupMutationBatch(
-    transaction: PersistenceTransaction,
+    transaction: IndexedDbTransaction,
     batchId: BatchId
   ): PersistencePromise<MutationBatch | null> {
     return mutationsStore(transaction)
@@ -257,7 +258,7 @@ export class IndexedDbMutationQueue implements MutationQueue {
   }
 
   getNextMutationBatchAfterBatchId(
-    transaction: PersistenceTransaction,
+    transaction: IndexedDbTransaction,
     batchId: BatchId
   ): PersistencePromise<MutationBatch | null> {
     // All batches with batchId <= this.metadata.lastAcknowledgedBatchId have
@@ -283,7 +284,7 @@ export class IndexedDbMutationQueue implements MutationQueue {
   }
 
   getAllMutationBatches(
-    transaction: PersistenceTransaction
+    transaction: IndexedDbTransaction
   ): PersistencePromise<MutationBatch[]> {
     const range = IDBKeyRange.bound(
       this.keyForBatchId(BATCHID_UNKNOWN),
@@ -297,7 +298,7 @@ export class IndexedDbMutationQueue implements MutationQueue {
   }
 
   getAllMutationBatchesThroughBatchId(
-    transaction: PersistenceTransaction,
+    transaction: IndexedDbTransaction,
     batchId: BatchId
   ): PersistencePromise<MutationBatch[]> {
     const range = IDBKeyRange.bound(
@@ -312,7 +313,7 @@ export class IndexedDbMutationQueue implements MutationQueue {
   }
 
   getAllMutationBatchesAffectingDocumentKey(
-    transaction: PersistenceTransaction,
+    transaction: IndexedDbTransaction,
     documentKey: DocumentKey
   ): PersistencePromise<MutationBatch[]> {
     // Scan the document-mutation index starting with a prefix starting with
@@ -363,7 +364,7 @@ export class IndexedDbMutationQueue implements MutationQueue {
   }
 
   getAllMutationBatchesAffectingQuery(
-    transaction: PersistenceTransaction,
+    transaction: IndexedDbTransaction,
     query: Query
   ): PersistencePromise<MutationBatch[]> {
     assert(
@@ -439,7 +440,7 @@ export class IndexedDbMutationQueue implements MutationQueue {
   }
 
   removeMutationBatches(
-    transaction: PersistenceTransaction,
+    transaction: IndexedDbTransaction,
     batches: MutationBatch[]
   ): PersistencePromise<void> {
     const txn = mutationsStore(transaction);
@@ -477,9 +478,7 @@ export class IndexedDbMutationQueue implements MutationQueue {
     return PersistencePromise.waitFor(promises);
   }
 
-  performConsistencyCheck(
-    txn: PersistenceTransaction
-  ): PersistencePromise<void> {
+  performConsistencyCheck(txn: IndexedDbTransaction): PersistencePromise<void> {
     return this.checkEmpty(txn).next(empty => {
       if (!empty) {
         return PersistencePromise.resolve();
@@ -517,7 +516,7 @@ export class IndexedDbMutationQueue implements MutationQueue {
   }
 
   containsKey(
-    txn: PersistenceTransaction,
+    txn: IndexedDbTransaction,
     key: DocumentKey
   ): PersistencePromise<boolean> {
     const indexKey = DbDocumentMutation.prefixForPath(this.userId, key.path);
@@ -561,10 +560,10 @@ function convertStreamToken(token: ProtoByteString): string {
  * Helper to get a typed SimpleDbStore for the mutations object store.
  */
 function mutationsStore(
-  txn: PersistenceTransaction
+  txn: IndexedDbTransaction
 ): SimpleDbStore<DbMutationBatchKey, DbMutationBatch> {
   return SimpleDb.getStore<DbMutationBatchKey, DbMutationBatch>(
-    txn,
+    txn.simpleDbTransaction,
     DbMutationBatch.store
   );
 }
@@ -573,10 +572,10 @@ function mutationsStore(
  * Helper to get a typed SimpleDbStore for the mutationQueues object store.
  */
 function documentMutationsStore(
-  txn: PersistenceTransaction
+  txn: IndexedDbTransaction
 ): SimpleDbStore<DbDocumentMutationKey, DbDocumentMutation> {
   return SimpleDb.getStore<DbDocumentMutationKey, DbDocumentMutation>(
-    txn,
+    txn.simpleDbTransaction,
     DbDocumentMutation.store
   );
 }
@@ -585,10 +584,10 @@ function documentMutationsStore(
  * Helper to get a typed SimpleDbStore for the mutationQueues object store.
  */
 function mutationQueuesStore(
-  txn: PersistenceTransaction
+  txn: IndexedDbTransaction
 ): SimpleDbStore<DbMutationQueueKey, DbMutationQueue> {
   return SimpleDb.getStore<DbMutationQueueKey, DbMutationQueue>(
-    txn,
+    txn.simpleDbTransaction,
     DbMutationQueue.store
   );
 }

packages/firestore/src/local/indexeddb_persistence.ts

@@ -59,6 +59,10 @@ const UNSUPPORTED_PLATFORM_ERROR_MSG =
   ' IndexedDB or is known to have an incomplete implementation. Offline' +
   ' persistence has been disabled.';
 
+export class IndexedDbTransaction {
+  constructor(readonly simpleDbTransaction: SimpleDbTransaction) {}
+}
+
 /**
  * An IndexedDB-backed instance of Persistence. Data is stored persistently
  * across sessions.
@@ -88,7 +92,7 @@ const UNSUPPORTED_PLATFORM_ERROR_MSG =
  * which acts as an indicator that another tab should go ahead and take the
  * owner lease immediately regardless of the current lease timestamp.
  */
-export class IndexedDbPersistence implements Persistence {
+export class IndexedDbPersistence implements Persistence<IndexedDbTransaction> {
   /**
    * The name of the main (and currently only) IndexedDB database. this name is
    * appended to the prefix provided to the IndexedDbPersistence constructor.
@@ -162,21 +166,21 @@ export class IndexedDbPersistence implements Persistence {
     return this._started;
   }
 
-  getMutationQueue(user: User): MutationQueue {
+  getMutationQueue(user: User): MutationQueue<IndexedDbTransaction> {
     return IndexedDbMutationQueue.forUser(user, this.serializer);
   }
 
-  getQueryCache(): QueryCache {
+  getQueryCache(): QueryCache<IndexedDbTransaction> {
     return new IndexedDbQueryCache(this.serializer);
   }
 
-  getRemoteDocumentCache(): RemoteDocumentCache {
+  getRemoteDocumentCache(): RemoteDocumentCache<IndexedDbTransaction> {
     return new IndexedDbRemoteDocumentCache(this.serializer);
   }
 
   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 +190,17 @@ 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.
+        const txn = new IndexedDbTransaction(simpleDbTxn);
+        return this.ensureOwnerLease(txn.simpleDbTransaction).next(() =>
+          operation(txn)
+        );
+      }
+    );
   }
 
   static isAvailable(): boolean {
@@ -326,12 +337,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