add support for TypedUpdateData + sanity tests

Author: Brian Chen

packages/firestore/src/exp/reference_impl.ts

@@ -45,6 +45,7 @@ import {
   NestedPartialWithFieldValue,
   Query,
   SetOptions,
+  TypedUpdateData,
   UpdateData,
   WithFieldValue
 } from '../lite/reference';
@@ -302,9 +303,9 @@ export function setDoc<T>(
  * @returns A Promise resolved once the data has been successfully written
  * to the backend (note that it won't resolve while you're offline).
  */
-export function updateDoc(
-  reference: DocumentReference<unknown>,
-  data: UpdateData
+export function updateDoc<T>(
+  reference: DocumentReference<T>,
+  data: TypedUpdateData<T>
 ): Promise<void>;
 /**
  * Updates fields in the document referred to by the specified
@@ -327,9 +328,9 @@ export function updateDoc(
   value: unknown,
   ...moreFieldsAndValues: unknown[]
 ): Promise<void>;
-export function updateDoc(
+export function updateDoc<T>(
   reference: DocumentReference<unknown>,
-  fieldOrUpdateData: string | FieldPath | UpdateData,
+  fieldOrUpdateData: string | FieldPath | TypedUpdateData<T>,
   value?: unknown,
   ...moreFieldsAndValues: unknown[]
 ): Promise<void> {

packages/firestore/src/lite/reference.ts

@@ -74,6 +74,39 @@ export interface UpdateData {
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   [fieldPath: string]: any;
 }
+// Represents an update object to Firestore document data, which can contain either fields like {a: 2}
+// or dot-separated paths such as {"a.b" : 2} (which updates the nested property "b" in map field "a").
+export type TypedUpdateData<T> = T extends Builtin
+  ? T
+  : T extends Map<infer K, infer V>
+  ? Map<TypedUpdateData<K>, TypedUpdateData<V>>
+  : T extends {}
+  ? { [K in keyof T]?: TypedUpdateData<T[K]> | FieldValue } &
+      NestedUpdateFields<T>
+  : Partial<T>;
+
+// For each field (e.g. "bar"), calculate its nested keys (e.g. {"bar.baz": T1, "bar.quax": T2}), and then
+// intersect them together to make one giant map containing all possible keys (all marked as optional).
+type NestedUpdateFields<T extends Record<string, any>> = UnionToIntersection<
+  {
+    [K in keyof T & string]: T[K] extends Record<string, any> // Only allow nesting for map values
+      ? AddPrefixToKeys<K, TypedUpdateData<T[K]>> // Recurse into map and add "bar." in front of every key
+      : never;
+  }[keyof T & string]
+>;
+
+// Return a new map where every key is prepended with Prefix + dot.
+type AddPrefixToKeys<Prefix extends string, T extends Record<string, any>> =
+  // Remap K => Prefix.K. See https://www.typescriptlang.org/docs/handbook/2/mapped-types.html#key-remapping-via-as
+  { [K in keyof T & string as `${Prefix}.${K}`]+?: T[K] };
+
+// This takes union type U = T1 | T2 | ... and returns a intersected type (T1 & T2 & ...)
+type UnionToIntersection<U> =
+  // Works because "multiple candidates for the same type variable in contra-variant positions causes an intersection type to be inferred"
+  // https://www.typescriptlang.org/docs/handbook/advanced-types.html#type-inference-in-conditional-types
+  (U extends any ? (k: U) => void : never) extends (k: infer I) => void
+    ? I
+    : never;
 
 /**
  * An options object that configures the behavior of {@link @firebase/firestore/lite#(setDoc:1)}, {@link

packages/firestore/src/lite/reference_impl.ts

@@ -44,6 +44,7 @@ import {
   NestedPartialWithFieldValue,
   Query,
   SetOptions,
+  TypedUpdateData,
   UpdateData,
   WithFieldValue
 } from './reference';
@@ -266,9 +267,9 @@ export function setDoc<T>(
  * @returns A Promise resolved once the data has been successfully written
  * to the backend.
  */
-export function updateDoc(
-  reference: DocumentReference<unknown>,
-  data: UpdateData
+export function updateDoc<T>(
+  reference: DocumentReference<T>,
+  data: TypedUpdateData<T>
 ): Promise<void>;
 /**
  * Updates fields in the document referred to by the specified
@@ -296,9 +297,9 @@ export function updateDoc(
   value: unknown,
   ...moreFieldsAndValues: unknown[]
 ): Promise<void>;
-export function updateDoc(
+export function updateDoc<T>(
   reference: DocumentReference<unknown>,
-  fieldOrUpdateData: string | FieldPath | UpdateData,
+  fieldOrUpdateData: string | FieldPath | TypedUpdateData<T>,
   value?: unknown,
   ...moreFieldsAndValues: unknown[]
 ): Promise<void> {

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

@@ -59,7 +59,8 @@ import {
   UpdateData,
   DocumentData,
   WithFieldValue,
-  NestedPartialWithFieldValue
+  NestedPartialWithFieldValue,
+  TypedUpdateData
 } from '../../src/lite/reference';
 import {
   addDoc,
@@ -69,7 +70,11 @@ import {
   setDoc,
   updateDoc
 } from '../../src/lite/reference_impl';
-import { snapshotEqual, QuerySnapshot } from '../../src/lite/snapshot';
+import {
+  snapshotEqual,
+  QuerySnapshot,
+  QueryDocumentSnapshot
+} from '../../src/lite/snapshot';
 import { Timestamp } from '../../src/lite/timestamp';
 import { runTransaction } from '../../src/lite/transaction';
 import { writeBatch } from '../../src/lite/write_batch';
@@ -348,9 +353,9 @@ interface MutationTester {
     data: NestedPartialWithFieldValue<T>,
     options: SetOptions
   ): Promise<void>;
-  update(
-    documentRef: DocumentReference<unknown>,
-    data: UpdateData
+  update<T>(
+    documentRef: DocumentReference<T>,
+    data: TypedUpdateData<T>
   ): Promise<void>;
   update(
     documentRef: DocumentReference<unknown>,
@@ -596,6 +601,144 @@ function genericMutationTests(
       });
     });
 
+    it('temporary sanity check tests', async () => {
+      class TestObject {
+        constructor(
+          readonly outerString: string,
+          readonly outerNum: number,
+          readonly outerArr: string[],
+          readonly nested: {
+            innerNested: {
+              innerNestedNum: number;
+              innerNestedString: string;
+            };
+            innerArr: number[];
+            timestamp: Timestamp;
+          }
+        ) {}
+      }
+
+      const testConverterMerge = {
+        toFirestore(testObj: WithFieldValue<TestObject>, options?: SetOptions) {
+          return { ...testObj };
+        },
+        fromFirestore(snapshot: QueryDocumentSnapshot): TestObject {
+          const data = snapshot.data();
+          return new TestObject(
+            data.outerString,
+            data.outerNum,
+            data.outerArr,
+            data.nested
+          );
+        }
+      };
+
+      return withTestDb(async db => {
+        const coll = collection(db, 'posts');
+        const ref = doc(coll, 'testobj').withConverter(testConverterMerge);
+
+        // Allow Field Values and nested partials.
+        await setDoc(
+          ref,
+          {
+            outerString: deleteField(),
+            nested: {
+              innerNested: {
+                innerNestedNum: increment(1)
+              },
+              innerArr: arrayUnion(2),
+              timestamp: serverTimestamp()
+            }
+          },
+          { merge: true }
+        );
+
+        // Checks for non-existent properties
+        await setDoc(
+          ref,
+          {
+            // @ts-expect-error
+            nonexistent: 'foo'
+          },
+          { merge: true }
+        );
+        await setDoc(
+          ref,
+          {
+            nested: {
+              // @ts-expect-error
+              nonexistent: 'foo'
+            }
+          },
+          { merge: true }
+        );
+
+        // Nested Partials are checked
+        await setDoc(
+          ref,
+          {
+            nested: {
+              innerNested: {
+                // @ts-expect-error
+                innerNestedNum: 'string'
+              },
+              // @ts-expect-error
+              innerArr: 2
+            }
+          },
+          { merge: true }
+        );
+        await setDoc(
+          ref,
+          {
+            // @ts-expect-error
+            nested: 3
+          },
+          { merge: true }
+        );
+
+        // Can use update to verify fields
+        await updateDoc(ref, {
+          // @ts-expect-error
+          outerString: 3,
+          // @ts-expect-error
+          outerNum: [],
+          outerArr: arrayUnion('foo'),
+          nested: {
+            innerNested: {
+              // @ts-expect-error
+              innerNestedNum: 'string'
+            },
+            // @ts-expect-error
+            innerArr: 2,
+            timestamp: serverTimestamp()
+          }
+        });
+
+        // Cannot update nonexistent fields
+        await updateDoc(ref, {
+          // @ts-expect-error
+          nonexistent: 'foo'
+        });
+        await updateDoc(ref, {
+          nested: {
+            // @ts-expect-error
+            nonexistent: 'foo'
+          }
+        });
+
+        // Can use update to check string separated fields
+        await updateDoc(ref, {
+          'nested.innerNested.innerNestedNum': 4,
+          // @ts-expect-error
+          'nested.innerNested.innerNestedString': 4,
+          // @ts-expect-error
+          'nested.innerArr': 3,
+          'nested.timestamp': serverTimestamp()
+        });
+      });
+    });
+
     it('supports partials with mergeFields', async () => {
       return withTestDb(async db => {
         const coll = collection(db, 'posts');