Compat class for DocumentSnapshot (#4051)

Author: schmidt-sebastian

.changeset/fast-impalas-tickle.md

@@ -0,0 +1,5 @@
+---
+"@firebase/firestore": patch
+---
+
+Fixed an issue that caused `DocumentReference`s in `DocumentSnapshot`s to be returned with the custom converter of the original `DocumentReference`. 

packages/firestore/exp/src/api/reference.ts

@@ -57,13 +57,13 @@ import {
   Unsubscribe
 } from '../../../src/api/observer';
 import {
-  getEventManager,
   executeQueryFromCache,
   executeQueryViaSnapshotListener,
-  readDocumentFromCache,
-  readDocumentViaSnapshotListener,
+  firestoreClientWrite,
+  getEventManager,
   getLocalStore,
-  firestoreClientWrite
+  readDocumentFromCache,
+  readDocumentViaSnapshotListener
 } from '../../../src/core/firestore_client';
 import {
   newQueryForPath,
@@ -80,6 +80,15 @@ import {
 } from '../../../src/core/event_manager';
 import { FirestoreError } from '../../../src/util/error';
 import { Compat } from '../../../src/compat/compat';
+import { ByteString } from '../../../src/util/byte_string';
+import { Bytes } from '../../../lite/src/api/bytes';
+import { AbstractUserDataWriter } from '../../../src/api/user_data_writer';
+
+export {
+  DocumentReference,
+  CollectionReference,
+  Query
+} from '../../../lite/src/api/reference';
 
 /**
  * An options object that can be passed to {@link onSnapshot()} and {@link
@@ -128,6 +137,21 @@ export function getDoc<T>(
   );
 }
 
+export class ExpUserDataWriter extends AbstractUserDataWriter {
+  constructor(protected firestore: FirebaseFirestore) {
+    super();
+  }
+
+  protected convertBytes(bytes: ByteString): Bytes {
+    return new Bytes(bytes);
+  }
+
+  protected convertReference(name: string): DocumentReference {
+    const key = this.convertDocumentKey(name, this.firestore._databaseId);
+    return new DocumentReference(this.firestore, /* converter= */ null, key);
+  }
+}
+
 /**
  * Reads the document referred to by this `DocumentReference` from cache.
  * Returns an error if the document is not currently cached.
@@ -140,6 +164,7 @@ export function getDocFromCache<T>(
 ): Promise<DocumentSnapshot<T>> {
   const firestore = cast(reference.firestore, FirebaseFirestore);
   const client = ensureFirestoreConfigured(firestore);
+  const userDataWriter = new ExpUserDataWriter(firestore);
 
   const deferred = new Deferred<Document | null>();
   firestore._queue.enqueueAndForget(async () => {
@@ -150,6 +175,7 @@ export function getDocFromCache<T>(
     doc =>
       new DocumentSnapshot(
         firestore,
+        userDataWriter,
         reference._key,
         doc,
         new SnapshotMetadata(
@@ -203,6 +229,7 @@ export function getDocFromServer<T>(
 export function getDocs<T>(query: Query<T>): Promise<QuerySnapshot<T>> {
   const firestore = cast(query.firestore, FirebaseFirestore);
   const client = ensureFirestoreConfigured(firestore);
+  const userDataWriter = new ExpUserDataWriter(firestore);
 
   validateHasExplicitOrderByForLimitToLast(query._query);
 
@@ -218,7 +245,7 @@ export function getDocs<T>(query: Query<T>): Promise<QuerySnapshot<T>> {
     );
   });
   return deferred.promise.then(
-    snapshot => new QuerySnapshot(firestore, query, snapshot)
+    snapshot => new QuerySnapshot(firestore, userDataWriter, query, snapshot)
   );
 }
 
@@ -233,14 +260,15 @@ export function getDocsFromCache<T>(
 ): Promise<QuerySnapshot<T>> {
   const firestore = cast(query.firestore, FirebaseFirestore);
   const client = ensureFirestoreConfigured(firestore);
+  const userDataWriter = new ExpUserDataWriter(firestore);
 
   const deferred = new Deferred<ViewSnapshot>();
   firestore._queue.enqueueAndForget(async () => {
     const localStore = await getLocalStore(client);
     await executeQueryFromCache(localStore, query._query, deferred);
   });
   return deferred.promise.then(
-    snapshot => new QuerySnapshot(firestore, query, snapshot)
+    snapshot => new QuerySnapshot(firestore, userDataWriter, query, snapshot)
   );
 }
 
@@ -255,6 +283,7 @@ export function getDocsFromServer<T>(
 ): Promise<QuerySnapshot<T>> {
   const firestore = cast(query.firestore, FirebaseFirestore);
   const client = ensureFirestoreConfigured(firestore);
+  const userDataWriter = new ExpUserDataWriter(firestore);
 
   const deferred = new Deferred<ViewSnapshot>();
   firestore._queue.enqueueAndForget(async () => {
@@ -268,7 +297,7 @@ export function getDocsFromServer<T>(
     );
   });
   return deferred.promise.then(
-    snapshot => new QuerySnapshot(firestore, query, snapshot)
+    snapshot => new QuerySnapshot(firestore, userDataWriter, query, snapshot)
   );
 }
 
@@ -701,12 +730,13 @@ export function onSnapshot<T>(
   } else {
     firestore = cast(reference.firestore, FirebaseFirestore);
     internalQuery = reference._query;
+    const userDataWriter = new ExpUserDataWriter(firestore);
 
     observer = {
       next: snapshot => {
         if (args[currArg]) {
           (args[currArg] as NextFn<QuerySnapshot<T>>)(
-            new QuerySnapshot(firestore, reference, snapshot)
+            new QuerySnapshot(firestore, userDataWriter, reference, snapshot)
           );
         }
       },
@@ -836,8 +866,10 @@ function convertToDocSnapshot<T>(
   );
   const doc = snapshot.docs.get(ref._key);
 
+  const userDataWriter = new ExpUserDataWriter(firestore);
   return new DocumentSnapshot(
     firestore,
+    userDataWriter,
     ref._key,
     doc,
     new SnapshotMetadata(snapshot.hasPendingWrites, snapshot.fromCache),

packages/firestore/exp/src/api/snapshot.ts

@@ -17,18 +17,14 @@
 
 import { DocumentKey } from '../../../src/model/document_key';
 import { Document } from '../../../src/model/document';
-import {
-  ServerTimestampBehavior,
-  UserDataWriter
-} from '../../../src/api/user_data_writer';
+import { AbstractUserDataWriter } from '../../../src/api/user_data_writer';
 import {
   DocumentSnapshot as LiteDocumentSnapshot,
   fieldPathFromArgument
 } from '../../../lite/src/api/snapshot';
 import { FirebaseFirestore } from './database';
 import {
   DocumentData,
-  DocumentReference,
   Query,
   queryEqual,
   SetOptions
@@ -41,9 +37,7 @@ import { Code, FirestoreError } from '../../../src/util/error';
 import { ViewSnapshot } from '../../../src/core/view_snapshot';
 import { FieldPath } from '../../../lite/src/api/field_path';
 import { SnapshotListenOptions } from './reference';
-import { Bytes } from '../../../lite/src/api/bytes';
-
-const DEFAULT_SERVER_TIMESTAMP_BEHAVIOR: ServerTimestampBehavior = 'none';
+import { UntypedFirestoreDataConverter } from '../../../src/api/user_data_reader';
 
 /**
  * Converter used by `withConverter()` to transform user objects of type `T`
@@ -187,12 +181,13 @@ export class DocumentSnapshot<T = DocumentData> extends LiteDocumentSnapshot<
 
   constructor(
     readonly _firestore: FirebaseFirestore,
+    userDataWriter: AbstractUserDataWriter,
     key: DocumentKey,
     document: Document | null,
     metadata: SnapshotMetadata,
-    converter: FirestoreDataConverter<T> | null
+    converter: UntypedFirestoreDataConverter<T> | null
   ) {
-    super(_firestore, key, document, converter);
+    super(_firestore, userDataWriter, key, document, converter);
     this._firestoreImpl = _firestore;
     this.metadata = metadata;
   }
@@ -219,29 +214,26 @@ export class DocumentSnapshot<T = DocumentData> extends LiteDocumentSnapshot<
    * @return An `Object` containing all fields in the document or `undefined` if
    * the document doesn't exist.
    */
