add support for set() with SetOptions when using FirestoreDataConverter (#3254)

Brian Chen · web-flow · commit 39ca8ecf9404 · 2020-06-26T12:39:34.000-07:00
diff --git a/.changeset/quiet-coats-type.md b/.changeset/quiet-coats-type.md
@@ -0,0 +1,7 @@
+---
+'firebase': minor
+'@firebase/firestore': minor
+'@firebase/firestore-types': minor
+---
+
+Added support for `set()` with merge options when using `FirestoreDataConverter`.
diff --git a/packages/firebase/index.d.ts b/packages/firebase/index.d.ts
@@ -7859,9 +7859,11 @@ declare namespace firebase.firestore {
     /**
      * Called by the Firestore SDK to convert a custom model object of type T
      * into a plain Javascript object (suitable for writing directly to the
-     * Firestore database).
+     * Firestore database). To use `set()` with `merge` and `mergeFields`,
+     * `toFirestore()` must be defined with `Partial<T>`.
      */
     toFirestore(modelObject: T): DocumentData;
+    toFirestore(modelObject: Partial<T>, options: SetOptions): DocumentData;
 
     /**
      * Called by the Firestore SDK to convert Firestore data into an object of
@@ -8310,10 +8312,21 @@ declare namespace firebase.firestore {
      */
     set<T>(
       documentRef: DocumentReference<T>,
-      data: T,
-      options?: SetOptions
+      data: Partial<T>,
+      options: SetOptions
     ): Transaction;
 
+    /**
+     * Writes to the document referred to by the provided `DocumentReference`.
+     * If the document does not exist yet, it will be created. If you pass
+     * `SetOptions`, the provided data can be merged into the existing document.
+     *
+     * @param documentRef A reference to the document to be set.
+     * @param data An object of the fields and values for the document.
+     * @return This `Transaction` instance. Used for chaining method calls.
+     */
+    set<T>(documentRef: DocumentReference<T>, data: T): Transaction;
+
     /**
      * Updates fields in the document referred to by the provided
      * `DocumentReference`. The update will fail if applied to a document that
@@ -8384,10 +8397,21 @@ declare namespace firebase.firestore {
      */
     set<T>(
       documentRef: DocumentReference<T>,
-      data: T,
-      options?: SetOptions
+      data: Partial<T>,
+      options: SetOptions
     ): WriteBatch;
 
+    /**
+     * Writes to the document referred to by the provided `DocumentReference`.
+     * If the document does not exist yet, it will be created. If you pass
+     * `SetOptions`, the provided data can be merged into the existing document.
+     *
+     * @param documentRef A reference to the document to be set.
+     * @param data An object of the fields and values for the document.
+     * @return This `WriteBatch` instance. Used for chaining method calls.
+     */
+    set<T>(documentRef: DocumentReference<T>, data: T): WriteBatch;
+
     /**
      * Updates fields in the document referred to by the provided
      * `DocumentReference`. The update will fail if applied to a document that
@@ -8566,7 +8590,18 @@ declare namespace firebase.firestore {
      * @return A Promise resolved once the data has been successfully written
      * to the backend (Note that it won't resolve while you're offline).
      */
-    set(data: T, options?: SetOptions): Promise<void>;
+    set(data: Partial<T>, options: SetOptions): Promise<void>;
+
+    /**
+     * Writes to the document referred to by this `DocumentReference`. If the
+     * document does not yet exist, it will be created. If you pass
+     * `SetOptions`, the provided data can be merged into an existing document.
+     *
+     * @param data A map of the fields and values for the document.
+     * @return A Promise resolved once the data has been successfully written
+     * to the backend (Note that it won't resolve while you're offline).
+     */
+    set(data: T): Promise<void>;
 
     /**
      * Updates fields in the document referred to by this `DocumentReference`.
diff --git a/packages/firestore-types/index.d.ts b/packages/firestore-types/index.d.ts
@@ -50,6 +50,7 @@ export function setLogLevel(logLevel: LogLevel): void;
 
 export interface FirestoreDataConverter<T> {
   toFirestore(modelObject: T): DocumentData;
+  toFirestore(modelObject: Partial<T>, options: SetOptions): DocumentData;
 
   fromFirestore(snapshot: QueryDocumentSnapshot, options: SnapshotOptions): T;
 }
@@ -146,9 +147,10 @@ export class Transaction {
 
   set<T>(
     documentRef: DocumentReference<T>,
-    data: T,
-    options?: SetOptions
+    data: Partial<T>,
+    options: SetOptions
   ): Transaction;
+  set<T>(documentRef: DocumentReference<T>, data: T): Transaction;
 
   update(documentRef: DocumentReference<any>, data: UpdateData): Transaction;
   update(
@@ -166,9 +168,10 @@ export class WriteBatch {
 
   set<T>(
     documentRef: DocumentReference<T>,
-    data: T,
-    options?: SetOptions
+    data: Partial<T>,
+    options: SetOptions
   ): WriteBatch;
+  set<T>(documentRef: DocumentReference<T>, data: T): WriteBatch;
 
   update(documentRef: DocumentReference<any>, data: UpdateData): WriteBatch;
   update(
@@ -208,7 +211,8 @@ export class DocumentReference<T = DocumentData> {
 
   isEqual(other: DocumentReference<T>): boolean;
 
-  set(data: T, options?: SetOptions): Promise<void>;
+  set(data: Partial<T>, options: SetOptions): Promise<void>;
+  set(data: T): Promise<void>;
 
   update(data: UpdateData): Promise<void>;
   update(
diff --git a/packages/firestore/lite/src/api/reference.ts b/packages/firestore/lite/src/api/reference.ts
@@ -456,16 +456,13 @@ export function setDoc<T>(
 ): Promise<void> {
   const ref = cast(reference, DocumentReference);
 
-  const [convertedValue] = applyFirestoreDataConverter(
-    ref._converter,
-    data,
-    'setDoc'
-  );
+  const convertedValue = applyFirestoreDataConverter(ref._converter, data);
   const dataReader = newUserDataReader(ref.firestore);
   const parsed = dataReader.parseSetData(
     'setDoc',
     ref._key,
     convertedValue,
+    ref._converter !== null,
     options
   );
 
@@ -548,14 +545,16 @@ export function addDoc<T>(
   const collRef = cast(reference, CollectionReference);
   const docRef = doc(collRef);
 
-  const [convertedValue] = applyFirestoreDataConverter(
-    collRef._converter,
-    data,
-    'addDoc'
-  );
+  const convertedValue = applyFirestoreDataConverter(collRef._converter, data);
 
   const dataReader = newUserDataReader(collRef.firestore);
-  const parsed = dataReader.parseSetData('addDoc', docRef._key, convertedValue);
+  const parsed = dataReader.parseSetData(
+    'addDoc',
+    docRef._key,
+    convertedValue,
+    docRef._converter !== null,
+    {}
+  );
 
   return collRef.firestore
     ._getDatastore()
diff --git a/packages/firestore/lite/src/api/transaction.ts b/packages/firestore/lite/src/api/transaction.ts
@@ -95,15 +95,12 @@ export class Transaction implements firestore.Transaction {
     options?: firestore.SetOptions
   ): Transaction {
     const ref = validateReference(documentRef, this._firestore);
-    const [convertedValue] = applyFirestoreDataConverter(
-      ref._converter,
-      value,
-      'Transaction.set'
-    );
+    const convertedValue = applyFirestoreDataConverter(ref._converter, value);
     const parsed = this._dataReader.parseSetData(
       'Transaction.set',
       ref._key,
       convertedValue,
+      ref._converter !== null,
       options
     );
     this._transaction.set(ref._key, parsed);
diff --git a/packages/firestore/lite/src/api/write_batch.ts b/packages/firestore/lite/src/api/write_batch.ts
@@ -59,16 +59,13 @@ export class WriteBatch implements firestore.WriteBatch {
     this.verifyNotCommitted();
     const ref = validateReference(documentRef, this._firestore);
 
-    const [convertedValue] = applyFirestoreDataConverter(
-      ref._converter,
-      value,
-      'WriteBatch.set'
-    );
+    const convertedValue = applyFirestoreDataConverter(ref._converter, value);
 
     const parsed = this._dataReader.parseSetData(
       'WriteBatch.set',
       ref._key,
       convertedValue,
+      ref._converter !== null,
       options
     );
     this._mutations = this._mutations.concat(
diff --git a/packages/firestore/src/api/database.ts b/packages/firestore/src/api/database.ts
@@ -726,9 +726,15 @@ export class Transaction implements firestore.Transaction {
       });
   }
 
+  set<T>(
+    documentRef: DocumentReference<T>,
+    data: Partial<T>,
+    options: firestore.SetOptions
+  ): Transaction;
+  set<T>(documentRef: DocumentReference<T>, data: T): Transaction;
   set<T>(
     documentRef: firestore.DocumentReference<T>,
-    value: T,
+    value: T | Partial<T>,
     options?: firestore.SetOptions
   ): Transaction {
     validateBetweenNumberOfArgs('Transaction.set', arguments, 2, 3);
@@ -738,15 +744,16 @@ export class Transaction implements firestore.Transaction {
       this._firestore
     );
     options = validateSetOptions('Transaction.set', options);
-    const [convertedValue, functionName] = applyFirestoreDataConverter(
+    const convertedValue = applyFirestoreDataConverter(
       ref._converter,
       value,
-      'Transaction.set'
+      options
     );
     const parsed = this._firestore._dataReader.parseSetData(
-      functionName,
+      'Transaction.set',
       ref._key,
       convertedValue,
+      ref._converter !== null,
       options
     );
     this._transaction.set(ref._key, parsed);
@@ -825,9 +832,15 @@ export class WriteBatch implements firestore.WriteBatch {
 
   constructor(private _firestore: Firestore) {}
 
+  set<T>(
+    documentRef: DocumentReference<T>,
+    data: Partial<T>,
+    options: firestore.SetOptions
+  ): WriteBatch;
+  set<T>(documentRef: DocumentReference<T>, data: T): WriteBatch;
   set<T>(
     documentRef: firestore.DocumentReference<T>,
-    value: T,
+    value: T | Partial<T>,
     options?: firestore.SetOptions
   ): WriteBatch {
     validateBetweenNumberOfArgs('WriteBatch.set', arguments, 2, 3);
@@ -838,15 +851,16 @@ export class WriteBatch implements firestore.WriteBatch {
       this._firestore
     );
     options = validateSetOptions('WriteBatch.set', options);
-    const [convertedValue, functionName] = applyFirestoreDataConverter(
+    const convertedValue = applyFirestoreDataConverter(
       ref._converter,
       value,
-      'WriteBatch.set'
+      options
     );
     const parsed = this._firestore._dataReader.parseSetData(
-      functionName,
+      'WriteBatch.set',
       ref._key,
       convertedValue,
+      ref._converter !== null,
       options
     );
     this._mutations = this._mutations.concat(
@@ -1032,22 +1046,21 @@ export class DocumentReference<T = firestore.DocumentData>
     );
   }
 
-  set(
-    value: firestore.DocumentData,
-    options?: firestore.SetOptions
-  ): Promise<void>;
-  set(value: T, options?: firestore.SetOptions): Promise<void> {
+  set(value: Partial<T>, options: firestore.SetOptions): Promise<void>;
+  set(value: T): Promise<void>;
+  set(value: T | Partial<T>, options?: firestore.SetOptions): Promise<void> {
     validateBetweenNumberOfArgs('DocumentReference.set', arguments, 1, 2);
     options = validateSetOptions('DocumentReference.set', options);
-    const [convertedValue, functionName] = applyFirestoreDataConverter(
+    const convertedValue = applyFirestoreDataConverter(
       this._converter,
       value,
-      'DocumentReference.set'
+      options
     );
     const parsed = this.firestore._dataReader.parseSetData(
-      functionName,
+      'DocumentReference.set',
       this._key,
       convertedValue,
+      this._converter !== null,
       options
     );
     return this._firestoreClient.write(
@@ -2584,16 +2597,22 @@ function resultChangeType(type: ChangeType): firestore.DocumentChangeType {
 export function applyFirestoreDataConverter<T>(
   converter: UntypedFirestoreDataConverter<T> | null,
   value: T,
-  functionName: string
-): [firestore.DocumentData, string] {
+  options?: firestore.SetOptions
+): firestore.DocumentData {
   let convertedValue;
   if (converter) {
-    convertedValue = converter.toFirestore(value);
-    functionName = 'toFirestore() in ' + functionName;
+    if (options && (options.merge || options.mergeFields)) {
+      // Cast to `any` in order to satisfy the union type constraint on
+      // toFirestore().
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      convertedValue = (converter as any).toFirestore(value, options);
+    } else {
+      convertedValue = converter.toFirestore(value);
+    }
   } else {
     convertedValue = value as firestore.DocumentData;
   }
-  return [convertedValue, functionName];
+  return convertedValue;
 }
 
 function contains(obj: object, key: string): obj is { key: unknown } {
diff --git a/packages/firestore/src/api/user_data_reader.ts b/packages/firestore/src/api/user_data_reader.ts
diff --git a/packages/firestore/test/integration/api/database.test.ts b/packages/firestore/test/integration/api/database.test.ts
