add support for union fields

Author: Brian Chen

.changeset/clean-cameras-check.md

@@ -2,4 +2,4 @@
 '@firebase/firestore': minor
 ---
 
-Fixed a bug where UpdateData did not recognize optional, dot-separated string fields
+Fixed a bug where `UpdateData` did not recognize optional, dot-separated string fields.

common/api-review/firestore-lite.api.md

@@ -32,6 +32,9 @@ export class Bytes {
     toUint8Array(): Uint8Array;
 }
 
+// @public
+export type ChildUpdateFields<T, K extends string> = T extends Record<string, any> ? AddPrefixToKeys<K, UpdateData<T>> : never;
+
 // @public
 export function collection(firestore: Firestore, path: string, ...pathSegments: string[]): CollectionReference<DocumentData>;
 
@@ -191,7 +194,7 @@ export { LogLevel }
 
 // @public
 export type NestedUpdateFields<T extends Record<string, unknown>> = UnionToIntersection<{
-    [K in keyof T & string]: T[K] extends Record<string, unknown> ? AddPrefixToKeys<K, UpdateData<T[K]>> : never;
+    [K in keyof T & string]: ChildUpdateFields<T[K], K>;
 }[keyof T & string]>;
 
 // @public
@@ -332,7 +335,7 @@ export class Transaction {
 export type UnionToIntersection<U> = (U extends unknown ? (k: U) => void : never) extends (k: infer I) => void ? I : never;
 
 // @public
-export type UpdateData<T> = T extends Primitive ? T : T extends Map<infer K, infer V> ? Map<UpdateData<K>, UpdateData<V>> : T extends {} ? {
+export type UpdateData<T> = T extends Primitive ? T : T extends {} ? {
     [K in keyof T]?: UpdateData<T[K]> | FieldValue;
 } & NestedUpdateFields<T> : Partial<T>;
 

common/api-review/firestore.api.md

@@ -35,6 +35,9 @@ export class Bytes {
 // @public
 export const CACHE_SIZE_UNLIMITED = -1;
 
+// @public
+export type ChildUpdateFields<T, K extends string> = T extends Record<string, any> ? AddPrefixToKeys<K, UpdateData<T>> : never;
+
 // @public
 export function clearIndexedDbPersistence(firestore: Firestore): Promise<void>;
 
@@ -265,7 +268,7 @@ export function namedQuery(firestore: Firestore, name: string): Promise<Query |
 
 // @public
 export type NestedUpdateFields<T extends Record<string, unknown>> = UnionToIntersection<{
-    [K in keyof T & string]: T[K] extends Record<string, unknown> ? AddPrefixToKeys<K, UpdateData<T[K]>> : never;
+    [K in keyof T & string]: ChildUpdateFields<T[K], K>;
 }[keyof T & string]>;
 
 // @public
@@ -481,7 +484,7 @@ export interface Unsubscribe {
 }
 
 // @public
-export type UpdateData<T> = T extends Primitive ? T : T extends Map<infer K, infer V> ? Map<UpdateData<K>, UpdateData<V>> : T extends {} ? {
+export type UpdateData<T> = T extends Primitive ? T : T extends {} ? {
     [K in keyof T]?: UpdateData<T[K]> | FieldValue;
 } & NestedUpdateFields<T> : Partial<T>;
 

packages/firestore/lite/index.ts

@@ -82,6 +82,7 @@ export {
 export {
   Primitive,
   NestedUpdateFields,
+  ChildUpdateFields,
   AddPrefixToKeys,
   UnionToIntersection
 } from '../src/lite-api/types';

packages/firestore/src/api.ts

@@ -136,6 +136,7 @@ export { AbstractUserDataWriter } from './lite-api/user_data_writer';
 export {
   Primitive,
   NestedUpdateFields,
+  ChildUpdateFields,
   AddPrefixToKeys,
   UnionToIntersection
 } from '../src/lite-api/types';

packages/firestore/src/lite-api/reference.ts

@@ -76,14 +76,11 @@ export type WithFieldValue<T> = T extends Primitive
  * reference nested fields within the document. FieldValues can be passed in
  * as property values.
  */
-export type UpdateData<T> = T extends undefined
-  ? never
-  : T extends Primitive
+export type UpdateData<T> = T extends Primitive
   ? T
   : T extends {}
   ? { [K in keyof T]?: UpdateData<T[K]> | FieldValue } & NestedUpdateFields<T>
   : Partial<T>;
-
 /**
  * An options object that configures the behavior of {@link @firebase/firestore/lite#(setDoc:1)}, {@link
  * @firebase/firestore/lite#(WriteBatch.set:1)} and {@link @firebase/firestore/lite#(Transaction.set:1)} calls. These calls can be

packages/firestore/src/lite-api/types.ts

@@ -35,15 +35,26 @@ export type NestedUpdateFields<T extends Record<string, unknown>> =
   UnionToIntersection<
     {
       // Check that T[K] extends Record to only allow nesting for map values.
-      [K in keyof T & string]: T[K] extends Record<string, unknown> | undefined
-        ? // Recurse into the map and add the prefix in front of each key
-          // (e.g. Prefix 'bar.' to create: 'bar.baz' and 'bar.qux'.
-          AddPrefixToKeys<K, UpdateData<T[K]>>
-        : // TypedUpdateData is always a map of values.
-          never;
+      // Union with `undefined` to allow for optional nested fields
+      [K in keyof T & string]: ChildUpdateFields<T[K], K>;
     }[keyof T & string] // Also include the generated prefix-string keys.
   >;
 
+/**
+ * Helper for calculating the nested fields for a given type T. This is needed
+ * to distribute union types such as `undefined | {...}` (happens for optional
+ * props) or `{a: A} | {b: B}`.
+ * https://www.typescriptlang.org/docs/handbook/advanced-types.html#distributive-conditional-types
+ */
+export type ChildUpdateFields<T, K extends string> =
+  // Only allow nesting for map values
+  T extends Record<string, any>
+    ? // Recurse into the map and add the prefix in front of each key
+      // (e.g. Prefix 'bar.' to create: 'bar.baz' and 'bar.qux'.
+      AddPrefixToKeys<K, UpdateData<T>>
+    : // UpdateData is always a map of values.
+      never;
+
 /**
  * Returns a new map where every key is prefixed with the outer key appended
  * to a dot.

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

@@ -1585,6 +1585,63 @@ describe('withConverter() support', () => {
         });
       });
 
+      it('supports union fields', () => {
+        interface TestObjectUnion {
+          optionalStr?: string;
+          nested?:
+            | {
+                requiredStr: string;
+              }
+            | { requiredStrObject: string };
+        }
+
+        const testConverterUnion = {
+          toFirestore(testObj: WithFieldValue<TestObjectUnion>) {
+            return { ...testObj };
+          },
+          fromFirestore(snapshot: QueryDocumentSnapshot): TestObjectUnion {
+            const data = snapshot.data();
+            return {
+              optionalStr: data.optionalStr,
+              nested: data.nested
+            };
+          }
+        };
+
+        return withTestDocAndInitialData(initialData, async docRef => {
+          const testDocRef: DocumentReference<TestObjectUnion> =
+            docRef.withConverter(testConverterUnion);
+
+          await updateDoc(testDocRef, {
+            optionalStr: 'foo'
+          });
+          await updateDoc(testDocRef, {
+            'optionalStr': 'foo'
+          });
+
+          await updateDoc(testDocRef, {
+            nested: {
+              requiredStr: 'foo'
+            }
+          });
+          await updateDoc(testDocRef, {
+            'nested.requiredStr': 'foo'
+          });
+          await updateDoc(testDocRef, {
+            // @ts-expect-error
+            'nested.requiredStr': 1
+          });
+          await updateDoc(testDocRef, {
+            'nested.requiredStrObject': 'foo'
+          });
+
+          await updateDoc(testDocRef, {
+            // @ts-expect-error
+            'nested.requiredStringObject': 1
+          });
+        });
+      });
+
       it('checks for nonexistent fields', () => {
         return withTestDocAndInitialData(initialData, async docRef => {
           const testDocRef: DocumentReference<TestObject> =