Compat class for WriteBatch

schmidt-sebastian · schmidt-sebastian · commit 4fc524dac72c · 2020-11-10T13:33:31.000-08:00
diff --git a/packages/firestore/exp/src/api/write_batch.ts b/packages/firestore/exp/src/api/write_batch.ts
@@ -20,6 +20,8 @@ import { FirebaseFirestore } from './database';
 import { executeWrite } from './reference';
 import { ensureFirestoreConfigured } from '../../../src/api/database';
 
+export { WriteBatch };
+
 /**
  * Creates a write batch, used for performing multiple writes as a single
  * atomic operation. The maximum number of writes allowed in a single WriteBatch
diff --git a/packages/firestore/exp/test/shim.ts b/packages/firestore/exp/test/shim.ts
@@ -127,63 +127,6 @@ export class Transaction
   }
 }
 
-export class WriteBatch
-  extends Compat<exp.WriteBatch>
-  implements legacy.WriteBatch {
-  set<T>(
-    documentRef: DocumentReference<T>,
-    data: T,
-    options?: legacy.SetOptions
-  ): WriteBatch {
-    if (options) {
-      validateSetOptions('WriteBatch.set', options);
-      this._delegate.set(documentRef._delegate, unwrap(data), options);
-    } else {
-      this._delegate.set(documentRef._delegate, unwrap(data));
-    }
-    return this;
-  }
-
-  update(
-    documentRef: DocumentReference<any>,
-    data: legacy.UpdateData
-  ): WriteBatch;
-  update(
-    documentRef: DocumentReference<any>,
-    field: string | FieldPath,
-    value: any,
-    ...moreFieldsAndValues: any[]
-  ): WriteBatch;
-  update(
-    documentRef: DocumentReference<any>,
-    dataOrField: any,
-    value?: any,
-    ...moreFieldsAndValues: any[]
-  ): WriteBatch {
-    if (arguments.length === 2) {
-      this._delegate.update(documentRef._delegate, unwrap(dataOrField));
-    } else {
-      this._delegate.update(
-        documentRef._delegate,
-        unwrap(dataOrField),
-        unwrap(value),
-        ...unwrap(moreFieldsAndValues)
-      );
-    }
-
-    return this;
-  }
-
-  delete(documentRef: DocumentReference<any>): WriteBatch {
-    this._delegate.delete(documentRef._delegate);
-    return this;
-  }
-
-  commit(): Promise<void> {
-    return this._delegate.commit();
-  }
-}
-
 export class DocumentSnapshot<T = legacy.DocumentData>
   extends Compat<exp.DocumentSnapshot<T>>
   implements legacy.DocumentSnapshot<T> {
diff --git a/packages/firestore/lite/src/api/reference.ts b/packages/firestore/lite/src/api/reference.ts
@@ -78,6 +78,7 @@ import {
 import { newSerializer } from '../../../src/platform/serializer';
 import { Code, FirestoreError } from '../../../src/util/error';
 import { getDatastore } from './components';
+import { Compat } from '../../../src/compat/compat';
 
 /**
  * Document data (for use with {@link setDoc()}) consists of fields mapped to
@@ -1086,6 +1087,12 @@ export function updateDoc(
 ): Promise<void> {
   const dataReader = newUserDataReader(reference.firestore);
 
+  // For Compat types, we have to "extract" the underlying types before
+  // performing validation.
+  if (fieldOrUpdateData instanceof Compat) {
+    fieldOrUpdateData = fieldOrUpdateData._delegate;
+  }
+
   let parsed: ParsedUpdateData;
   if (
     typeof fieldOrUpdateData === 'string' ||
diff --git a/packages/firestore/lite/src/api/transaction.ts b/packages/firestore/lite/src/api/transaction.ts
@@ -43,6 +43,7 @@ import {
 } from './reference';
 import { FieldPath } from './field_path';
 import { getDatastore } from './components';
+import { Compat } from '../../../src/compat/compat';
 
 // TODO(mrschmidt) Consider using `BaseTransaction` as the base class in the
 // legacy SDK.
@@ -194,6 +195,12 @@ export class Transaction {
   ): this {
     const ref = validateReference(documentRef, this._firestore);
 
+    // For Compat types, we have to "extract" the underlying types before
+    // performing validation.
+    if (fieldOrUpdateData instanceof Compat) {
+      fieldOrUpdateData = fieldOrUpdateData._delegate;
+    }
+
     let parsed;
     if (
       typeof fieldOrUpdateData === 'string' ||
diff --git a/packages/firestore/lite/src/api/write_batch.ts b/packages/firestore/lite/src/api/write_batch.ts
@@ -38,6 +38,7 @@ import { FirebaseFirestore } from './database';
 import { invokeCommitRpc } from '../../../src/remote/datastore';
 import { FieldPath } from './field_path';
 import { getDatastore } from './components';
+import { Compat } from '../../../src/compat/compat';
 
 /**
  * A write batch, used to perform multiple writes as a single atomic unit.
@@ -155,8 +156,13 @@ export class WriteBatch {
     this.verifyNotCommitted();
     const ref = validateReference(documentRef, this._firestore);
 
-    let parsed;
+    // For Compat types, we have to "extract" the underlying types before
+    // performing validation.
+    if (fieldOrUpdateData instanceof Compat) {
+      fieldOrUpdateData = fieldOrUpdateData._delegate;
+    }
 
+    let parsed;
     if (
       typeof fieldOrUpdateData === 'string' ||
       fieldOrUpdateData instanceof FieldPath
@@ -233,9 +239,12 @@ export class WriteBatch {
 }
 
 export function validateReference<T>(
-  documentRef: DocumentReference<T>,
+  documentRef: DocumentReference<T> | Compat<DocumentReference<T>>,
   firestore: FirebaseFirestore
 ): DocumentReference<T> {
+  if (documentRef instanceof Compat) {
+    documentRef = documentRef._delegate;
+  }
   if (documentRef.firestore !== firestore) {
     throw new FirestoreError(
       Code.INVALID_ARGUMENT,
diff --git a/packages/firestore/src/api/database.ts b/packages/firestore/src/api/database.ts
@@ -27,8 +27,7 @@ import {
   firestoreClientGetDocumentsFromLocalCache,
   firestoreClientGetDocumentsViaSnapshotListener,
   firestoreClientListen,
-  firestoreClientTransaction,
-  firestoreClientWrite
+  firestoreClientTransaction
 } from '../core/firestore_client';
 import {
   Bound,
@@ -59,7 +58,6 @@ import { Transaction as InternalTransaction } from '../core/transaction';
 import { ChangeType, ViewSnapshot } from '../core/view_snapshot';
 import { Document, MaybeDocument, NoDocument } from '../model/document';
 import { DocumentKey } from '../model/document_key';
-import { DeleteMutation, Mutation, Precondition } from '../model/mutation';
 import { FieldPath, ResourcePath } from '../model/path';
 import { isServerTimestamp } from '../model/server_timestamps';
 import { refValue } from '../model/values';
@@ -118,11 +116,13 @@ import {
   getDocFromCache,
   getDocFromServer,
   getDoc,
-  onSnapshot
+  onSnapshot,
+  executeWrite
 } from '../../exp/src/api/reference';
 import { DocumentSnapshot as ExpDocumentSnapshot } from '../../exp/src/api/snapshot';
 import { LRU_COLLECTION_DISABLED } from '../local/lru_garbage_collector';
 import { Compat } from '../compat/compat';
+import { WriteBatch as ExpWriteBatch } from '../../exp/src/api/write_batch';
 
 import {
   CollectionReference as PublicCollectionReference,
@@ -399,7 +399,11 @@ export class Firestore
 
   batch(): PublicWriteBatch {
     ensureFirestoreConfigured(this._delegate);
-    return new WriteBatch(this);
+    return new WriteBatch(
+      new ExpWriteBatch(this._delegate, mutations =>
+        executeWrite(this._delegate, mutations)
+      )
+    );
   }
 }
 
@@ -590,15 +594,9 @@ export class Transaction implements PublicTransaction {
   }
 }
 
-export class WriteBatch implements PublicWriteBatch {
-  private _mutations = [] as Mutation[];
-  private _committed = false;
-  private _dataReader: UserDataReader;
-
-  constructor(private _firestore: Firestore) {
-    this._dataReader = newUserDataReader(this._firestore._delegate);
-  }
-
+export class WriteBatch
+  extends Compat<ExpWriteBatch>
+  implements PublicWriteBatch {
   set<T>(
     documentRef: DocumentReference<T>,
     data: Partial<T>,
@@ -607,38 +605,22 @@ export class WriteBatch implements PublicWriteBatch {
   set<T>(documentRef: DocumentReference<T>, data: T): WriteBatch;
   set<T>(
     documentRef: PublicDocumentReference<T>,
-    value: T | Partial<T>,
+    data: T | Partial<T>,
     options?: PublicSetOptions
   ): WriteBatch {
-    this.verifyNotCommitted();
-    const ref = validateReference(
-      'WriteBatch.set',
-      documentRef,
-      this._firestore
-    );
-    options = validateSetOptions('WriteBatch.set', options);
-    const convertedValue = applyFirestoreDataConverter(
-      ref._converter,
-      value,
-      options
-    );
-    const parsed = parseSetData(
-      this._dataReader,
-      'WriteBatch.set',
-      ref._key,
-      convertedValue,
-      ref._converter !== null,
-      options
-    );
-    this._mutations = this._mutations.concat(
-      parsed.toMutations(ref._key, Precondition.none())
-    );
+    const ref = castReference(documentRef);
+    if (options) {
+      validateSetOptions('WriteBatch.set', options);
+      this._delegate.set(ref, data, options);
+    } else {
+      this._delegate.set(ref, data);
+    }
     return this;
   }
 
   update(
     documentRef: PublicDocumentReference<unknown>,
-    value: PublicUpdateData
+    data: PublicUpdateData
   ): WriteBatch;
   update(
     documentRef: PublicDocumentReference<unknown>,
@@ -648,83 +630,32 @@ export class WriteBatch implements PublicWriteBatch {
   ): WriteBatch;
   update(
     documentRef: PublicDocumentReference<unknown>,
-    fieldOrUpdateData: string | PublicFieldPath | PublicUpdateData,
+    dataOrField: string | PublicFieldPath | PublicUpdateData,
     value?: unknown,
     ...moreFieldsAndValues: unknown[]
   ): WriteBatch {
-    this.verifyNotCommitted();
-    const ref = validateReference(
-      'WriteBatch.update',
-      documentRef,
-      this._firestore
-    );
-
-    // For Compat types, we have to "extract" the underlying types before
-    // performing validation.
-    if (fieldOrUpdateData instanceof Compat) {
-      fieldOrUpdateData = (fieldOrUpdateData as Compat<ExpFieldPath>)._delegate;
-    }
-
-    let parsed;
-    if (
-      typeof fieldOrUpdateData === 'string' ||
-      fieldOrUpdateData instanceof ExpFieldPath
-    ) {
-      parsed = parseUpdateVarargs(
-        this._dataReader,
-        'WriteBatch.update',
-        ref._key,
-        fieldOrUpdateData,
-        value,
-        moreFieldsAndValues
-      );
+    const ref = castReference(documentRef);
+    if (arguments.length === 2) {
+      this._delegate.update(ref, dataOrField as PublicUpdateData);
     } else {
-      parsed = parseUpdateData(
-        this._dataReader,
-        'WriteBatch.update',
-        ref._key,
-        fieldOrUpdateData
+      this._delegate.update(
+        ref,
+        dataOrField as string | ExpFieldPath,
+        value,
+        ...moreFieldsAndValues
       );
     }
-
-    this._mutations = this._mutations.concat(
-      parsed.toMutations(ref._key, Precondition.exists(true))
-    );
     return this;
   }
 
   delete(documentRef: PublicDocumentReference<unknown>): WriteBatch {
-    this.verifyNotCommitted();
-    const ref = validateReference(
-      'WriteBatch.delete',
-      documentRef,
-      this._firestore
-    );
-    this._mutations = this._mutations.concat(
-      new DeleteMutation(ref._key, Precondition.none())
-    );
+    const ref = castReference(documentRef);
+    this._delegate.delete(ref);
     return this;
   }
 
   commit(): Promise<void> {
-    this.verifyNotCommitted();
-    this._committed = true;
-    if (this._mutations.length > 0) {
-      const client = ensureFirestoreConfigured(this._firestore._delegate);
-      return firestoreClientWrite(client, this._mutations);
-    }
-
-    return Promise.resolve();
-  }
-
-  private verifyNotCommitted(): void {
-    if (this._committed) {
-      throw new FirestoreError(
-        Code.FAILED_PRECONDITION,
-        'A write batch can no longer be used after commit() ' +
-          'has been called.'
-      );
-    }
+    return this._delegate.commit();
   }
 }
 
@@ -2030,6 +1961,15 @@ export class CollectionReference<T = PublicDocumentData>
   }
 }
 
+function castReference<T>(
+  documentRef: PublicDocumentReference<T>
+): ExpDocumentReference<T> {
+  if (documentRef instanceof Compat) {
+    documentRef = documentRef._delegate;
+  }
+  return cast<ExpDocumentReference<T>>(documentRef, ExpDocumentReference);
+}
+
 function validateReference<T>(
   methodName: string,
   documentRef: PublicDocumentReference<T>,