-  data(options?: SnapshotOptions): T | undefined {
+  data(options: SnapshotOptions = {}): T | undefined {
     if (!this._document) {
       return undefined;
     } else if (this._converter) {
       // We only want to use the converter and create a new DocumentSnapshot
       // if a converter has been provided.
       const snapshot = new QueryDocumentSnapshot(
         this._firestore,
+        this._userDataWriter,
         this._key,
         this._document,
         this.metadata,
         /* converter= */ null
       );
       return this._converter.fromFirestore(snapshot, options);
     } else {
-      const userDataWriter = new UserDataWriter(
-        this._firestoreImpl._databaseId,
-        options?.serverTimestamps || DEFAULT_SERVER_TIMESTAMP_BEHAVIOR,
-        key =>
-          new DocumentReference(this._firestore, /* converter= */ null, key),
-        bytes => new Bytes(bytes)
-      );
-      return userDataWriter.convertValue(this._document.toProto()) as T;
+      return this._userDataWriter.convertValue(
+        this._document.toProto(),
+        options.serverTimestamps
+      ) as T;
     }
   }
 
@@ -269,13 +261,10 @@ export class DocumentSnapshot<T = DocumentData> extends LiteDocumentSnapshot<
         .data()
         .field(fieldPathFromArgument('DocumentSnapshot.get', fieldPath));
       if (value !== null) {
-        const userDataWriter = new UserDataWriter(
-          this._firestoreImpl._databaseId,
-          options.serverTimestamps || DEFAULT_SERVER_TIMESTAMP_BEHAVIOR,
-          key => new DocumentReference(this._firestore, this._converter, key),
-          bytes => new Bytes(bytes)
+        return this._userDataWriter.convertValue(
+          value,
+          options.serverTimestamps
         );
-        return userDataWriter.convertValue(value);
       }
     }
     return undefined;
@@ -339,6 +328,7 @@ export class QuerySnapshot<T = DocumentData> {
 
   constructor(
     readonly _firestore: FirebaseFirestore,
+    readonly _userDataWriter: AbstractUserDataWriter,
     query: Query<T>,
     readonly _snapshot: ViewSnapshot
   ) {
@@ -431,6 +421,7 @@ export class QuerySnapshot<T = DocumentData> {
   ): QueryDocumentSnapshot<T> {
     return new QueryDocumentSnapshot<T>(
       this._firestore,
+      this._userDataWriter,
       doc.key,
       doc,
       new SnapshotMetadata(hasPendingWrites, fromCache),

packages/firestore/exp/src/api/transaction.ts

@@ -29,6 +29,7 @@ import { Transaction as InternalTransaction } from '../../../src/core/transactio
 import { validateReference } from '../../../lite/src/api/write_batch';
 import { getDatastore } from '../../../lite/src/api/components';
 import { DocumentReference } from '../../../lite/src/api/reference';
+import { ExpUserDataWriter } from './reference';
 
 /**
  * A reference to a transaction.
@@ -56,12 +57,14 @@ export class Transaction extends LiteTransaction {
    */
   get<T>(documentRef: DocumentReference<T>): Promise<DocumentSnapshot<T>> {
     const ref = validateReference<T>(documentRef, this._firestore);
+    const userDataWriter = new ExpUserDataWriter(this._firestore);
     return super
       .get(documentRef)
       .then(
         liteDocumentSnapshot =>
           new DocumentSnapshot(
             this._firestore,
+            userDataWriter,
             ref._key,
             liteDocumentSnapshot._document,
             new SnapshotMetadata(

packages/firestore/exp/test/shim.ts

@@ -49,6 +49,8 @@ import { Compat } from '../../src/compat/compat';
 import {
   Firestore,
   DocumentReference,
+  DocumentSnapshot,
+  QueryDocumentSnapshot,
   wrapObserver,
   extractSnapshotOptions
 } from '../../src/api/database';
@@ -184,48 +186,6 @@ export class WriteBatch
   }
 }
 
-export class DocumentSnapshot<T = legacy.DocumentData>
-  extends Compat<exp.DocumentSnapshot<T>>
-  implements legacy.DocumentSnapshot<T> {
-  constructor(
-    private readonly _firestore: Firestore,
-    delegate: exp.DocumentSnapshot<T>
-  ) {
-    super(delegate);
-  }
-
-  readonly ref = new DocumentReference<T>(this._firestore, this._delegate.ref);
-  readonly id = this._delegate.id;
-  readonly metadata = this._delegate.metadata;
-
-  get exists(): boolean {
-    return this._delegate.exists();
-  }
-
-  data(options?: legacy.SnapshotOptions): T | undefined {
-    return wrap(this._firestore, this._delegate.data(options));
-  }
-
-  get(fieldPath: string | FieldPath, options?: legacy.SnapshotOptions): any {
-    return wrap(
-      this._firestore,
-      this._delegate.get(unwrap(fieldPath), options)
-    );
-  }
-
-  isEqual(other: DocumentSnapshot<T>): boolean {
-    return snapshotEqual(this._delegate, other._delegate);
-  }
-}
-
-export class QueryDocumentSnapshot<T = legacy.DocumentData>
-  extends DocumentSnapshot<T>
-  implements legacy.QueryDocumentSnapshot<T> {
-  data(options?: legacy.SnapshotOptions): T {
-    return this._delegate.data(options)!;
-  }
-}
-
 export class Query<T = legacy.DocumentData>
   extends Compat<exp.Query<T>>
   implements legacy.Query<T> {
@@ -496,30 +456,6 @@ export class Blob extends Compat<BytesExp> implements legacy.Blob {
   }
 }
 
-/**
- * Takes document data that uses the firestore-exp API types and replaces them
- * with the API types defined in this shim.
- */
-function wrap(firestore: Firestore, value: any): any {
-  if (Array.isArray(value)) {
-    return value.map(v => wrap(firestore, v));
-  } else if (value instanceof FieldPathExp) {
-    return new FieldPath(...value._internalPath.toArray());
-  } else if (value instanceof BytesExp) {
-    return new Blob(value);
-  } else if (isPlainObject(value)) {
-    const obj: any = {};
-    for (const key in value) {
-      if (value.hasOwnProperty(key)) {
-        obj[key] = wrap(firestore, value[key]);
-      }
-    }
-    return obj;
-  } else {
-    return value;
-  }
-}
-
 /**
  * Takes user data that uses API types from this shim and replaces them
  * with the the firestore-exp API types.