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.

integration/firestore/firebase_export.ts

@@ -50,11 +50,22 @@ export function usesFunctionalApi(): false {
   return false;
 }
 
-const Firestore = firebase.firestore.Firestore;
+const Blob = firebase.firestore.Blob;
+const DocumentReference = firebase.firestore.DocumentReference;
 const FieldPath = firebase.firestore.FieldPath;
-const Timestamp = firebase.firestore.Timestamp;
-const GeoPoint = firebase.firestore.GeoPoint;
 const FieldValue = firebase.firestore.FieldValue;
-const Blob = firebase.firestore.Blob;
+const Firestore = firebase.firestore.Firestore;
+const GeoPoint = firebase.firestore.GeoPoint;
+const QueryDocumentSnapshot = firebase.firestore.QueryDocumentSnapshot;
+const Timestamp = firebase.firestore.Timestamp;
 
-export { Firestore, FieldValue, FieldPath, Timestamp, Blob, GeoPoint };
+export {
+  Blob,
+  DocumentReference,
+  FieldPath,
+  FieldValue,
+  Firestore,
+  GeoPoint,
+  QueryDocumentSnapshot,
+  Timestamp
+};

integration/firestore/firebase_export_memory.ts

@@ -50,11 +50,22 @@ export function usesFunctionalApi(): false {
   return false;
 }
 
-const Firestore = firebase.firestore.Firestore;
+const Blob = firebase.firestore.Blob;
+const DocumentReference = firebase.firestore.DocumentReference;
 const FieldPath = firebase.firestore.FieldPath;
-const Timestamp = firebase.firestore.Timestamp;
-const GeoPoint = firebase.firestore.GeoPoint;
 const FieldValue = firebase.firestore.FieldValue;
-const Blob = firebase.firestore.Blob;
+const Firestore = firebase.firestore.Firestore;
+const GeoPoint = firebase.firestore.GeoPoint;
+const QueryDocumentSnapshot = firebase.firestore.QueryDocumentSnapshot;
+const Timestamp = firebase.firestore.Timestamp;
 
-export { Firestore, FieldValue, FieldPath, Timestamp, Blob, GeoPoint };
+export {
+  Blob,
+  DocumentReference,
+  FieldPath,
+  FieldValue,
+  Firestore,
+  GeoPoint,
+  QueryDocumentSnapshot,
+  Timestamp
+};

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,86 @@ export class WriteBatch
   }
 }
 
+/**
+ * Wraps a `PublicFirestoreDataConverter` translating the types from the
+ * experimental SDK into corresponding types from the Classic SDK before passing
+ * them to the wrapped converter.
+ */
+class FirestoreDataConverter<U>
+  extends Compat<PublicFirestoreDataConverter<U>>
+  implements UntypedFirestoreDataConverter<U> {
+  private static readonly INSTANCES = new WeakMap();
+
+  private constructor(
+    private readonly _firestore: Firestore,
+    private readonly _userDataWriter: UserDataWriter,
+    delegate: PublicFirestoreDataConverter<U>
+  ) {
+    super(delegate);
+  }
+
+  fromFirestore(
+    snapshot: ExpQueryDocumentSnapshot,
+    options?: PublicSnapshotOptions
+  ): U {
+    const expSnapshot = new ExpQueryDocumentSnapshot(
+      this._firestore._delegate,
+      this._userDataWriter,
+      snapshot._key,
+      snapshot._document,
+      snapshot.metadata,
+      /* converter= */ null
+    );
+    return this._delegate.fromFirestore(
+      new QueryDocumentSnapshot(this._firestore, expSnapshot),
+      options ?? {}
+    );
+  }
+
+  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 `FirestoreDataConverter` for the given instances
+  // of `Firestore` and `PublicFirestoreDataConverter` so that isEqual() will
+  // compare equal for two objects created with the same converter instance.
+  static getInstance<U>(
+    firestore: Firestore,
+    converter: PublicFirestoreDataConverter<U>
+  ): FirestoreDataConverter<U> {
+    const converterMapByFirestore = FirestoreDataConverter.INSTANCES;
+    let untypedConverterByConverter = converterMapByFirestore.get(firestore);
+    if (!untypedConverterByConverter) {
+      untypedConverterByConverter = new WeakMap();
+      converterMapByFirestore.set(firestore, untypedConverterByConverter);
+    }
+
+    let instance = untypedConverterByConverter.get(converter);
+    if (!instance) {
+      instance = new FirestoreDataConverter(
+        firestore,
+        new UserDataWriter(firestore),
+        converter
+      );
+      untypedConverterByConverter.set(converter, instance);
+    }
+
+    return instance;
+  }
+}
+
 /**
  * A reference to a particular document in a collection in the database.
  */
@@ -753,7 +834,7 @@ export class DocumentReference<T = PublicDocumentData>
     return new DocumentReference<U>(
       this.firestore,
       this._delegate.withConverter(
-        converter as UntypedFirestoreDataConverter<U>
+        FirestoreDataConverter.getInstance(this.firestore, converter)
       )
     );
   }
