Compat class for WriteBatch (#4053)

* Compat class for WriteBatch * Fix merge * Lint * Create blue-geese-approve.md
firebase · Nov 13, 2020 · 6c6c49a · 6c6c49a
1 parent ad44c19
commit 6c6c49a
Show file tree

Hide file tree

Showing 7 changed files with 74 additions and 161 deletions.
diff --git a/.changeset/blue-geese-approve.md b/.changeset/blue-geese-approve.md
@@ -0,0 +1,5 @@
+---
+"@firebase/firestore": patch
+---
+
+Internal changes for the upcoming modular API.
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
@@ -128,63 +128,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 Query<T = legacy.DocumentData>
   extends Compat<exp.Query<T>>
   implements legacy.Query<T> {

diff --git a/packages/firestore/lite/src/api/reference.ts b/packages/firestore/lite/src/api/reference.ts
@@ -81,6 +81,7 @@ import { getDatastore } from './components';
 import { ByteString } from '../../../src/util/byte_string';
 import { Bytes } from './bytes';
 import { AbstractUserDataWriter } from '../../../src/api/user_data_writer';
+import { Compat } from '../../../src/compat/compat';
 
 /**
  * Document data (for use with {@link setDoc()}) consists of fields mapped to
@@ -1109,6 +1110,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
@@ -44,6 +44,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.
@@ -198,6 +199,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
@@ -28,8 +28,7 @@ import {
   firestoreClientGetNamedQuery,
   firestoreClientListen,
   firestoreClientLoadBundle,
-  firestoreClientTransaction,
-  firestoreClientWrite
+  firestoreClientTransaction
 } from '../core/firestore_client';
 import {
   Bound,
@@ -59,7 +58,6 @@ import { Transaction as InternalTransaction } from '../core/transaction';
 import { 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';
@@ -121,10 +119,12 @@ import {
   getDoc,
   onSnapshot,
   DocumentReference as ExpDocumentReference,
-  Query as ExpQuery
+  Query as ExpQuery,
+  executeWrite
 } from '../../exp/src/api/reference';
 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,
@@ -403,7 +403,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)
+      )
+    );
   }
 }
 
@@ -637,15 +641,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>,
@@ -654,38 +652,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>,
@@ -695,83 +677,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();
   }
 }
 
@@ -2016,6 +1947,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>,