@schmidt-sebastian

Return Firestore-Compat types from QuerySnapshot.docs

Author: schmidt-sebastian

packages/firestore/src/api/database.ts

@@ -967,8 +967,11 @@ export class QueryDocumentSnapshot<T = PublicDocumentData>
 export class Query<T = PublicDocumentData>
   extends Compat<ExpQuery<T>>
   implements PublicQuery<T> {
+  private readonly _userDataWriter: UserDataWriter;
+
   constructor(readonly firestore: Firestore, delegate: ExpQuery<T>) {
     super(delegate);
+    this._userDataWriter = new UserDataWriter(firestore);
   }
 
   where(
@@ -1076,7 +1079,18 @@ export class Query<T = PublicDocumentData>
     } else {
       query = getDocs(this._delegate);
     }
-    return query.then(result => new QuerySnapshot(this.firestore, result));
+    return query.then(
+      result =>
+        new QuerySnapshot(
+          this.firestore,
+          new ExpQuerySnapshot<T>(
+            this.firestore._delegate,
+            this._userDataWriter,
+            this._delegate,
+            result._snapshot
+          )
+        )
+    );
   }
 
   onSnapshot(observer: PartialObserver<PublicQuerySnapshot<T>>): Unsubscribe;
@@ -1100,7 +1114,16 @@ export class Query<T = PublicDocumentData>
     const options = extractSnapshotOptions(args);
     const observer = wrapObserver<QuerySnapshot<T>, ExpQuerySnapshot<T>>(
       args,
-      snap => new QuerySnapshot(this.firestore, snap)
+      snap =>
+        new QuerySnapshot(
+          this.firestore,
+          new ExpQuerySnapshot<T>(
+            this.firestore._delegate,
+            this._userDataWriter,
+            this._delegate,
+            snap._snapshot
+          )
+        )
     );
     return onSnapshot(this._delegate, options, observer);
   }

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

@@ -20,6 +20,7 @@ import { expect } from 'chai';
 import { addEqualityMatcher } from '../../util/equality_matcher';
 import * as firebaseExport from '../util/firebase_export';
 import { apiDescribe, withTestDb, withTestDoc } from '../util/helpers';
+import { EventsAccumulator } from '../util/events_accumulator';
 
 const Blob = firebaseExport.Blob;
 const GeoPoint = firebaseExport.GeoPoint;
@@ -28,17 +29,31 @@ const Timestamp = firebaseExport.Timestamp;
 apiDescribe('Firestore', (persistence: boolean) => {
   addEqualityMatcher();
 
-  function expectRoundtrip(
+  async function expectRoundtrip(
     db: firestore.FirebaseFirestore,
     data: {}
   ): Promise<void> {
-    const doc = db.collection('rooms').doc();
-    return doc
-      .set(data)
-      .then(() => doc.get())
-      .then(snapshot => {
-        expect(snapshot.data()).to.deep.equal(data);
-      });
+    const collection = db.collection(db.collection('a').doc().id);
+    const doc = collection.doc();
+    await doc.set(data);
+
+    let docSnapshot = await doc.get();
+    expect(docSnapshot.data()).to.deep.equal(data);
+
+    // TODO(b/174486484): If we wait for the backend response, the "can read
+    // and write number fields" test fails since we store the backend proto
+    // directly in IndexedDb, which turns -0.0 into 0.0
+    let querySnapshot = await collection.get({ source: 'cache' });
+    docSnapshot = querySnapshot.docs[0];
+    expect(docSnapshot.data()).to.deep.equal(data);
+
+    const eventsAccumulator = new EventsAccumulator<firestore.QuerySnapshot>();
+    const unlisten = collection.onSnapshot(eventsAccumulator.storeEvent);
+    querySnapshot = await eventsAccumulator.awaitEvent();
+    docSnapshot = querySnapshot.docs[0];
+    expect(docSnapshot.data()).to.deep.equal(data);
+
+    unlisten();
   }
 
   it('can read and write null fields', () => {

packages/firestore/test/util/equality_matcher.ts

@@ -16,6 +16,7 @@
  */
 
 import { use } from 'chai';
+import { Indexable } from '../../src/util/misc';
 
 /**
  * Duck-typed interface for objects that have an isEqual() method.
@@ -57,6 +58,15 @@ function customDeepEqual(
       return customMatcher.equalsFn(left, right);
     }
   }
+  if (left && typeof left === 'object' && right && typeof right === 'object') {
+    // The `isEqual` check below returns true if firestore-exp types are
+    // compared with API types from Firestore classic. We do want to
+    // differentiate between these types in our tests to ensure that the we do
+    // not return firestore-exp types in the classic SDK.
+    if ((left as Indexable).constructor !== (right as Indexable).constructor) {
+      return false;
+    }
+  }
   if (typeof left === 'object' && left && 'isEqual' in left) {
     return (left as Equatable<unknown>).isEqual(right);
   }