feat!: add QoL improvements for converters (#1595)

googleapis · Aug 24, 2021 · 39e5138 · 39e5138
1 parent b6c5a1a
commit 39e5138
Show file tree

Hide file tree

Showing 9 changed files with 793 additions and 73 deletions.
diff --git a/dev/src/bulk-writer.ts b/dev/src/bulk-writer.ts
@@ -562,7 +562,7 @@ export class BulkWriter {
    */
   create<T>(
     documentRef: firestore.DocumentReference<T>,
-    data: T
+    data: firestore.WithFieldValue<T>
   ): Promise<WriteResult> {
     this._verifyNotClosed();
     return this._enqueue(documentRef, 'create', bulkCommitBatch =>
@@ -654,13 +654,20 @@ export class BulkWriter {
    */
   set<T>(
     documentRef: firestore.DocumentReference<T>,
-    data: T | Partial<T>,
+    data: firestore.PartialWithFieldValue<T>,
     options?: firestore.SetOptions
   ): Promise<WriteResult> {
     this._verifyNotClosed();
-    return this._enqueue(documentRef, 'set', bulkCommitBatch =>
-      bulkCommitBatch.set(documentRef, data, options)
-    );
+    return this._enqueue(documentRef, 'set', bulkCommitBatch => {
+      if (options) {
+        return bulkCommitBatch.set(documentRef, data, options);
+      } else {
+        return bulkCommitBatch.set(
+          documentRef,
+          data as firestore.WithFieldValue<T>
+        );
+      }
+    });
   }
 
   /**
@@ -706,7 +713,7 @@ export class BulkWriter {
    */
   update<T>(
     documentRef: firestore.DocumentReference<T>,
-    dataOrField: firestore.UpdateData | string | FieldPath,
+    dataOrField: firestore.UpdateData<T> | string | FieldPath,
     ...preconditionOrValues: Array<
       {lastUpdateTime?: Timestamp} | unknown | string | FieldPath
     >

diff --git a/dev/src/document-change.ts b/dev/src/document-change.ts
@@ -27,7 +27,7 @@ export type DocumentChangeType = 'added' | 'removed' | 'modified';
  * @class DocumentChange
  */
 export class DocumentChange<T = firestore.DocumentData>
-  implements firestore.DocumentChange
+  implements firestore.DocumentChange<T>
 {
   private readonly _type: DocumentChangeType;
   private readonly _document: QueryDocumentSnapshot<T>;

diff --git a/dev/src/reference.ts b/dev/src/reference.ts
@@ -132,6 +132,7 @@ export class DocumentReference<T = firestore.DocumentData>
    *
    * @param _firestore The Firestore Database client.
    * @param _path The Path of this reference.
+   * @param _converter The converter to use when serializing data.
    */
   constructor(
     private readonly _firestore: Firestore,
@@ -356,7 +357,7 @@ export class DocumentReference<T = firestore.DocumentData>
    *   console.log(`Failed to create document: ${err}`);
    * });
    */
-  create(data: T): Promise<WriteResult> {
+  create(data: firestore.WithFieldValue<T>): Promise<WriteResult> {
     const writeBatch = new WriteBatch(this._firestore);
     return writeBatch
       .create(this, data)
@@ -395,8 +396,11 @@ export class DocumentReference<T = firestore.DocumentData>
       .then(([writeResult]) => writeResult);
   }
 
-  set(data: Partial<T>, options: firestore.SetOptions): Promise<WriteResult>;
-  set(data: T): Promise<WriteResult>;
+  set(
+    data: firestore.PartialWithFieldValue<T>,
+    options: firestore.SetOptions
+  ): Promise<WriteResult>;
+  set(data: firestore.WithFieldValue<T>): Promise<WriteResult>;
   /**
    * Writes to the document referred to by this DocumentReference. If the
    * document does not yet exist, it will be created. If you pass
@@ -421,14 +425,16 @@ export class DocumentReference<T = firestore.DocumentData>
    * });
    */
   set(
-    data: T | Partial<T>,
+    data: firestore.PartialWithFieldValue<T>,
     options?: firestore.SetOptions
   ): Promise<WriteResult> {
-    const writeBatch = new WriteBatch(this._firestore);
-    return writeBatch
-      .set(this, data, options)
-      .commit()
-      .then(([writeResult]) => writeResult);
+    let writeBatch = new WriteBatch(this._firestore);
+    if (options) {
+      writeBatch = writeBatch.set(this, data, options);
+    } else {
+      writeBatch = writeBatch.set(this, data as firestore.WithFieldValue<T>);
+    }
+    return writeBatch.commit().then(([writeResult]) => writeResult);
   }
 
   /**
@@ -461,7 +467,7 @@ export class DocumentReference<T = firestore.DocumentData>
    * });
    */
   update(
-    dataOrField: firestore.UpdateData | string | firestore.FieldPath,
+    dataOrField: firestore.UpdateData<T> | string | firestore.FieldPath,
     ...preconditionOrValues: Array<
       unknown | string | firestore.FieldPath | firestore.Precondition
     >
@@ -2609,7 +2615,7 @@ export class CollectionReference<T = firestore.DocumentData>
    *   console.log(`Added document with name: ${documentReference.id}`);
    * });
    */
-  add(data: T): Promise<DocumentReference<T>> {
+  add(data: firestore.WithFieldValue<T>): Promise<DocumentReference<T>> {
     const firestoreData = this._queryOptions.converter.toFirestore(data);
     validateDocumentData(
       'data',

diff --git a/dev/src/transaction.ts b/dev/src/transaction.ts
@@ -212,17 +212,23 @@ export class Transaction implements firestore.Transaction {
    *   });
    * });
    */
-  create<T>(documentRef: firestore.DocumentReference<T>, data: T): Transaction {
+  create<T>(
+    documentRef: firestore.DocumentReference<T>,
+    data: firestore.WithFieldValue<T>
+  ): Transaction {
     this._writeBatch.create(documentRef, data);
     return this;
   }
 
   set<T>(
     documentRef: firestore.DocumentReference<T>,
-    data: Partial<T>,
+    data: firestore.PartialWithFieldValue<T>,
     options: firestore.SetOptions
   ): Transaction;
-  set<T>(documentRef: firestore.DocumentReference<T>, data: T): Transaction;
+  set<T>(
+    documentRef: firestore.DocumentReference<T>,
+    data: firestore.WithFieldValue<T>
+  ): Transaction;
   /**
    * Writes to the document referred to by the provided
    * [DocumentReference]{@link DocumentReference}. If the document
@@ -252,10 +258,14 @@ export class Transaction implements firestore.Transaction {
    */
   set<T>(
     documentRef: firestore.DocumentReference<T>,
-    data: T | Partial<T>,
+    data: firestore.PartialWithFieldValue<T>,
     options?: firestore.SetOptions
   ): Transaction {
-    this._writeBatch.set(documentRef, data, options);
+    if (options) {
+      this._writeBatch.set(documentRef, data, options);
+    } else {
+      this._writeBatch.set(documentRef, data as firestore.WithFieldValue<T>);
+    }
     return this;
   }
 
@@ -299,7 +309,7 @@ export class Transaction implements firestore.Transaction {
    */
   update<T>(
     documentRef: firestore.DocumentReference<T>,
-    dataOrField: firestore.UpdateData | string | firestore.FieldPath,
+    dataOrField: firestore.UpdateData<T> | string | firestore.FieldPath,
     ...preconditionOrValues: Array<
       firestore.Precondition | unknown | string | firestore.FieldPath
     >

diff --git a/dev/src/types.ts b/dev/src/types.ts
@@ -14,7 +14,11 @@
  * limitations under the License.
  */
 
-import {FirestoreDataConverter, DocumentData} from '@google-cloud/firestore';
+import {
+  FirestoreDataConverter,
+  DocumentData,
+  WithFieldValue,
+} from '@google-cloud/firestore';
 
 import {CallOptions} from 'google-gax';
 import {Duplex} from 'stream';
@@ -114,7 +118,7 @@ export type RBTree = any;
  * @internal
  */
 const defaultConverterObj: FirestoreDataConverter<DocumentData> = {
-  toFirestore(modelObject: DocumentData): DocumentData {
+  toFirestore(modelObject: WithFieldValue<DocumentData>): DocumentData {
     return modelObject;
   },
   fromFirestore(snapshot: QueryDocumentSnapshot): DocumentData {

diff --git a/dev/src/write-batch.ts b/dev/src/write-batch.ts
@@ -185,9 +185,14 @@ export class WriteBatch implements firestore.WriteBatch {
    *   console.log('Successfully executed batch.');
    * });
    */
-  create<T>(documentRef: firestore.DocumentReference<T>, data: T): WriteBatch {
+  create<T>(
+    documentRef: firestore.DocumentReference<T>,
+    data: firestore.WithFieldValue<T>
+  ): WriteBatch {
     const ref = validateDocumentReference('documentRef', documentRef);
-    const firestoreData = ref._converter.toFirestore(data);
+    const firestoreData = ref._converter.toFirestore(
+      data as firestore.WithFieldValue<T>
+    );
     validateDocumentData(
       'data',
       firestoreData,
@@ -268,14 +273,12 @@ export class WriteBatch implements firestore.WriteBatch {
 
   set<T>(
     documentRef: firestore.DocumentReference<T>,
-    data: Partial<T>,
+    data: firestore.PartialWithFieldValue<T>,
     options: firestore.SetOptions
   ): WriteBatch;
-  set<T>(documentRef: firestore.DocumentReference<T>, data: T): WriteBatch;
   set<T>(
     documentRef: firestore.DocumentReference<T>,
-    data: T | Partial<T>,
-    options?: firestore.SetOptions
+    data: firestore.WithFieldValue<T>
   ): WriteBatch;
   /**
    * Write to the document referred to by the provided
@@ -308,7 +311,7 @@ export class WriteBatch implements firestore.WriteBatch {
    */
   set<T>(
     documentRef: firestore.DocumentReference<T>,
-    data: T | Partial<T>,
+    data: firestore.PartialWithFieldValue<T>,
     options?: firestore.SetOptions
   ): WriteBatch {
     validateSetOptions('options', options, {optional: true});
@@ -405,7 +408,7 @@ export class WriteBatch implements firestore.WriteBatch {
    */
   update<T = firestore.DocumentData>(
     documentRef: firestore.DocumentReference<T>,
-    dataOrField: firestore.UpdateData | string | firestore.FieldPath,
+    dataOrField: firestore.UpdateData<T> | string | firestore.FieldPath,
     ...preconditionOrValues: Array<
       | {lastUpdateTime?: firestore.Timestamp}
       | unknown
@@ -470,15 +473,16 @@ export class WriteBatch implements firestore.WriteBatch {
         // eslint-disable-next-line prefer-rest-params
         validateMaxNumberOfArguments('update', arguments, 3);
 
-        const data = dataOrField as firestore.UpdateData;
-        Object.entries(data).forEach(([key, value]) => {
-          // Skip `undefined` values (can be hit if `ignoreUndefinedProperties`
-          // is set)
-          if (value !== undefined) {
-            validateFieldPath(key, key);
-            updateMap.set(FieldPath.fromArgument(key), value);
+        Object.entries(dataOrField as firestore.UpdateData<T>).forEach(
+          ([key, value]) => {
+            // Skip `undefined` values (can be hit if `ignoreUndefinedProperties`
+            // is set)
+            if (value !== undefined) {
+              validateFieldPath(key, key);
+              updateMap.set(FieldPath.fromArgument(key), value);
+            }
           }
-        });
+        );
 
         if (preconditionOrValues.length > 0) {
           validateUpdatePrecondition(