Accept Compat types in firestore-exp API

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

@@ -53,6 +53,7 @@ import {
   indexedDbStoragePrefix
 } from '../../../src/local/indexeddb_persistence';
 import { PersistenceSettings } from '../../../exp-types';
+import { cast } from '../../../src/util/input_validation';
 
 /** DOMException error code constants. */
 const DOM_EXCEPTION_INVALID_STATE = 11;
@@ -169,6 +170,7 @@ export function enableIndexedDbPersistence(
   firestore: FirebaseFirestore,
   persistenceSettings?: PersistenceSettings
 ): Promise<void> {
+  firestore = cast(firestore, FirebaseFirestore);
   verifyNotInitialized(firestore);
 
   const client = ensureFirestoreConfigured(firestore);
@@ -212,6 +214,7 @@ export function enableIndexedDbPersistence(
 export function enableMultiTabIndexedDbPersistence(
   firestore: FirebaseFirestore
 ): Promise<void> {
+  firestore = cast(firestore, FirebaseFirestore);
   verifyNotInitialized(firestore);
 
   const client = ensureFirestoreConfigured(firestore);
@@ -366,6 +369,7 @@ export function clearIndexedDbPersistence(
 export function waitForPendingWrites(
   firestore: FirebaseFirestore
 ): Promise<void> {
+  firestore = cast(firestore, FirebaseFirestore);
   const client = ensureFirestoreConfigured(firestore);
   return firestoreClientWaitForPendingWrites(client);
 }
@@ -377,6 +381,7 @@ export function waitForPendingWrites(
  * @return A promise that is resolved once the network has been enabled.
  */
 export function enableNetwork(firestore: FirebaseFirestore): Promise<void> {
+  firestore = cast(firestore, FirebaseFirestore);
   const client = ensureFirestoreConfigured(firestore);
   return firestoreClientEnableNetwork(client);
 }
@@ -390,6 +395,7 @@ export function enableNetwork(firestore: FirebaseFirestore): Promise<void> {
  * @return A promise that is resolved once the network has been disabled.
  */
 export function disableNetwork(firestore: FirebaseFirestore): Promise<void> {
+  firestore = cast(firestore, FirebaseFirestore);
   const client = ensureFirestoreConfigured(firestore);
   return firestoreClientDisableNetwork(client);
 }

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

@@ -109,6 +109,7 @@ export interface SnapshotListenOptions {
 export function getDoc<T>(
   reference: DocumentReference<T>
 ): Promise<DocumentSnapshot<T>> {
+  reference = cast<DocumentReference<T>>(reference, DocumentReference);
   const firestore = cast(reference.firestore, FirebaseFirestore);
   const client = ensureFirestoreConfigured(firestore);
 
@@ -138,6 +139,7 @@ export function getDoc<T>(
 export function getDocFromCache<T>(
   reference: DocumentReference<T>
 ): Promise<DocumentSnapshot<T>> {
+  reference = cast<DocumentReference<T>>(reference, DocumentReference);
   const firestore = cast(reference.firestore, FirebaseFirestore);
   const client = ensureFirestoreConfigured(firestore);
 
@@ -171,6 +173,7 @@ export function getDocFromCache<T>(
 export function getDocFromServer<T>(
   reference: DocumentReference<T>
 ): Promise<DocumentSnapshot<T>> {
+  reference = cast<DocumentReference<T>>(reference, DocumentReference);
   const firestore = cast(reference.firestore, FirebaseFirestore);
   const client = ensureFirestoreConfigured(firestore);
 
@@ -201,6 +204,7 @@ export function getDocFromServer<T>(
  * @return A Promise that will be resolved with the results of the query.
  */
 export function getDocs<T>(query: Query<T>): Promise<QuerySnapshot<T>> {
+  query = cast<Query<T>>(query, Query);
   const firestore = cast(query.firestore, FirebaseFirestore);
   const client = ensureFirestoreConfigured(firestore);
 
@@ -231,6 +235,7 @@ export function getDocs<T>(query: Query<T>): Promise<QuerySnapshot<T>> {
 export function getDocsFromCache<T>(
   query: Query<T>
 ): Promise<QuerySnapshot<T>> {
+  query = cast<Query<T>>(query, Query);
   const firestore = cast(query.firestore, FirebaseFirestore);
   const client = ensureFirestoreConfigured(firestore);
 
@@ -253,6 +258,7 @@ export function getDocsFromCache<T>(
 export function getDocsFromServer<T>(
   query: Query<T>
 ): Promise<QuerySnapshot<T>> {
+  query = cast<Query<T>>(query, Query);
   const firestore = cast(query.firestore, FirebaseFirestore);
   const client = ensureFirestoreConfigured(firestore);
 
@@ -306,6 +312,7 @@ export function setDoc<T>(
   data: T,
   options?: SetOptions
 ): Promise<void> {
+  reference = cast<DocumentReference<T>>(reference, DocumentReference);
   const firestore = cast(reference.firestore, FirebaseFirestore);
 
   const convertedValue = applyFirestoreDataConverter(
@@ -370,6 +377,7 @@ export function updateDoc(
   value?: unknown,
   ...moreFieldsAndValues: unknown[]
 ): Promise<void> {
+  reference = cast<DocumentReference<unknown>>(reference, DocumentReference);
   const firestore = cast(reference.firestore, FirebaseFirestore);
 
   const dataReader = newUserDataReader(firestore);
@@ -659,6 +667,10 @@ export function onSnapshot<T>(
   reference: Query<T> | DocumentReference<T>,
   ...args: unknown[]
 ): Unsubscribe {
+  if (reference instanceof Compat) {
+    reference = reference._delegate;
+  }
+
   let options: SnapshotListenOptions = {
     includeMetadataChanges: false
   };
@@ -691,7 +703,11 @@ export function onSnapshot<T>(
       next: snapshot => {
         if (args[currArg]) {
           (args[currArg] as NextFn<DocumentSnapshot<T>>)(
-            convertToDocSnapshot(firestore, reference, snapshot)
+            convertToDocSnapshot(
+              firestore,
+              reference as DocumentReference<T>,
+              snapshot
+            )
           );
         }
       },
@@ -706,7 +722,7 @@ export function onSnapshot<T>(
       next: snapshot => {
         if (args[currArg]) {
           (args[currArg] as NextFn<QuerySnapshot<T>>)(
-            new QuerySnapshot(firestore, reference, snapshot)
+            new QuerySnapshot(firestore, reference as Query<T>, snapshot)
           );
         }
       },
@@ -789,6 +805,7 @@ export function onSnapshotsInSync(
   firestore: FirebaseFirestore,
   arg: unknown
 ): Unsubscribe {
+  firestore = cast(firestore, FirebaseFirestore);
   const client = ensureFirestoreConfigured(firestore);
 
   const observer = isPartialObserver(arg)

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 { cast } from '../../../src/util/input_validation';
 
 /**
  * A reference to a transaction.
@@ -94,6 +95,7 @@ export function runTransaction<T>(
   firestore: FirebaseFirestore,
   updateFunction: (transaction: Transaction) => Promise<T>
 ): Promise<T> {
+  firestore = cast(firestore, FirebaseFirestore);
   ensureFirestoreConfigured(firestore);
 
   const deferred = new Deferred<T>();

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

@@ -19,6 +19,7 @@ import { WriteBatch } from '../../../lite/src/api/write_batch';
 import { FirebaseFirestore } from './database';
 import { executeWrite } from './reference';
 import { ensureFirestoreConfigured } from '../../../src/api/database';
+import { cast } from '../../../src/util/input_validation';
 
 /**
  * Creates a write batch, used for performing multiple writes as a single
@@ -32,6 +33,7 @@ import { ensureFirestoreConfigured } from '../../../src/api/database';
  * writes.
  */
 export function writeBatch(firestore: FirebaseFirestore): WriteBatch {
+  firestore = cast(firestore, FirebaseFirestore);
   ensureFirestoreConfigured(firestore);
   return new WriteBatch(firestore, mutations =>
     executeWrite(firestore, mutations)

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

@@ -35,7 +35,10 @@ import {
   LRU_DEFAULT_CACHE_SIZE_BYTES,
   LRU_MINIMUM_CACHE_SIZE_BYTES
 } from '../../../src/local/lru_garbage_collector';
-import { validateIsNotUsedTogether } from '../../../src/util/input_validation';
+import {
+  cast,
+  validateIsNotUsedTogether
+} from '../../../src/util/input_validation';
 
 declare module '@firebase/component' {
   interface NameServiceMapping {
@@ -316,6 +319,7 @@ export function getFirestore(app: FirebaseApp): FirebaseFirestore {
  * terminated.
  */
 export function terminate(firestore: FirebaseFirestore): Promise<void> {
+  firestore = cast(firestore, FirebaseFirestore);
   _removeServiceInstance(firestore.app, 'firestore/lite');
   return firestore._delete();
 }

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

@@ -70,6 +70,7 @@ import {
 } from '../../../src/api/database';
 import { FieldPath } from './field_path';
 import {
+  cast,
   validateCollectionPath,
   validateDocumentPath,
   validateNonEmptyArgument,
@@ -78,6 +79,7 @@ import {
 import { newSerializer } from '../../../src/platform/serializer';
 import { Code, FirestoreError } from '../../../src/util/error';
 import { getDatastore } from './components';
+import { Compat } from '../../../src/compat/compat';
 
 /**
  * Document data (for use with {@link setDoc()}) consists of fields mapped to
@@ -733,6 +735,10 @@ export function collection(
   path: string,
   ...pathSegments: string[]
 ): CollectionReference<DocumentData> {
+  if (parent instanceof Compat) {
+    parent = parent._delegate;
+  }
+
   validateNonEmptyArgument('collection', 'path', path);
   if (parent instanceof FirebaseFirestore) {
     const absolutePath = ResourcePath.fromString(path, ...pathSegments);
@@ -780,6 +786,8 @@ export function collectionGroup(
   firestore: FirebaseFirestore,
   collectionId: string
 ): Query<DocumentData> {
+  firestore = cast(firestore, FirebaseFirestore);
+
   validateNonEmptyArgument('collectionGroup', 'collection id', collectionId);
   if (collectionId.indexOf('/') >= 0) {
     throw new FirestoreError(
@@ -858,6 +866,10 @@ export function doc<T>(
   path?: string,
   ...pathSegments: string[]
 ): DocumentReference {
+  if (parent instanceof Compat) {
+    parent = parent._delegate;
+  }
+
   // We allow omission of 'pathString' but explicitly prohibit passing in both
   // 'undefined' and 'null'.
   if (arguments.length === 1) {
@@ -912,6 +924,7 @@ export function doc<T>(
 export function getDoc<T>(
   reference: DocumentReference<T>
 ): Promise<DocumentSnapshot<T>> {
+  reference = cast<DocumentReference<T>>(reference, DocumentReference);
   const datastore = getDatastore(reference.firestore);
   return invokeBatchGetDocumentsRpc(datastore, [reference._key]).then(
     result => {
@@ -940,6 +953,7 @@ export function getDoc<T>(
  * @return A Promise that will be resolved with the results of the query.
  */
 export function getDocs<T>(query: Query<T>): Promise<QuerySnapshot<T>> {
+  query = cast<Query<T>>(query, Query);
   validateHasExplicitOrderByForLimitToLast(query._query);
 
   const datastore = getDatastore(query.firestore);
@@ -1009,6 +1023,7 @@ export function setDoc<T>(
   data: T,
   options?: SetOptions
 ): Promise<void> {
+  reference = cast<DocumentReference<T>>(reference, DocumentReference);
   const convertedValue = applyFirestoreDataConverter(
     reference._converter,
     data,
@@ -1084,6 +1099,7 @@ export function updateDoc(
   value?: unknown,
   ...moreFieldsAndValues: unknown[]
 ): Promise<void> {
+  reference = cast<DocumentReference<unknown>>(reference, DocumentReference);
   const dataReader = newUserDataReader(reference.firestore);
 
   let parsed: ParsedUpdateData;
@@ -1127,7 +1143,10 @@ export function updateDoc(
  * @return A Promise resolved once the document has been successfully
  * deleted from the backend.
  */
-export function deleteDoc(reference: DocumentReference): Promise<void> {
+export function deleteDoc(
+  reference: DocumentReference<unknown>
+): Promise<void> {
+  reference = cast<DocumentReference<unknown>>(reference, DocumentReference);
   const datastore = getDatastore(reference.firestore);
   return invokeCommitRpc(datastore, [
     new DeleteMutation(reference._key, Precondition.none())
@@ -1152,6 +1171,7 @@ export function addDoc<T>(
   reference: CollectionReference<T>,
   data: T
 ): Promise<DocumentReference<T>> {
+  reference = cast<CollectionReference<T>>(reference, CollectionReference);
   const docRef = doc(reference);
 
   const convertedValue = applyFirestoreDataConverter(
@@ -1188,6 +1208,13 @@ export function refEqual<T>(
   left: DocumentReference<T> | CollectionReference<T>,
   right: DocumentReference<T> | CollectionReference<T>
 ): boolean {
+  if (left instanceof Compat) {
+    left = left._delegate;
+  }
+  if (right instanceof Compat) {
+    right = right._delegate;
+  }
+
   if (
     (left instanceof DocumentReference ||
       left instanceof CollectionReference) &&
@@ -1212,6 +1239,13 @@ export function refEqual<T>(
  * Firestore database.
  */
 export function queryEqual<T>(left: Query<T>, right: Query<T>): boolean {
+  if (left instanceof Compat) {
+    left = left._delegate;
+  }
+  if (right instanceof Compat) {
+    right = right._delegate;
+  }
+
   if (left instanceof Query && right instanceof Query) {
     return (
       left.firestore === right.firestore &&

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

@@ -34,6 +34,7 @@ import {
 } from '../../../src/api/user_data_reader';
 import { arrayEquals } from '../../../src/util/misc';
 import { Bytes } from './bytes';
+import { Compat } from '../../../src/compat/compat';
 
 /**
  * Converter used by `withConverter()` to transform user objects of type `T`
@@ -297,6 +298,13 @@ export function snapshotEqual<T>(
   left: DocumentSnapshot<T> | QuerySnapshot<T>,
   right: DocumentSnapshot<T> | QuerySnapshot<T>
 ): boolean {
+  if (left instanceof Compat) {
+    left = left._delegate;
+  }
+  if (right instanceof Compat) {
+    right = right._delegate;
+  }
+
   if (left instanceof DocumentSnapshot && right instanceof DocumentSnapshot) {
     return (
       left._firestore === right._firestore &&

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

@@ -43,6 +43,7 @@ import {
 } from './reference';
 import { FieldPath } from './field_path';
 import { getDatastore } from './components';
+import { cast } from '../../../src/util/input_validation';
 
 // TODO(mrschmidt) Consider using `BaseTransaction` as the base class in the
 // legacy SDK.
@@ -253,6 +254,7 @@ export function runTransaction<T>(
   firestore: FirebaseFirestore,
   updateFunction: (transaction: Transaction) => Promise<T>
 ): Promise<T> {
+  firestore = cast(firestore, FirebaseFirestore);
   const datastore = getDatastore(firestore);
   const deferred = new Deferred<T>();
   new TransactionRunner<T>(

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

@@ -38,6 +38,7 @@ import { FirebaseFirestore } from './database';
 import { invokeCommitRpc } from '../../../src/remote/datastore';
 import { FieldPath } from './field_path';
 import { getDatastore } from './components';
+import { cast } from '../../../src/util/input_validation';
 
 /**
  * A write batch, used to perform multiple writes as a single atomic unit.
@@ -260,6 +261,7 @@ export function validateReference<T>(
  * writes.
  */
 export function writeBatch(firestore: FirebaseFirestore): WriteBatch {
+  firestore = cast(firestore, FirebaseFirestore);
   const datastore = getDatastore(firestore);
   return new WriteBatch(firestore, writes =>
     invokeCommitRpc(datastore, writes)

packages/firestore/src/api/database.ts

@@ -817,6 +817,7 @@ export class DocumentReference<T = PublicDocumentData>
     if (other instanceof Compat) {
       other = other._delegate;
     }
+
     if (!(other instanceof ExpDocumentReference)) {
       return false;
     }