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

@@ -131,6 +131,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);
 
@@ -175,6 +176,7 @@ export class ExpUserDataWriter extends AbstractUserDataWriter {
 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);
   const userDataWriter = new ExpUserDataWriter(firestore);
@@ -210,6 +212,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);
 
@@ -240,6 +243,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);
   const userDataWriter = new ExpUserDataWriter(firestore);
@@ -271,6 +275,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);
   const userDataWriter = new ExpUserDataWriter(firestore);
@@ -294,6 +299,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);
   const userDataWriter = new ExpUserDataWriter(firestore);
@@ -348,6 +354,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(
@@ -412,6 +419,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);
@@ -701,6 +709,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
   };
@@ -733,23 +745,28 @@ 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
+            )
           );
         }
       },
       error: args[currArg + 1] as ErrorFn,
       complete: args[currArg + 2] as CompleteFn
     };
   } else {
-    firestore = cast(reference.firestore, FirebaseFirestore);
-    internalQuery = reference._query;
+    const query = cast<Query<T>>(reference, Query);
+    firestore = cast(query.firestore, FirebaseFirestore);
+    internalQuery = query._query;
     const userDataWriter = new ExpUserDataWriter(firestore);
 
     observer = {
       next: snapshot => {
         if (args[currArg]) {
           (args[currArg] as NextFn<QuerySnapshot<T>>)(
-            new QuerySnapshot(firestore, userDataWriter, reference, snapshot)
+            new QuerySnapshot(firestore, userDataWriter, query, snapshot)
           );
         }
       },
@@ -832,6 +849,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/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';
 
 export { WriteBatch };
 
@@ -34,6 +35,7 @@ export { WriteBatch };
  * 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

@@ -74,6 +74,7 @@ import { DeleteMutation, Precondition } from '../../../src/model/mutation';
 import { applyFirestoreDataConverter } from '../../../src/api/database';
 import { FieldPath } from './field_path';
 import {
+  cast,
   validateCollectionPath,
   validateDocumentPath,
   validateNonEmptyArgument,
@@ -1147,6 +1148,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);
@@ -1194,6 +1199,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(
@@ -1272,6 +1279,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) {
@@ -1341,6 +1352,7 @@ export class LiteUserDataWriter extends AbstractUserDataWriter {
 export function getDoc<T>(
   reference: DocumentReference<T>
 ): Promise<DocumentSnapshot<T>> {
+  reference = cast<DocumentReference<T>>(reference, DocumentReference);
   const datastore = getDatastore(reference.firestore);
   const userDataWriter = new LiteUserDataWriter(reference.firestore);
 
@@ -1372,6 +1384,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);
@@ -1443,6 +1456,7 @@ export function setDoc<T>(
   data: T,
   options?: SetOptions
 ): Promise<void> {
+  reference = cast<DocumentReference<T>>(reference, DocumentReference);
   const convertedValue = applyFirestoreDataConverter(
     reference._converter,
     data,
@@ -1518,6 +1532,7 @@ export function updateDoc(
   value?: unknown,
   ...moreFieldsAndValues: unknown[]
 ): Promise<void> {
+  reference = cast<DocumentReference<unknown>>(reference, DocumentReference);
   const dataReader = newUserDataReader(reference.firestore);
 
   // For Compat types, we have to "extract" the underlying types before
@@ -1567,7 +1582,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())
@@ -1592,6 +1610,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(
@@ -1628,6 +1647,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) &&
@@ -1652,6 +1678,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

@@ -286,6 +286,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

@@ -44,6 +44,7 @@ import {
 } from './reference';
 import { FieldPath } from './field_path';
 import { getDatastore } from './components';
+import { cast } from '../../../src/util/input_validation';
 import { Compat } from '../../../src/compat/compat';
 
 // TODO(mrschmidt) Consider using `BaseTransaction` as the base class in the
@@ -264,6 +265,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';
 import { Compat } from '../../../src/compat/compat';
 
 /**
@@ -269,6 +270,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

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

packages/firestore/src/api/field_path.ts

@@ -55,6 +55,7 @@ export class FieldPath extends Compat<ExpFieldPath> implements PublicFieldPath {
     if (other instanceof Compat) {
       other = other._delegate;
     }
+
     if (!(other instanceof ExpFieldPath)) {
       return false;
     }