Add WriteBatch

packages/firestore/lite/index.d.ts

@@ -140,10 +140,11 @@ export class Transaction {
 
   get<T>(documentRef: DocumentReference<T>): Promise<DocumentSnapshot<T>>;
 
+  set<T>(documentRef: DocumentReference<T>, data: T): Transaction;
   set<T>(
     documentRef: DocumentReference<T>,
-    data: T,
-    options?: SetOptions
+    data: Partial<T>,
+    options: SetOptions
   ): Transaction;
 
   update(documentRef: DocumentReference<any>, data: UpdateData): Transaction;
@@ -160,10 +161,11 @@ export class Transaction {
 export class WriteBatch {
   private constructor();
 
+  set<T>(documentRef: DocumentReference<T>, data: T): WriteBatch;
   set<T>(
     documentRef: DocumentReference<T>,
-    data: T,
-    options?: SetOptions
+    data: Partial<T>,
+    options: SetOptions
   ): WriteBatch;
 
   update(documentRef: DocumentReference<any>, data: UpdateData): WriteBatch;

packages/firestore/lite/index.node.ts

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

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

@@ -318,17 +318,17 @@ export function setDoc<T>(
 }
 
 export function updateDoc(
-  reference: firestore.DocumentReference,
+  reference: firestore.DocumentReference<unknown>,
   data: firestore.UpdateData
 ): Promise<void>;
 export function updateDoc(
-  reference: firestore.DocumentReference,
+  reference: firestore.DocumentReference<unknown>,
   field: string | firestore.FieldPath,
   value: unknown,
   ...moreFieldsAndValues: unknown[]
 ): Promise<void>;
 export function updateDoc(
-  reference: firestore.DocumentReference,
+  reference: firestore.DocumentReference<unknown>,
   fieldOrUpdateData: string | firestore.FieldPath | firestore.UpdateData,
   value?: unknown,
   ...moreFieldsAndValues: unknown[]
@@ -418,7 +418,7 @@ export function addDoc<T>(
     .then(() => docRef);
 }
 
-function newUserDataReader(
+export function newUserDataReader(
   databaseId: DatabaseId,
   settings: firestore.Settings
 ): UserDataReader {

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

@@ -0,0 +1,179 @@
+/**
+ * @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 {
+  DeleteMutation,
+  Mutation,
+  Precondition
+} from '../../../src/model/mutation';
+import { Code, FirestoreError } from '../../../src/util/error';
+import { applyFirestoreDataConverter } from '../../../src/api/database';
+import {
+  DocumentKeyReference,
+  UserDataReader
+} from '../../../src/api/user_data_reader';
+import { cast } from './util';
+import { DocumentReference, newUserDataReader } from './reference';
+import { Firestore } from './database';
+import { invokeCommitRpc } from '../../../src/remote/datastore';
+import { FieldPath } from './field_path';
+
+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 _mutations = [] as Mutation[];
+  private _committed = false;
+  private _dataReader: UserDataReader;
+
+  constructor(private _firestore: Firestore) {
+    // 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!
+    );
+  }
+
+  set<T>(documentRef: firestore.DocumentReference<T>, value: T): WriteBatch;
+  set<T>(
+    documentRef: firestore.DocumentReference<T>,
+    value: Partial<T>,
+    options: firestore.SetOptions
+  ): WriteBatch;
+  set<T>(
+    documentRef: firestore.DocumentReference<T>,
+    value: T,
+    options?: firestore.SetOptions
+  ): WriteBatch {
+    this.verifyNotCommitted();
+    const ref = validateReference(documentRef, this._firestore);
+
+    const [convertedValue] = applyFirestoreDataConverter(
+      ref._converter,
+      value,
+      'WriteBatch.set'
+    );
+
+    const parsed = this._dataReader.parseSetData(
+      'WriteBatch.set',
+      convertedValue,
+      options
+    );
+    this._mutations = this._mutations.concat(
+      parsed.toMutations(ref._key, Precondition.none())
+    );
+    return this;
+  }
+
+  update(
+    documentRef: firestore.DocumentReference<unknown>,
+    value: firestore.UpdateData
+  ): WriteBatch;
+  update(
+    documentRef: firestore.DocumentReference<unknown>,
+    field: string | firestore.FieldPath,
+    value: unknown,
+    ...moreFieldsAndValues: unknown[]
+  ): WriteBatch;
+  update(
+    documentRef: firestore.DocumentReference<unknown>,
+    fieldOrUpdateData: string | firestore.FieldPath | firestore.UpdateData,
+    value?: unknown,
+    ...moreFieldsAndValues: unknown[]
+  ): WriteBatch {
+    this.verifyNotCommitted();
+    const ref = validateReference(documentRef, this._firestore);
+
+    let parsed;
+
+    if (
+      typeof fieldOrUpdateData === 'string' ||
+      fieldOrUpdateData instanceof FieldPath
+    ) {
+      parsed = this._dataReader.parseUpdateVarargs(
+        'WriteBatch.update',
+        fieldOrUpdateData,
+        value,
+        moreFieldsAndValues
+      );
+    } else {
+      parsed = this._dataReader.parseUpdateData(
+        'WriteBatch.update',
+        fieldOrUpdateData
+      );
+    }
+
+    this._mutations = this._mutations.concat(
+      parsed.toMutations(ref._key, Precondition.exists(true))
+    );
+    return this;
+  }
+
+  delete(documentRef: firestore.DocumentReference<unknown>): WriteBatch {
+    this.verifyNotCommitted();
+    const ref = validateReference(documentRef, this._firestore);
+    this._mutations = this._mutations.concat(
+      new DeleteMutation(ref._key, Precondition.none())
+    );
+    return this;
+  }
+
+  commit(): Promise<void> {
+    this.verifyNotCommitted();
+    this._committed = true;
+    if (this._mutations.length > 0) {
+      return this._firestore
+        ._ensureClientConfigured()
+        .then(datastore => invokeCommitRpc(datastore, 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.'
+      );
+    }
+  }
+}
+
+export function validateReference<T>(
+  documentRef: firestore.DocumentReference<T>,
+  firestore: Firestore
+): DocumentKeyReference<T> {
+  if (documentRef.firestore !== firestore) {
+    throw new FirestoreError(
+      Code.INVALID_ARGUMENT,
+      'Provided document reference is from a different Firestore instance.'
+    );
+  } else {
+    return cast(documentRef, DocumentReference) as DocumentReference<T>;
+  }
+}
+
+export function writeBatch(
+  firestore: firestore.FirebaseFirestore
+): firestore.WriteBatch {
+  return new WriteBatch(cast(firestore, Firestore));
+}