add support for WriteBatch and Transaction

Author: Brian Chen

packages/firestore/src/api/database.ts

@@ -441,21 +441,24 @@ export class Transaction implements PublicTransaction, Compat<ExpTransaction> {
 
   set<T>(
     documentRef: DocumentReference<T>,
-    data: Partial<T>,
+    data: NestedPartialWithFieldValue<T>,
     options: PublicSetOptions
   ): Transaction;
-  set<T>(documentRef: DocumentReference<T>, data: T): Transaction;
+  set<T>(
+    documentRef: DocumentReference<T>,
+    data: WithFieldValue<T>
+  ): Transaction;
   set<T>(
     documentRef: PublicDocumentReference<T>,
-    data: T | Partial<T>,
+    data: WithFieldValue<T> | NestedPartialWithFieldValue<T>,
     options?: PublicSetOptions
   ): Transaction {
     const ref = castReference(documentRef);
     if (options) {
       validateSetOptions('Transaction.set', options);
-      this._delegate.set(ref, data, options);
+      this._delegate.set(ref, data as NestedPartialWithFieldValue<T>, options);
     } else {
-      this._delegate.set(ref, data);
+      this._delegate.set(ref, data as WithFieldValue<T>);
     }
     return this;
   }
@@ -502,21 +505,24 @@ export class WriteBatch implements PublicWriteBatch, Compat<ExpWriteBatch> {
   constructor(readonly _delegate: ExpWriteBatch) {}
   set<T>(
     documentRef: DocumentReference<T>,
-    data: Partial<T>,
+    data: NestedPartialWithFieldValue<T>,
     options: PublicSetOptions
   ): WriteBatch;
-  set<T>(documentRef: DocumentReference<T>, data: T): WriteBatch;
+  set<T>(
+    documentRef: DocumentReference<T>,
+    data: WithFieldValue<T>
+  ): WriteBatch;
   set<T>(
     documentRef: PublicDocumentReference<T>,
-    data: T | Partial<T>,
+    data: WithFieldValue<T> | NestedPartialWithFieldValue<T>,
     options?: PublicSetOptions
   ): WriteBatch {
     const ref = castReference(documentRef);
     if (options) {
       validateSetOptions('WriteBatch.set', options);
-      this._delegate.set(ref, data, options);
+      this._delegate.set(ref, data as NestedPartialWithFieldValue<T>, options);
     } else {
-      this._delegate.set(ref, data);
+      this._delegate.set(ref, data as WithFieldValue<T>);
     }
     return this;
   }

packages/firestore/src/lite/transaction.ts

@@ -27,7 +27,14 @@ import { Deferred } from '../util/promise';
 import { getDatastore } from './components';
 import { Firestore } from './database';
 import { FieldPath } from './field_path';
-import { DocumentReference, SetOptions, UpdateData } from './reference';
+import {
+  DocumentReference,
+  NestedPartialWithFieldValue,
+  SetOptions,
+  TypedUpdateData,
+  UpdateData,
+  WithFieldValue
+} from './reference';
 import {
   applyFirestoreDataConverter,
   LiteUserDataWriter
@@ -114,7 +121,7 @@ export class Transaction {
    * @param data - An object of the fields and values for the document.
    * @returns This `Transaction` instance. Used for chaining method calls.
    */
-  set<T>(documentRef: DocumentReference<T>, data: T): this;
+  set<T>(documentRef: DocumentReference<T>, data: WithFieldValue<T>): this;
   /**
    * Writes to the document referred to by the provided {@link
    * DocumentReference}. If the document does not exist yet, it will be created.
@@ -128,12 +135,12 @@ export class Transaction {
    */
   set<T>(
     documentRef: DocumentReference<T>,
-    data: Partial<T>,
+    data: NestedPartialWithFieldValue<T>,
     options: SetOptions
   ): this;
   set<T>(
     documentRef: DocumentReference<T>,
-    value: T,
+    value: WithFieldValue<T> | NestedPartialWithFieldValue<T>,
     options?: SetOptions
   ): this {
     const ref = validateReference(documentRef, this._firestore);
@@ -165,7 +172,7 @@ export class Transaction {
    * within the document.
    * @returns This `Transaction` instance. Used for chaining method calls.
    */
-  update(documentRef: DocumentReference<unknown>, data: UpdateData): this;
+  update<T>(documentRef: DocumentReference<T>, data: TypedUpdateData<T>): this;
   /**
    * Updates fields in the document referred to by the provided {@link
    * DocumentReference}. The update will fail if applied to a document that does
@@ -186,9 +193,9 @@ export class Transaction {
     value: unknown,
     ...moreFieldsAndValues: unknown[]
   ): this;
-  update(
-    documentRef: DocumentReference<unknown>,
-    fieldOrUpdateData: string | FieldPath | UpdateData,
+  update<T>(
+    documentRef: DocumentReference<T>,
+    fieldOrUpdateData: string | FieldPath | TypedUpdateData<T>,
     value?: unknown,
     ...moreFieldsAndValues: unknown[]
   ): this {

packages/firestore/src/lite/write_batch.ts

@@ -25,7 +25,14 @@ import { cast } from '../util/input_validation';
 import { getDatastore } from './components';
 import { Firestore } from './database';
 import { FieldPath } from './field_path';
-import { DocumentReference, SetOptions, UpdateData } from './reference';
+import {
+  DocumentReference,
+  NestedPartialWithFieldValue,
+  SetOptions,
+  TypedUpdateData,
+  UpdateData,
+  WithFieldValue
+} from './reference';
 import { applyFirestoreDataConverter } from './reference_impl';
 import {
   newUserDataReader,
@@ -67,7 +74,10 @@ export class WriteBatch {
    * @param data - An object of the fields and values for the document.
    * @returns This `WriteBatch` instance. Used for chaining method calls.
    */
-  set<T>(documentRef: DocumentReference<T>, data: T): WriteBatch;
+  set<T>(
+    documentRef: DocumentReference<T>,
+    data: WithFieldValue<T>
+  ): WriteBatch;
   /**
    * Writes to the document referred to by the provided {@link
    * DocumentReference}. If the document does not exist yet, it will be created.
@@ -81,12 +91,12 @@ export class WriteBatch {
    */
   set<T>(
     documentRef: DocumentReference<T>,
-    data: Partial<T>,
+    data: NestedPartialWithFieldValue<T>,
     options: SetOptions
   ): WriteBatch;
   set<T>(
     documentRef: DocumentReference<T>,
-    data: T,
+    data: WithFieldValue<T> | NestedPartialWithFieldValue<T>,
     options?: SetOptions
   ): WriteBatch {
     this._verifyNotCommitted();
@@ -120,7 +130,10 @@ export class WriteBatch {
    * within the document.
    * @returns This `WriteBatch` instance. Used for chaining method calls.
    */
-  update(documentRef: DocumentReference<unknown>, data: UpdateData): WriteBatch;
+  update<T>(
+    documentRef: DocumentReference<T>,
+    data: TypedUpdateData<T>
+  ): WriteBatch;
   /**
    * Updates fields in the document referred to by this {@link
    * DocumentReference}. The update will fail if applied to a document that does
@@ -141,9 +154,9 @@ export class WriteBatch {
     value: unknown,
     ...moreFieldsAndValues: unknown[]
   ): WriteBatch;
-  update(
-    documentRef: DocumentReference<unknown>,
-    fieldOrUpdateData: string | FieldPath | UpdateData,
+  update<T>(
+    documentRef: DocumentReference<T>,
+    fieldOrUpdateData: string | FieldPath | TypedUpdateData<T>,
     value?: unknown,
     ...moreFieldsAndValues: unknown[]
   ): WriteBatch {

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

@@ -382,7 +382,7 @@ describe('WriteBatch', () => {
 
     set<T>(
       ref: DocumentReference<T>,
-      data: T | Partial<T>,
+      data: WithFieldValue<T> | NestedPartialWithFieldValue<T>,
       options?: SetOptions
     ): Promise<void> {
       const batch = writeBatch(ref.firestore);
@@ -392,9 +392,9 @@ describe('WriteBatch', () => {
       return batch.commit();
     }
 
-    update(
-      ref: DocumentReference<unknown>,
-      dataOrField: UpdateData | string | FieldPath,
+    update<T>(
+      ref: DocumentReference<T>,
+      dataOrField: TypedUpdateData<T> | string | FieldPath,
       value?: unknown,
       ...moreFieldsAndValues: unknown[]
     ): Promise<void> {
@@ -445,21 +445,21 @@ describe('Transaction', () => {
 
     set<T>(
       ref: DocumentReference<T>,
-      data: T | Partial<T>,
+      data: WithFieldValue<T> | NestedPartialWithFieldValue<T>,
       options?: SetOptions
     ): Promise<void> {
       return runTransaction(ref.firestore, async transaction => {
         if (options) {
-          transaction.set(ref, data, options);
+          transaction.set(ref, data as NestedPartialWithFieldValue<T>, options);
         } else {
-          transaction.set(ref, data);
+          transaction.set(ref, data as WithFieldValue<T>);
         }
       });
     }
 
-    update(
-      ref: DocumentReference<unknown>,
-      dataOrField: UpdateData | string | FieldPath,
+    update<T>(
+      ref: DocumentReference<T>,
+      dataOrField: TypedUpdateData<T> | string | FieldPath,
       value?: unknown,
       ...moreFieldsAndValues: unknown[]
     ): Promise<void> {