Return correct DocumentSnapshot type in exp SDK (#3360)

.changeset/gorgeous-beers-build.md

@@ -0,0 +1,2 @@
+---
+---

packages/firestore/exp/index.node.ts

@@ -54,7 +54,7 @@ export {
   parent
 } from '../lite/src/api/reference';
 
-export { runTransaction, Transaction } from '../lite/src/api/transaction';
+export { runTransaction, Transaction } from './src/api/transaction';
 
 export {
   getDoc,

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

@@ -0,0 +1,81 @@
+/**
+ * @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 '../../index';
+
+import { Transaction as LiteTransaction } from '../../../lite/src/api/transaction';
+import { DocumentSnapshot } from './snapshot';
+import { TransactionRunner } from '../../../src/core/transaction_runner';
+import { AsyncQueue } from '../../../src/util/async_queue';
+import { cast } from '../../../lite/src/api/util';
+import { Firestore } from './database';
+import { Deferred } from '../../../src/util/promise';
+import { SnapshotMetadata } from '../../../src/api/database';
+import { Transaction as InternalTransaction } from '../../../src/core/transaction';
+import { validateReference } from '../../../lite/src/api/write_batch';
+
+export class Transaction extends LiteTransaction
+  implements firestore.Transaction {
+  // This class implements the same logic as the Transaction API in the Lite SDK
+  // but is subclassed in order to return its own DocumentSnapshot types.
+
+  constructor(
+    protected readonly _firestore: Firestore,
+    _transaction: InternalTransaction
+  ) {
+    super(_firestore, _transaction);
+  }
+
+  get<T>(
+    documentRef: firestore.DocumentReference<T>
+  ): Promise<DocumentSnapshot<T>> {
+    const ref = validateReference<T>(documentRef, this._firestore);
+    return super
+      .get(documentRef)
+      .then(
+        liteDocumentSnapshot =>
+          new DocumentSnapshot(
+            this._firestore,
+            ref._key,
+            liteDocumentSnapshot._document,
+            new SnapshotMetadata(
+              /* hasPendingWrites= */ false,
+              /* fromCache= */ false
+            ),
+            ref._converter
+          )
+      );
+  }
+}
+
+export function runTransaction<T>(
+  firestore: firestore.FirebaseFirestore,
+  updateFunction: (transaction: firestore.Transaction) => Promise<T>
+): Promise<T> {
+  const firestoreClient = cast(firestore, Firestore);
+  return firestoreClient._getDatastore().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/transaction.ts

@@ -42,22 +42,26 @@ 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.
+// TODO(mrschmidt) Consider using `BaseTransaction` as the base class in the
+// legacy SDK.
+export class Transaction {
+  // This is the tree-shakeable version of the Transaction class used in the
+  // legacy SDK. The class is a close copy but takes different input and output
+  // types. The firestore-exp SDK further extends this class to return its API
+  // type.
 
   private readonly _dataReader: UserDataReader;
 
   constructor(
-    private readonly _firestore: Firestore,
+    protected readonly _firestore: Firestore,
     private readonly _transaction: InternalTransaction
   ) {
     this._dataReader = newUserDataReader(_firestore);
   }
 
   get<T>(
     documentRef: firestore.DocumentReference<T>
-  ): Promise<firestore.DocumentSnapshot<T>> {
+  ): Promise<DocumentSnapshot<T>> {
     const ref = validateReference(documentRef, this._firestore);
     return this._transaction
       .lookup([ref._key])
@@ -67,14 +71,14 @@ export class Transaction implements firestore.Transaction {
         }
         const doc = docs[0];
         if (doc instanceof NoDocument) {
-          return new DocumentSnapshot<T>(
+          return new DocumentSnapshot(
             this._firestore,
             ref._key,
             null,
             ref._converter
           );
         } else if (doc instanceof Document) {
-          return new DocumentSnapshot<T>(
+          return new DocumentSnapshot(
             this._firestore,
             doc.key,
             doc,
@@ -88,17 +92,17 @@ export class Transaction implements firestore.Transaction {
       });
   }
 
-  set<T>(documentRef: firestore.DocumentReference<T>, value: T): Transaction;
+  set<T>(documentRef: firestore.DocumentReference<T>, value: T): this;
   set<T>(
     documentRef: firestore.DocumentReference<T>,
     value: Partial<T>,
     options: firestore.SetOptions
-  ): Transaction;
+  ): this;
   set<T>(
     documentRef: firestore.DocumentReference<T>,
     value: T,
     options?: firestore.SetOptions
-  ): Transaction {
+  ): this {
     const ref = validateReference(documentRef, this._firestore);
     const convertedValue = applyFirestoreDataConverter(
       ref._converter,
@@ -120,19 +124,19 @@ export class Transaction implements firestore.Transaction {
   update(
     documentRef: firestore.DocumentReference<unknown>,
     value: firestore.UpdateData
-  ): Transaction;
+  ): this;
   update(
     documentRef: firestore.DocumentReference<unknown>,
     field: string | ExternalFieldPath,
     value: unknown,
     ...moreFieldsAndValues: unknown[]
-  ): Transaction;
+  ): this;
   update(
     documentRef: firestore.DocumentReference<unknown>,
     fieldOrUpdateData: string | ExternalFieldPath | firestore.UpdateData,
     value?: unknown,
     ...moreFieldsAndValues: unknown[]
-  ): Transaction {
+  ): this {
     const ref = validateReference(documentRef, this._firestore);
 
     let parsed;
@@ -161,7 +165,7 @@ export class Transaction implements firestore.Transaction {
     return this;
   }
 
-  delete(documentRef: firestore.DocumentReference<unknown>): Transaction {
+  delete(documentRef: firestore.DocumentReference<unknown>): this {
     const ref = validateReference(documentRef, this._firestore);
     this._transaction.delete(ref._key);
     return this;