WIP: working NestedParial and WithFieldValue

Brian Chen · Brian Chen · commit f5a00fcd4cde · 2021-08-10T13:55:43.000-05:00
diff --git a/packages/firestore/src/api/database.ts b/packages/firestore/src/api/database.ts
@@ -127,6 +127,7 @@ import {
   NextFn,
   PartialObserver
 } from './observer';
+import { NestedPartialWithFieldValue, WithFieldValue } from '../lite/reference';
 
 /**
  * A persistence provider for either memory-only or IndexedDB persistence.
@@ -597,19 +598,19 @@ class FirestoreDataConverter<U>
     );
   }
 
-  toFirestore(modelObject: U): PublicDocumentData;
+  toFirestore(modelObject: WithFieldValue<U>): PublicDocumentData;
   toFirestore(
-    modelObject: Partial<U>,
+    modelObject: NestedPartialWithFieldValue<U>,
     options: PublicSetOptions
   ): PublicDocumentData;
   toFirestore(
-    modelObject: U | Partial<U>,
+    modelObject: WithFieldValue<U> | NestedPartialWithFieldValue<U>,
     options?: PublicSetOptions
   ): PublicDocumentData {
     if (!options) {
       return this._delegate.toFirestore(modelObject as U);
     } else {
-      return this._delegate.toFirestore(modelObject, options);
+      return this._delegate.toFirestore(modelObject as Partial<U>, options);
     }
   }
 
@@ -733,7 +734,15 @@ export class DocumentReference<T = PublicDocumentData>
   set(value: T | Partial<T>, options?: PublicSetOptions): Promise<void> {
     options = validateSetOptions('DocumentReference.set', options);
     try {
-      return setDoc(this._delegate, value, options);
+      if (options) {
+        return setDoc(
+          this._delegate,
+          value as NestedPartialWithFieldValue<T>,
+          options
+        );
+      } else {
+        return setDoc(this._delegate, value as WithFieldValue<T>);
+      }
     } catch (e) {
       throw replaceFunctionName(e, 'setDoc()', 'DocumentReference.set()');
     }
diff --git a/packages/firestore/src/exp/reference_impl.ts b/packages/firestore/src/exp/reference_impl.ts
@@ -42,9 +42,11 @@ import {
   CollectionReference,
   doc,
   DocumentReference,
+  NestedPartialWithFieldValue,
   Query,
   SetOptions,
-  UpdateData
+  UpdateData,
+  WithFieldValue
 } from '../lite/reference';
 import { applyFirestoreDataConverter } from '../lite/reference_impl';
 import {
@@ -243,7 +245,7 @@ export function getDocsFromServer<T>(
  */
 export function setDoc<T>(
   reference: DocumentReference<T>,
-  data: T
+  data: WithFieldValue<T>
 ): Promise<void>;
 /**
  * Writes to the document referred to by the specified `DocumentReference`. If
@@ -258,12 +260,12 @@ export function setDoc<T>(
  */
 export function setDoc<T>(
   reference: DocumentReference<T>,
-  data: Partial<T>,
+  data: NestedPartialWithFieldValue<T>,
   options: SetOptions
 ): Promise<void>;
 export function setDoc<T>(
   reference: DocumentReference<T>,
-  data: T,
+  data: WithFieldValue<T> | NestedPartialWithFieldValue<T>,
   options?: SetOptions
 ): Promise<void> {
   reference = cast<DocumentReference<T>>(reference, DocumentReference);
@@ -393,7 +395,7 @@ export function deleteDoc(
  */
 export function addDoc<T>(
   reference: CollectionReference<T>,
-  data: T
+  data: WithFieldValue<T>
 ): Promise<DocumentReference<T>> {
   const firestore = cast(reference.firestore, Firestore);
 
diff --git a/packages/firestore/src/exp/snapshot.ts b/packages/firestore/src/exp/snapshot.ts
@@ -18,7 +18,14 @@
 import { newQueryComparator } from '../core/query';
 import { ChangeType, ViewSnapshot } from '../core/view_snapshot';
 import { FieldPath } from '../lite/field_path';
-import { DocumentData, Query, queryEqual, SetOptions } from '../lite/reference';
+import {
+  DocumentData,
+  NestedPartialWithFieldValue,
+  Query,
+  queryEqual,
+  SetOptions,
+  WithFieldValue
+} from '../lite/reference';
 import {
   DocumentSnapshot as LiteDocumentSnapshot,
   fieldPathFromArgument,
@@ -84,15 +91,18 @@ export interface FirestoreDataConverter<T>
    * Firestore database). To use `set()` with `merge` and `mergeFields`,
    * `toFirestore()` must be defined with `Partial<T>`.
    */
-  toFirestore(modelObject: T): DocumentData;
+  toFirestore(modelObject: WithFieldValue<T>): DocumentData;
 
   /**
    * 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). Used with {@link (setDoc:1)}, {@link (WriteBatch.set:1)}
    * and {@link (Transaction.set:1)} with `merge:true` or `mergeFields`.
    */
-  toFirestore(modelObject: Partial<T>, options: SetOptions): DocumentData;
+  toFirestore(
+    modelObject: NestedPartialWithFieldValue<T>,
+    options: SetOptions
+  ): DocumentData;
 
   /**
    * Called by the Firestore SDK to convert Firestore data into an object of
diff --git a/packages/firestore/src/lite/reference.ts b/packages/firestore/src/lite/reference.ts
@@ -36,6 +36,7 @@ import { AutoId } from '../util/misc';
 
 import { Firestore } from './database';
 import { FieldPath } from './field_path';
+import { FieldValue } from './field_value';
 import { FirestoreDataConverter } from './snapshot';
 
 /**
@@ -48,6 +49,21 @@ export interface DocumentData {
   [field: string]: any;
 }
 
+type Primitive = string | number | boolean | bigint | undefined | null;
+// eslint-disable-next-line @typescript-eslint/ban-types
+type Builtin = Primitive | Function;
+
+/** Like Partial but recursive */
+export type NestedPartialWithFieldValue<T> = T extends Builtin
+  ? T
+  : T extends Map<infer K, infer V>
+  ? Map<NestedPartialWithFieldValue<K>, NestedPartialWithFieldValue<V>>
+  : T extends {}
+  ? { [K in keyof T]?: NestedPartialWithFieldValue<T[K]> | FieldValue }
+  : Partial<T>;
+
+export type WithFieldValue<T> = { [P in keyof T]: T[P] | FieldValue };
+
 /**
  * Update data (for use with {@link @firebase/firestore/lite#(updateDoc:1)}) consists of field paths (e.g.
  * 'foo' or 'foo.baz') mapped to values. Fields that contain dots reference
diff --git a/packages/firestore/src/lite/reference_impl.ts b/packages/firestore/src/lite/reference_impl.ts
@@ -41,9 +41,11 @@ import {
   CollectionReference,
   doc,
   DocumentReference,
+  NestedPartialWithFieldValue,
   Query,
   SetOptions,
-  UpdateData
+  UpdateData,
+  WithFieldValue
 } from './reference';
 import {
   DocumentSnapshot,
@@ -197,7 +199,7 @@ export function getDocs<T>(query: Query<T>): Promise<QuerySnapshot<T>> {
  */
 export function setDoc<T>(
   reference: DocumentReference<T>,
-  data: T
+  data: WithFieldValue<T>
 ): Promise<void>;
 /**
  * Writes to the document referred to by the specified `DocumentReference`. If
@@ -217,12 +219,12 @@ export function setDoc<T>(
  */
 export function setDoc<T>(
   reference: DocumentReference<T>,
-  data: Partial<T>,
+  data: NestedPartialWithFieldValue<T>,
   options: SetOptions
 ): Promise<void>;
 export function setDoc<T>(
   reference: DocumentReference<T>,
-  data: T,
+  data: WithFieldValue<T> | NestedPartialWithFieldValue<T>,
   options?: SetOptions
 ): Promise<void> {
   reference = cast<DocumentReference<T>>(reference, DocumentReference);
diff --git a/packages/firestore/src/lite/snapshot.ts b/packages/firestore/src/lite/snapshot.ts
@@ -27,9 +27,11 @@ import { FieldPath } from './field_path';
 import {
   DocumentData,
   DocumentReference,
+  NestedPartialWithFieldValue,
   Query,
   queryEqual,
-  SetOptions
+  SetOptions,
+  WithFieldValue
 } from './reference';
 import {
   fieldPathFromDotSeparatedString,
@@ -83,15 +85,18 @@ export interface FirestoreDataConverter<T> {
    * Firestore database). Used with {@link @firebase/firestore/lite#(setDoc:1)}, {@link @firebase/firestore/lite#(WriteBatch.set:1)}
    * and {@link @firebase/firestore/lite#(Transaction.set:1)}.
    */
-  toFirestore(modelObject: T): DocumentData;
+  toFirestore(modelObject: WithFieldValue<T>): DocumentData;
 
   /**
    * 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). Used with {@link @firebase/firestore/lite#(setDoc:1)}, {@link @firebase/firestore/lite#(WriteBatch.set:1)}
    * and {@link @firebase/firestore/lite#(Transaction.set:1)} with `merge:true` or `mergeFields`.
    */
-  toFirestore(modelObject: Partial<T>, options: SetOptions): DocumentData;
+  toFirestore(
+    modelObject: NestedPartialWithFieldValue<T>,
+    options: SetOptions
+  ): DocumentData;
 
   /**
    * Called by the Firestore SDK to convert Firestore data into an object of
diff --git a/packages/firestore/src/lite/user_data_reader.ts b/packages/firestore/src/lite/user_data_reader.ts
@@ -63,7 +63,11 @@ import { Firestore } from './database';
 import { FieldPath } from './field_path';
 import { FieldValue } from './field_value';
 import { GeoPoint } from './geo_point';
-import { DocumentReference } from './reference';
+import {
+  DocumentReference,
+  NestedPartialWithFieldValue,
+  WithFieldValue
+} from './reference';
 import { Timestamp } from './timestamp';
 
 const RESERVED_FIELD_REGEX = /^__.*__$/;
@@ -73,8 +77,11 @@ const RESERVED_FIELD_REGEX = /^__.*__$/;
  * lite, firestore-exp and classic SDK.
  */
 export interface UntypedFirestoreDataConverter<T> {
-  toFirestore(modelObject: T): DocumentData;
-  toFirestore(modelObject: Partial<T>, options: SetOptions): DocumentData;
+  toFirestore(modelObject: WithFieldValue<T>): DocumentData;
+  toFirestore(
+    modelObject: NestedPartialWithFieldValue<T>,
+    options: SetOptions
+  ): DocumentData;
   fromFirestore(snapshot: unknown, options?: unknown): T;
 }
 
diff --git a/packages/firestore/test/integration/api/database.test.ts b/packages/firestore/test/integration/api/database.test.ts
@@ -32,6 +32,7 @@ import {
   withTestDocAndInitialData
 } from '../util/helpers';
 import { DEFAULT_SETTINGS, DEFAULT_PROJECT_ID } from '../util/settings';
+import { NestedPartialWithFieldValue } from '../../../src/lite/reference';
 
 use(chaiAsPromised);
 
@@ -1351,7 +1352,7 @@ apiDescribe('Database', (persistence: boolean) => {
 
     const postConverterMerge = {
       toFirestore(
-        post: Partial<Post>,
+        post: NestedPartialWithFieldValue<Post>,
         options?: firestore.SetOptions
       ): firestore.DocumentData {
         if (options && (options.merge || options.mergeFields)) {
diff --git a/packages/firestore/test/lite/helpers.ts b/packages/firestore/test/lite/helpers.ts
@@ -26,7 +26,8 @@ import {
   DocumentData,
   CollectionReference,
   DocumentReference,
-  SetOptions
+  SetOptions,
+  NestedPartialWithFieldValue
 } from '../../src/lite/reference';
 import { setDoc } from '../../src/lite/reference_impl';
 import { FirestoreSettings } from '../../src/lite/settings';
@@ -102,7 +103,11 @@ export function withTestCollection(
 
 // Used for testing the FirestoreDataConverter.
 export class Post {
-  constructor(readonly title: string, readonly author: string) {}
+  constructor(
+    readonly title: string,
+    readonly author: string,
+    readonly id = 1
+  ) {}
   byline(): string {
     return this.title + ', by ' + this.author;
   }
@@ -119,7 +124,10 @@ export const postConverter = {
 };
 
 export const postConverterMerge = {
-  toFirestore(post: Partial<Post>, options?: SetOptions): DocumentData {
+  toFirestore(
+    post: NestedPartialWithFieldValue<Post>,
+    options?: SetOptions
+  ): DocumentData {
     if (
       options &&
       ((options as { merge: true }).merge ||
diff --git a/packages/firestore/test/lite/integration.test.ts b/packages/firestore/test/lite/integration.test.ts
@@ -57,7 +57,9 @@ import {
   collectionGroup,
   SetOptions,
   UpdateData,
-  DocumentData
+  DocumentData,
+  WithFieldValue,
+  NestedPartialWithFieldValue
 } from '../../src/lite/reference';
 import {
   addDoc,
@@ -337,10 +339,13 @@ describe('getDoc()', () => {
  * DocumentReference-based mutation API.
  */
 interface MutationTester {
-  set<T>(documentRef: DocumentReference<T>, data: T): Promise<void>;
   set<T>(
     documentRef: DocumentReference<T>,
-    data: Partial<T>,
+    data: WithFieldValue<T>
+  ): Promise<void>;
+  set<T>(
+    documentRef: DocumentReference<T>,
+    data: NestedPartialWithFieldValue<T>,
     options: SetOptions
   ): Promise<void>;
   update(
@@ -580,7 +585,11 @@ function genericMutationTests(
         const coll = collection(db, 'posts');
         const ref = doc(coll, 'post').withConverter(postConverterMerge);
         await setDoc(ref, new Post('walnut', 'author'));
-        await setDoc(ref, { title: 'olive' }, { merge: true });
+        await setDoc(
+          ref,
+          { title: 'olive', id: increment(2) },
+          { merge: true }
+        );
         const postDoc = await getDoc(ref);
         expect(postDoc.get('title')).to.equal('olive');
         expect(postDoc.get('author')).to.equal('author');
