Add Transaction (#3153)

Author: schmidt-sebastian

packages/firestore/lite/index.node.ts

@@ -60,6 +60,8 @@ export { DocumentSnapshot, QueryDocumentSnapshot } from './src/api/snapshot';
 
 export { WriteBatch, writeBatch } from './src/api/write_batch';
 
+export { Transaction, runTransaction } from './src/api/transaction';
+
 export { setLogLevel } from '../src/util/log';
 
 export function registerFirestore(): void {

packages/firestore/lite/src/api/transaction.ts

@@ -0,0 +1,181 @@
+/**
+ * @license
+ * Copyright 2020 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import * as firestore from '../../';
+
+import { UserDataReader } from '../../../src/api/user_data_reader';
+import { Transaction as InternalTransaction } from '../../../src/core/transaction';
+import {
+  Document,
+  MaybeDocument,
+  NoDocument
+} from '../../../src/model/document';
+import { fail } from '../../../src/util/assert';
+import { applyFirestoreDataConverter } from '../../../src/api/database';
+import { DocumentSnapshot } from './snapshot';
+import { Firestore } from './database';
+import { TransactionRunner } from '../../../src/core/transaction_runner';
+import { AsyncQueue } from '../../../src/util/async_queue';
+import { Deferred } from '../../../src/util/promise';
+import { FieldPath as ExternalFieldPath } from '../../../src/api/field_path';
+import { validateReference } from './write_batch';
+import { newUserDataReader } from './reference';
+import { FieldPath } from './field_path';
+import { cast } from './util';
+
+export class Transaction implements firestore.Transaction {
+  // This is the lite version of the Transaction API used in the legacy SDK. The
+  // class is a close copy but takes different input types.
+
+  private readonly _dataReader: UserDataReader;
+
+  constructor(
+    private readonly _firestore: Firestore,
+    private readonly _transaction: InternalTransaction
+  ) {
+    // Kick off configuring the client, which freezes the settings.
+    // eslint-disable-next-line @typescript-eslint/no-floating-promises
+    _firestore._ensureClientConfigured();
+    this._dataReader = newUserDataReader(
+      _firestore._databaseId,
+      _firestore._settings!
+    );
+  }
+
+  get<T>(
+    documentRef: firestore.DocumentReference<T>
+  ): Promise<firestore.DocumentSnapshot<T>> {
+    const ref = validateReference(documentRef, this._firestore);
+    return this._transaction
+      .lookup([ref._key])
+      .then((docs: MaybeDocument[]) => {
+        if (!docs || docs.length !== 1) {
+          return fail('Mismatch in docs returned from document lookup.');
+        }
+        const doc = docs[0];
+        if (doc instanceof NoDocument) {
+          return new DocumentSnapshot<T>(
+            this._firestore,
+            ref._key,
+            null,
+            ref._converter
+          );
+        } else if (doc instanceof Document) {
+          return new DocumentSnapshot<T>(
+            this._firestore,
+            doc.key,
+            doc,
+            ref._converter
+          );
+        } else {
+          throw fail(
+            `BatchGetDocumentsRequest returned unexpected document type: ${doc.constructor.name}`
+          );
+        }
+      });
+  }
+
+  set<T>(documentRef: firestore.DocumentReference<T>, value: T): Transaction;
+  set<T>(
+    documentRef: firestore.DocumentReference<T>,
+    value: Partial<T>,
+    options: firestore.SetOptions
+  ): Transaction;
+  set<T>(
+    documentRef: firestore.DocumentReference<T>,
+    value: T,
+    options?: firestore.SetOptions
+  ): Transaction {
+    const ref = validateReference(documentRef, this._firestore);
+    const [convertedValue] = applyFirestoreDataConverter(
+      ref._converter,
+      value,
+      'Transaction.set'
+    );
+    const parsed = this._dataReader.parseSetData(
+      'Transaction.set',
+      convertedValue,
+      options
+    );
+    this._transaction.set(ref._key, parsed);
+    return this;
+  }
+
+  update(
+    documentRef: firestore.DocumentReference<unknown>,
+    value: firestore.UpdateData
+  ): Transaction;
+  update(
+    documentRef: firestore.DocumentReference<unknown>,
+    field: string | ExternalFieldPath,
+    value: unknown,
+    ...moreFieldsAndValues: unknown[]
+  ): Transaction;
+  update(
+    documentRef: firestore.DocumentReference<unknown>,
+    fieldOrUpdateData: string | ExternalFieldPath | firestore.UpdateData,
+    value?: unknown,
+    ...moreFieldsAndValues: unknown[]
+  ): Transaction {
+    const ref = validateReference(documentRef, this._firestore);
+
+    let parsed;
+    if (
+      typeof fieldOrUpdateData === 'string' ||
+      fieldOrUpdateData instanceof FieldPath
+    ) {
+      parsed = this._dataReader.parseUpdateVarargs(
+        'Transaction.update',
+        fieldOrUpdateData,
+        value,
+        moreFieldsAndValues
+      );
+    } else {
+      parsed = this._dataReader.parseUpdateData(
+        'Transaction.update',
+        fieldOrUpdateData
+      );
+    }
+
+    this._transaction.update(ref._key, parsed);
+    return this;
+  }
+
+  delete(documentRef: firestore.DocumentReference<unknown>): Transaction {
+    const ref = validateReference(documentRef, this._firestore);
+    this._transaction.delete(ref._key);
+    return this;
+  }
+}
+
+export function runTransaction<T>(
+  firestore: firestore.FirebaseFirestore,
+  updateFunction: (transaction: firestore.Transaction) => Promise<T>
+): Promise<T> {
+  const firestoreClient = cast(firestore, Firestore);
+  return firestoreClient._ensureClientConfigured().then(async datastore => {
+    const deferred = new Deferred<T>();
+    new TransactionRunner<T>(
+      new AsyncQueue(),
+      datastore,
+      internalTransaction =>
+        updateFunction(new Transaction(firestoreClient, internalTransaction)),
+      deferred
+    ).run();
+    return deferred.promise;
+  });
+}

packages/firestore/lite/src/api/write_batch.ts

@@ -37,11 +37,11 @@ export class WriteBatch implements firestore.WriteBatch {
   // This is the lite version of the WriteBatch API used in the legacy SDK. The
   // class is a close copy but takes different input types.
 
+  private readonly _dataReader: UserDataReader;
   private _mutations = [] as Mutation[];
   private _committed = false;
-  private _dataReader: UserDataReader;
 
-  constructor(private _firestore: Firestore) {
+  constructor(private readonly _firestore: Firestore) {
     // Kick off configuring the client, which freezes the settings.
     // eslint-disable-next-line @typescript-eslint/no-floating-promises
     _firestore._ensureClientConfigured();

packages/firestore/lite/test/integration.test.ts

@@ -52,9 +52,9 @@ import {
   DEFAULT_SETTINGS
 } from '../../test/integration/util/settings';
 import { writeBatch } from '../src/api/write_batch';
+import { runTransaction } from '../src/api/transaction';
 import { expectEqual, expectNotEqual } from '../../test/util/helpers';
 import { FieldValue } from '../../src/api/field_value';
-
 use(chaiAsPromised);
 
 describe('Firestore', () => {
@@ -324,7 +324,95 @@ describe('WriteBatch', () => {
   });
 });
 
-function genericMutationTests(op: MutationTester): void {
+describe('Transaction', () => {
+  class TransactionTester implements MutationTester {
+    delete(ref: firestore.DocumentReference<unknown>): Promise<void> {
+      return runTransaction(ref.firestore, async transaction => {
+        transaction.delete(ref);
+      });
+    }
+
+    set<T>(
+      ref: firestore.DocumentReference<T>,
+      data: T | Partial<T>,
+      options?: firestore.SetOptions
+    ): Promise<void> {
+      const args = Array.from(arguments);
+      return runTransaction(ref.firestore, async transaction => {
+        // TODO(mrschmidt): Find a way to remove the `any` cast here
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        (transaction.set as any).apply(transaction, args);
+      });
+    }
+
+    update(
+      ref: firestore.DocumentReference<unknown>,
+      dataOrField: firestore.UpdateData | string | firestore.FieldPath,
+      value?: unknown,
+      ...moreFieldsAndValues: unknown[]
+    ): Promise<void> {
+      const args = Array.from(arguments);
+      return runTransaction(ref.firestore, async transaction => {
+        // TODO(mrschmidt): Find a way to remove the `any` cast here
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        (transaction.update as any).apply(transaction, args);
+      });
+    }
+  }
+
+  genericMutationTests(
+    new TransactionTester(),
+    /* validationUsesPromises= */ true
+  );
+
+  it('can read and then write', () => {
+    return withTestDocAndInitialData({ counter: 1 }, async doc => {
+      await runTransaction(doc.firestore, async transaction => {
+        const snap = await transaction.get(doc);
+        transaction.update(doc, 'counter', snap.get('counter') + 1);
+      });
+      const result = await getDoc(doc);
+      expect(result.get('counter')).to.equal(2);
+    });
+  });
+
+  it('can read non-existing doc then write', () => {
+    return withTestDoc(async doc => {
+      await runTransaction(doc.firestore, async transaction => {
+        const snap = await transaction.get(doc);
+        expect(snap.exists()).to.be.false;
+        transaction.set(doc, { counter: 1 });
+      });
+      const result = await getDoc(doc);
+      expect(result.get('counter')).to.equal(1);
+    });
+  });
+
+  it('retries when document is modified', () => {
+    return withTestDoc(async doc => {
+      let retryCounter = 0;
+      await runTransaction(doc.firestore, async transaction => {
+        ++retryCounter;
+        await transaction.get(doc);
+
+        if (retryCounter === 1) {
+          // Out of band modification that doesn't use the transaction
+          await setDoc(doc, { counter: 'invalid' });
+        }
+
+        transaction.set(doc, { counter: 1 });
+      });
+      expect(retryCounter).to.equal(2);
+      const result = await getDoc(doc);
+      expect(result.get('counter')).to.equal(1);
+    });
+  });
+});
+
+function genericMutationTests(
+  op: MutationTester,
+  validationUsesPromises: boolean = false
+): void {
   const setDoc = op.set;
   const updateDoc = op.update;
   const deleteDoc = op.delete;
@@ -379,9 +467,17 @@ function genericMutationTests(op: MutationTester): void {
 
     it('throws when user input fails validation', () => {
       return withTestDoc(async docRef => {
-        expect(() => setDoc(docRef, { val: undefined })).to.throw(
-          /Function .* called with invalid data. Unsupported field value: undefined \(found in field val\)/
-        );
+        if (validationUsesPromises) {
+          return expect(
+            setDoc(docRef, { val: undefined })
+          ).to.eventually.be.rejectedWith(
+            /Function .* called with invalid data. Unsupported field value: undefined \(found in field val\)/
+          );
+        } else {
+          expect(() => setDoc(docRef, { val: undefined })).to.throw(
+            /Function .* called with invalid data. Unsupported field value: undefined \(found in field val\)/
+          );
+        }
       });
     });
 
@@ -418,9 +514,17 @@ function genericMutationTests(op: MutationTester): void {
 
     it('throws when user input fails validation', () => {
       return withTestDoc(async docRef => {
-        expect(() => updateDoc(docRef, { val: undefined })).to.throw(
-          /Function .* called with invalid data. Unsupported field value: undefined \(found in field val\)/
-        );
+        if (validationUsesPromises) {
+          return expect(
+            updateDoc(docRef, { val: undefined })
+          ).to.eventually.be.rejectedWith(
+            /Function .* called with invalid data. Unsupported field value: undefined \(found in field val\)/
+          );
+        } else {
+          expect(() => updateDoc(docRef, { val: undefined })).to.throw(
+            /Function .* called with invalid data. Unsupported field value: undefined \(found in field val\)/
+          );
+        }
       });
     });
   });

packages/firestore/src/util/promise.ts

@@ -23,7 +23,7 @@ export interface Rejecter {
   (reason?: Error): void;
 }
 
-export class Deferred<R> {
+export class Deferred<R = void> {
   promise: Promise<R>;
   // Assigned synchronously in constructor by Promise constructor callback.
   resolve!: Resolver<R>;