@@ -1062,7 +1143,7 @@ export class Query<T = PublicDocumentData>
     return new Query<U>(
       this.firestore,
       this._delegate.withConverter(
-        converter as UntypedFirestoreDataConverter<U>
+        FirestoreDataConverter.getInstance(this.firestore, converter)
       )
     );
   }
@@ -1206,7 +1287,7 @@ export class CollectionReference<T = PublicDocumentData>
     return new CollectionReference<U>(
       this.firestore,
       this._delegate.withConverter(
-        converter as UntypedFirestoreDataConverter<U>
+        FirestoreDataConverter.getInstance(this.firestore, converter)
       )
     );
   }

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

@@ -39,6 +39,8 @@ const newTestFirestore = firebaseExport.newTestFirestore;
 const Timestamp = firebaseExport.Timestamp;
 const FieldPath = firebaseExport.FieldPath;
 const FieldValue = firebaseExport.FieldValue;
+const DocumentReference = firebaseExport.DocumentReference;
+const QueryDocumentSnapshot = firebaseExport.QueryDocumentSnapshot;
 
 const MEMORY_ONLY_BUILD =
   typeof process !== 'undefined' &&
@@ -1326,7 +1328,11 @@ apiDescribe('Database', (persistence: boolean) => {
   // only to web.
   apiDescribe('withConverter() support', (persistence: boolean) => {
     class Post {
-      constructor(readonly title: string, readonly author: string) {}
+      constructor(
+        readonly title: string,
+        readonly author: string,
+        readonly ref: firestore.DocumentReference | null = null
+      ) {}
       byline(): string {
         return this.title + ', by ' + this.author;
       }
@@ -1340,8 +1346,9 @@ apiDescribe('Database', (persistence: boolean) => {
         snapshot: firestore.QueryDocumentSnapshot,
         options: firestore.SnapshotOptions
       ): Post {
+        expect(snapshot).to.be.an.instanceof(QueryDocumentSnapshot);
         const data = snapshot.data(options);
-        return new Post(data.title, data.author);
+        return new Post(data.title, data.author, snapshot.ref);
       }
     };
 
@@ -1369,7 +1376,7 @@ apiDescribe('Database', (persistence: boolean) => {
         options: firestore.SnapshotOptions
       ): Post {
         const data = snapshot.data();
-        return new Post(data.title, data.author);
+        return new Post(data.title, data.author, snapshot.ref);
       }
     };
 
@@ -1552,7 +1559,7 @@ apiDescribe('Database', (persistence: boolean) => {
             expect(options).to.deep.equal({ serverTimestamps: 'estimate' });
 
             const data = snapshot.data(options);
-            return new Post(data.title, data.author);
+            return new Post(data.title, data.author, snapshot.ref);
           }
         });
 
@@ -1590,6 +1597,53 @@ 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(postConverter);
+        await docRef.set(new Post('post', 'author'));
+        const docSnapshot = await docRef.get();
+        const ref = docSnapshot.data()!.ref!;
+        expect(ref).to.be.an.instanceof(DocumentReference);
+        expect(untypedDocRef.isEqual(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(postConverter);
+        const docRef = collection.doc();
+        await docRef.set(new Post('post', 'author', docRef));
+        const querySnapshot = await collection.get();
+        expect(querySnapshot.size).to.equal(1);
+        const ref = querySnapshot.docs[0].data().ref!;
+        expect(ref).to.be.an.instanceof(DocumentReference);
+        const untypedDocRef = untypedCollection.doc(docRef.id);
+        expect(untypedDocRef.isEqual(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(postConverter);
+        await docRef.set(new Post('post', 'author', docRef));
+        const query = untypedCollection
+          .where(FieldPath.documentId(), '==', docRef.id)
+          .withConverter(postConverter);
+        const querySnapshot = await query.get();
+        expect(querySnapshot.size).to.equal(1);
+        const ref = querySnapshot.docs[0].data().ref!;
+        expect(ref).to.be.an.instanceof(DocumentReference);
+        expect(untypedDocRef.isEqual(ref)).to.be.true;
+      });
+    });
   });
 
   it('can set and get data with auto detect long polling enabled', () => {

packages/firestore/test/integration/util/firebase_export.ts

@@ -25,7 +25,11 @@ import { FirebaseApp } from '@firebase/app-types';
 import * as firestore from '@firebase/firestore-types';
 
 import { Blob } from '../../../src/api/blob';
-import { Firestore } from '../../../src/api/database';
+import {
+  Firestore,
+  DocumentReference,
+  QueryDocumentSnapshot
+} from '../../../src/api/database';
 import { FieldPath } from '../../../src/api/field_path';
 import { FieldValue } from '../../../src/api/field_value';
 import { GeoPoint } from '../../../src/api/geo_point';
@@ -68,4 +72,13 @@ export function newTestFirestore(
   return firestore;
 }
 
-export { Firestore, FieldValue, FieldPath, Timestamp, Blob, GeoPoint };
+export {
+  Firestore,
+  FieldValue,
+  FieldPath,
+  Timestamp,
+  Blob,
+  GeoPoint,
+  DocumentReference,
+  QueryDocumentSnapshot
+};