Fix FirestoreDataConverter.fromFirestore() being called with QueryDocumentSnapshot from exp

.changeset/neat-yaks-kneel.md

@@ -0,0 +1,5 @@
+---
+'@firebase/firestore': patch
+---
+
+Fixes FirestoreDataConverter.fromFirestore() being called with an incorrect "snapshot" object.

packages/firestore/src/api/database.ts

@@ -99,6 +99,7 @@ import {
 import {
   DocumentChange as ExpDocumentChange,
   DocumentSnapshot as ExpDocumentSnapshot,
+  QueryDocumentSnapshot as ExpQueryDocumentSnapshot,
   QuerySnapshot as ExpQuerySnapshot,
   snapshotEqual,
   SnapshotMetadata
@@ -552,6 +553,61 @@ export class WriteBatch
   }
 }
 
+class UntypedFirestoreDataConverterAdapter<U>
+  extends Compat<PublicFirestoreDataConverter<U>>
+  implements UntypedFirestoreDataConverter<U> {
+  private static readonly INSTANCES = new WeakMap();
+
+  private constructor(
+    private readonly _firestore: Firestore,
+    delegate: PublicFirestoreDataConverter<U>
+  ) {
+    super(delegate);
+  }
+
+  fromFirestore(snapshot: unknown, options?: unknown): U {
+    return this._delegate.fromFirestore(
+      new QueryDocumentSnapshot<U>(
+        this._firestore,
+        snapshot as ExpQueryDocumentSnapshot<U>
+      ),
+      options as PublicSnapshotOptions
+    );
+  }
+
+  toFirestore(modelObject: U): PublicDocumentData;
+  toFirestore(
+    modelObject: Partial<U>,
+    options: PublicSetOptions
+  ): PublicDocumentData;
+  toFirestore(
+    modelObject: U | Partial<U>,
+    options?: PublicSetOptions
+  ): PublicDocumentData {
+    if (!options) {
+      return this._delegate.toFirestore(modelObject as U);
+    } else {
+      return this._delegate.toFirestore(modelObject, options);
+    }
+  }
+
+  // Use the same instance of UntypedFirestoreDataConverterAdapter for a given
+  // instance of PublicFirestoreDataConverter so that isEqual() will compare
+  // equal for two objects created with the same converter instance.
+  static getInstance<U>(
+    firestore: Firestore,
+    converter: PublicFirestoreDataConverter<U>
+  ): UntypedFirestoreDataConverterAdapter<U> {
+    const map = UntypedFirestoreDataConverterAdapter.INSTANCES;
+    let instance = map.get(converter);
+    if (!instance) {
+      instance = new UntypedFirestoreDataConverterAdapter(firestore, converter);
+      map.set(converter, instance);
+    }
+    return instance;
+  }
+}
+
 /**
  * A reference to a particular document in a collection in the database.
  */
@@ -753,7 +809,10 @@ export class DocumentReference<T = PublicDocumentData>
     return new DocumentReference<U>(
       this.firestore,
       this._delegate.withConverter(
-        converter as UntypedFirestoreDataConverter<U>
+        UntypedFirestoreDataConverterAdapter.getInstance(
+          this.firestore,
+          converter
+        )
       )
     );
   }
@@ -1062,7 +1121,10 @@ export class Query<T = PublicDocumentData>
     return new Query<U>(
       this.firestore,
       this._delegate.withConverter(
-        converter as UntypedFirestoreDataConverter<U>
+        UntypedFirestoreDataConverterAdapter.getInstance(
+          this.firestore,
+          converter
+        )
       )
     );
   }
@@ -1206,7 +1268,10 @@ export class CollectionReference<T = PublicDocumentData>
     return new CollectionReference<U>(
       this.firestore,
       this._delegate.withConverter(
-        converter as UntypedFirestoreDataConverter<U>
+        UntypedFirestoreDataConverterAdapter.getInstance(
+          this.firestore,
+          converter
+        )
       )
     );
   }

packages/firestore/test/integration/api/database.test.ts

@@ -1373,6 +1373,25 @@ apiDescribe('Database', (persistence: boolean) => {
       }
     };
 
+    class ModelWithRef {
+      constructor(readonly ref: firestore.DocumentReference<ModelWithRef>) {}
+    }
+
+    const modelWithRefConverter = {
+      toFirestore(model: ModelWithRef): firestore.DocumentData {
+        // Remove the "ref" attribute, since it's not meant to be stored in
+        // the Firestore database.
+        const { ['ref']: _, ...rest } = model;
+        return rest;
+      },
+      fromFirestore(
+        snapshot: firestore.QueryDocumentSnapshot<ModelWithRef>,
+        options: firestore.SnapshotOptions
+      ): ModelWithRef {
+        return new ModelWithRef(snapshot.ref);
+      }
+    };
+
     it('for DocumentReference.withConverter()', () => {
       return withTestDb(persistence, async db => {
         const docRef = db
@@ -1590,6 +1609,51 @@ apiDescribe('Database', (persistence: boolean) => {
         expect(docRef.isEqual(docRef2)).to.be.false;
       });
     });
+
+    it('Correct snapshot specified to fromFirestore() when registered with DocumentReference', () => {
+      return withTestDb(persistence, async db => {
+        const untypedDocRef = db.collection('/models').doc();
+        const docRef = untypedDocRef.withConverter(modelWithRefConverter);
+        await docRef.set(new ModelWithRef(docRef));
+        const docSnapshot = await docRef.get();
+        expect(untypedDocRef.isEqual(docSnapshot.data()!.ref)).to.be.true;
+      });
+    });
+
+    it('Correct snapshot specified to fromFirestore() when registered with CollectionReference', () => {
+      return withTestDb(persistence, async db => {
+        const untypedCollection = db
+          .collection('/models')
+          .doc()
+          .collection('sub');
+        const collection = untypedCollection.withConverter(
+          modelWithRefConverter
+        );
+        const docRef = collection.doc();
+        await docRef.set(new ModelWithRef(docRef));
+        const querySnapshot = await collection.get();
+        expect(querySnapshot.size).to.equal(1);
+        const untypedDocRef = untypedCollection.doc(docRef.id);
+        expect(untypedDocRef.isEqual(querySnapshot.docs[0].data().ref)).to.be
+          .true;
+      });
+    });
+
+    it('Correct snapshot specified to fromFirestore() when registered with Query', () => {
+      return withTestDb(persistence, async db => {
+        const untypedCollection = db.collection('/models');
+        const untypedDocRef = untypedCollection.doc();
+        const docRef = untypedDocRef.withConverter(modelWithRefConverter);
+        await docRef.set(new ModelWithRef(docRef));
+        const query = untypedCollection
+          .where(FieldPath.documentId(), '==', docRef.id)
+          .withConverter(modelWithRefConverter);
+        const querySnapshot = await query.get();
+        expect(querySnapshot.size).to.equal(1);
+        expect(untypedDocRef.isEqual(querySnapshot.docs[0].data().ref)).to.be
+          .true;
+      });
+    });
   });
 
   it('can set and get data with auto detect long polling enabled', () => {