add set() overrides to lite sdk

Author: Brian Chen

packages/firestore/lite/index.d.ts

@@ -45,6 +45,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): T;
 }
 
@@ -182,9 +183,10 @@ export class WriteBatch {
   commit(): Promise<void>;
 }
 
-export type SetOptions =
-  | { merge: true }
-  | { mergeFields: Array<string | FieldPath> };
+export interface SetOptions {
+  readonly merge?: boolean;
+  readonly mergeFields?: Array<string | FieldPath>;
+}
 
 export class DocumentReference<T = DocumentData> {
   private constructor();

packages/firestore/lite/src/api/write_batch.ts

@@ -60,7 +60,6 @@ export class WriteBatch implements firestore.WriteBatch {
     const ref = validateReference(documentRef, this._firestore);
 
     const convertedValue = applyFirestoreDataConverter(ref._converter, value);
-
     const parsed = this._dataReader.parseSetData(
       'WriteBatch.set',
       ref._key,
@@ -162,7 +161,10 @@ export function validateReference<T>(
       'Provided document reference is from a different Firestore instance.'
     );
   } else {
-    return cast(documentRef, DocumentReference) as DocumentReference<T>;
+    return (cast(
+      documentRef,
+      DocumentReference
+    ) as unknown) as DocumentReference<T>;
   }
 }
 

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

@@ -881,7 +881,8 @@ describe('equality', () => {
       expect(refEqual(coll1a, coll2)).to.be.false;
 
       const coll1c = collection(firestore, 'a').withConverter({
-        toFirestore: data => data as firestore.DocumentData,
+        toFirestore: (data: firestore.DocumentData) =>
+          data as firestore.DocumentData,
         fromFirestore: snap => snap.data()
       });
       expect(refEqual(coll1a, coll1c)).to.be.false;
@@ -900,7 +901,8 @@ describe('equality', () => {
       expect(refEqual(doc1a, doc2)).to.be.false;
 
       const doc1c = collection(firestore, 'a').withConverter({
-        toFirestore: data => data as firestore.DocumentData,
+        toFirestore: (data: firestore.DocumentData) =>
+          data as firestore.DocumentData,
         fromFirestore: snap => snap.data()
       });
       expect(refEqual(doc1a, doc1c)).to.be.false;
@@ -985,6 +987,35 @@ describe('withConverter() support', () => {
     }
   };
 
+  const postConverterMerge = {
+    toFirestore(
+      post: Partial<Post>,
+      options?: firestore.SetOptions
+    ): firestore.DocumentData {
+      if (
+        options &&
+        ((options as { merge: true }).merge ||
+          (options as { mergeFields: Array<string | number> }).mergeFields)
+      ) {
+        expect(post).to.not.be.an.instanceof(Post);
+      } else {
+        expect(post).to.be.an.instanceof(Post);
+      }
+      const result: firestore.DocumentData = {};
+      if (post.title) {
+        result.title = post.title;
+      }
+      if (post.author) {
+        result.author = post.author;
+      }
+      return result;
+    },
+    fromFirestore(snapshot: firestore.QueryDocumentSnapshot): Post {
+      const data = snapshot.data();
+      return new Post(data.title, data.author);
+    }
+  };
+
   it('for DocumentReference.withConverter()', () => {
     return withTestDoc(async docRef => {
       docRef = docRef.withConverter(postConverter);
@@ -1045,4 +1076,110 @@ describe('withConverter() support', () => {
       expect(refEqual(docRef, docRef2)).to.be.false;
     });
   });
+
+  it('requires the correct converter for Partial usage', async () => {
+    return withTestDb(async db => {
+      const coll = collection(db, 'posts');
+      const ref = doc(coll, 'post').withConverter(postConverter);
+      const batch = writeBatch(db);
+      expect(() =>
+        batch.set(ref, { title: 'olive' }, { merge: true })
+      ).to.throw(
+        'Function WriteBatch.set() called with invalid data. Unsupported ' +
+          'field value: undefined (found in field author in document posts/post)'
+      );
+    });
+  });
+
+  it('WriteBatch.set() supports partials with merge', async () => {
+    return withTestDb(async db => {
+      const coll = collection(db, 'posts');
+      const ref = doc(coll, 'post').withConverter(postConverterMerge);
+      await setDoc(ref, new Post('walnut', 'author'));
+      const batch = writeBatch(db);
+      batch.set(ref, { title: 'olive' }, { merge: true });
+      await batch.commit();
+      const postDoc = await getDoc(ref);
+      expect(postDoc.get('title')).to.equal('olive');
+      expect(postDoc.get('author')).to.equal('author');
+    });
+  });
+
+  it('WriteBatch.set() supports partials with mergeFields', async () => {
+    return withTestDb(async db => {
+      const coll = collection(db, 'posts');
+      const ref = doc(coll, 'post').withConverter(postConverterMerge);
+      await setDoc(ref, new Post('walnut', 'author'));
+      const batch = writeBatch(db);
+      batch.set(ref, { title: 'olive' }, { merge: true });
+      await batch.commit();
+      const postDoc = await getDoc(ref);
+      expect(postDoc.get('title')).to.equal('olive');
+      expect(postDoc.get('author')).to.equal('author');
+    });
+  });
+
+  it('Transaction.set() supports partials with merge', async () => {
+    return withTestDb(async db => {
+      const coll = collection(db, 'posts');
+      const ref = doc(coll, 'post').withConverter(postConverterMerge);
+      await setDoc(ref, new Post('walnut', 'author'));
+      await runTransaction(db, async tx => {
+        tx.set(
+          ref,
+          { title: 'olive', author: 'person' },
+          { mergeFields: ['title'] }
+        );
+      });
+      const postDoc = await getDoc(ref);
+      expect(postDoc.get('title')).to.equal('olive');
+      expect(postDoc.get('author')).to.equal('author');
+    });
+  });
+
+  it('Transaction.set() supports partials with mergeFields', async () => {
+    return withTestDb(async db => {
+      const coll = collection(db, 'posts');
+      const ref = doc(coll, 'post').withConverter(postConverterMerge);
+      await setDoc(ref, new Post('walnut', 'author'));
+      await runTransaction(db, async tx => {
+        tx.set(
+          ref,
+          { title: 'olive', author: 'person' },
+          { mergeFields: ['title'] }
+        );
+      });
+      const postDoc = await getDoc(ref);
+      expect(postDoc.get('title')).to.equal('olive');
+      expect(postDoc.get('author')).to.equal('author');
+    });
+  });
+
+  it('DocumentReference.set() supports partials with merge', async () => {
+    return withTestDb(async db => {
+      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 });
+      const postDoc = await getDoc(ref);
+      expect(postDoc.get('title')).to.equal('olive');
+      expect(postDoc.get('author')).to.equal('author');
+    });
+  });
+
+  it('DocumentReference.set() supports partials with mergeFields', async () => {
+    return withTestDb(async db => {
+      const coll = collection(db, 'posts');
+      const ref = doc(coll, 'post').withConverter(postConverterMerge);
+      await setDoc(ref, new Post('walnut', 'author'));
+      await setDoc(
+        ref,
+        { title: 'olive', author: 'writer' },
+        { mergeFields: ['title'] }
+      );
+      const postDoc = await getDoc(ref);
+      expect(postDoc.get('title')).to.equal('olive');
+      expect(postDoc.get('author')).to.equal('author');
+    });
+  });
 });

packages/firestore/src/api/user_data_reader.ts

@@ -58,6 +58,10 @@ const RESERVED_FIELD_REGEX = /^__.*__$/;
  */
 export interface UntypedFirestoreDataConverter<T> {
   toFirestore(modelObject: T): firestore.DocumentData;
+  toFirestore(
+    modelObject: Partial<T>,
+    options: firestore.SetOptions
+  ): firestore.DocumentData;
   fromFirestore(snapshot: unknown, options?: unknown): T;
 }