[3/4] Integrate realtime ppl into SDK cache management and tests

Author: wu-hui

packages/firestore/src/api/pipeline_impl.ts

@@ -15,8 +15,21 @@
  * limitations under the License.
  */
 
-import { Pipeline } from '../api/pipeline';
-import { firestoreClientExecutePipeline } from '../core/firestore_client';
+// Re-adding necessary imports that were removed previously
+import {
+  CompleteFn,
+  ErrorFn,
+  isPartialObserver,
+  NextFn,
+  PartialObserver
+} from '../api/observer';
+import {
+  firestoreClientExecutePipeline,
+  firestoreClientListen
+} from '../core/firestore_client';
+import { ListenerDataSource } from '../core/event_manager';
+import { toCorePipeline } from '../core/pipeline-util';
+import { ViewSnapshot } from '../core/view_snapshot';
 import {
   StructuredPipeline,
   StructuredPipelineOptions
@@ -31,10 +44,15 @@ import {
   UserDataReader,
   UserDataSource
 } from '../lite-api/user_data_reader';
+import { FirestoreError } from '../util/error';
 import { cast } from '../util/input_validation';
 
 import { ensureFirestoreConfigured, Firestore } from './database';
+import { Pipeline } from './pipeline'; // Keep this specific Pipeline import if needed alongside LitePipeline
+import { RealtimePipeline } from './realtime_pipeline';
 import { DocumentReference } from './reference';
+import { SnapshotListenOptions, Unsubscribe } from './reference_impl';
+import { RealtimePipelineSnapshot } from './snapshot';
 import { ExpUserDataWriter } from './user_data_writer';
 
 declare module './database' {
@@ -49,6 +67,7 @@ declare module './database' {
    */
   interface Firestore {
     pipeline(): PipelineSource<Pipeline>;
+    realtimePipeline(): PipelineSource<RealtimePipeline>;
   }
 }
 
@@ -179,3 +198,115 @@ Firestore.prototype.pipeline = function (): PipelineSource<Pipeline> {
     }
   );
 };
+
+Firestore.prototype.realtimePipeline =
+  function (): PipelineSource<RealtimePipeline> {
+    const userDataReader = newUserDataReader(this);
+    return new PipelineSource<RealtimePipeline>(
+      this._databaseId,
+      userDataReader,
+      (stages: Stage[]) => {
+        return new RealtimePipeline(
+          this,
+          newUserDataReader(this),
+          new ExpUserDataWriter(this),
+          stages
+        );
+      }
+    );
+  };
+
+/**
+ * @internal
+ * @private
+ */
+export function _onRealtimePipelineSnapshot(
+  pipeline: RealtimePipeline,
+  observer: {
+    next?: (snapshot: RealtimePipelineSnapshot) => void;
+    error?: (error: FirestoreError) => void;
+    complete?: () => void;
+  }
+): Unsubscribe;
+/**
+ * @internal
+ * @private
+ */
+export function _onRealtimePipelineSnapshot(
+  pipeline: RealtimePipeline,
+  options: SnapshotListenOptions,
+  observer: {
+    next?: (snapshot: RealtimePipelineSnapshot) => void;
+    error?: (error: FirestoreError) => void;
+    complete?: () => void;
+  }
+): Unsubscribe;
+/**
+ * @internal
+ * @private
+ */
+export function _onRealtimePipelineSnapshot(
+  pipeline: RealtimePipeline,
+  onNext: (snapshot: RealtimePipelineSnapshot) => void,
+  onError?: (error: FirestoreError) => void,
+  onComplete?: () => void
+): Unsubscribe;
+/**
+ * @internal
+ * @private
+ */
+export function _onRealtimePipelineSnapshot(
+  pipeline: RealtimePipeline,
+  options: SnapshotListenOptions,
+  onNext: (snapshot: RealtimePipelineSnapshot) => void,
+  onError?: (error: FirestoreError) => void,
+  onComplete?: () => void
+): Unsubscribe;
+export function _onRealtimePipelineSnapshot(
+  pipeline: RealtimePipeline,
+  ...args: unknown[]
+): Unsubscribe {
+  let options: SnapshotListenOptions = {
+    includeMetadataChanges: false,
+    source: 'default'
+  };
+  let currArg = 0;
+  if (typeof args[currArg] === 'object' && !isPartialObserver(args[currArg])) {
+    options = args[currArg] as SnapshotListenOptions;
+    currArg++;
+  }
+
+  const internalOptions = {
+    includeMetadataChanges: options.includeMetadataChanges,
+    source: options.source as ListenerDataSource
+  };
+
+  let userObserver: PartialObserver<RealtimePipelineSnapshot>;
+  if (isPartialObserver(args[currArg])) {
+    userObserver = args[currArg] as PartialObserver<RealtimePipelineSnapshot>;
+  } else {
+    userObserver = {
+      next: args[currArg] as NextFn<RealtimePipelineSnapshot>,
+      error: args[currArg + 1] as ErrorFn,
+      complete: args[currArg + 2] as CompleteFn
+    };
+  }
+
+  const client = ensureFirestoreConfigured(pipeline._db as Firestore);
+  const observer = {
+    next: (snapshot: ViewSnapshot) => {
+      if (userObserver.next) {
+        userObserver.next(new RealtimePipelineSnapshot(pipeline, snapshot));
+      }
+    },
+    error: userObserver.error,
+    complete: userObserver.complete
+  };
+
+  return firestoreClientListen(
+    client,
+    toCorePipeline(pipeline),
+    internalOptions, // Pass parsed options here
+    observer
+  );
+}

packages/firestore/src/api/reference_impl.ts

@@ -35,7 +35,8 @@ import {
   firestoreClientListen,
   firestoreClientWrite
 } from '../core/firestore_client';
-import { newQueryForPath, Query as InternalQuery } from '../core/query';
+import { QueryOrPipeline, toCorePipeline } from '../core/pipeline-util';
+import { Query as InternalQuery, newQueryForPath } from '../core/query';
 import { ViewSnapshot } from '../core/view_snapshot';
 import { FieldPath } from '../lite-api/field_path';
 import { validateHasExplicitOrderByForLimitToLast } from '../lite-api/query';
@@ -69,9 +70,12 @@ import {
   DocumentSnapshot,
   FirestoreDataConverter,
   QuerySnapshot,
+  RealtimePipelineSnapshot,
   SnapshotMetadata
 } from './snapshot';
 import { ExpUserDataWriter } from './user_data_writer';
+import { RealtimePipeline } from './realtime_pipeline';
+import { CorePipeline } from '../core/pipeline';
 
 /**
  * An options object that can be passed to {@link (onSnapshot:1)} and {@link
@@ -197,6 +201,10 @@ export function getDocFromServer<
  *
  * @returns A `Promise` that will be resolved with the results of the query.
  */
+export function getDocs<AppModelType, DbModelType extends DocumentData>(
+  query: Query<AppModelType, DbModelType>
+): Promise<QuerySnapshot<AppModelType, DbModelType>>;
+
 export function getDocs<AppModelType, DbModelType extends DocumentData>(
   query: Query<AppModelType, DbModelType>
 ): Promise<QuerySnapshot<AppModelType, DbModelType>> {
@@ -214,7 +222,7 @@ export function getDocs<AppModelType, DbModelType extends DocumentData>(
       new QuerySnapshot<AppModelType, DbModelType>(
         firestore,
         userDataWriter,
-        query,
+        query as Query<AppModelType, DbModelType>,
         snapshot
       )
   );
@@ -642,6 +650,7 @@ export function onSnapshot<AppModelType, DbModelType extends DocumentData>(
   onError?: (error: FirestoreError) => void,
   onCompletion?: () => void
 ): Unsubscribe;
+
 export function onSnapshot<AppModelType, DbModelType extends DocumentData>(
   reference:
     | Query<AppModelType, DbModelType>
@@ -1061,6 +1070,94 @@ export function onSnapshotResume<
   }
 }
 
+export function onPipelineSnapshot(
+  query: RealtimePipeline,
+  observer: {
+    next?: (snapshot: RealtimePipelineSnapshot) => void;
+    error?: (error: FirestoreError) => void;
+    complete?: () => void;
+  }
+): Unsubscribe;
+export function onPipelineSnapshot(
+  query: RealtimePipeline,
+  options: SnapshotListenOptions,
+  observer: {
+    next?: (snapshot: RealtimePipelineSnapshot) => void;
+    error?: (error: FirestoreError) => void;
+    complete?: () => void;
+  }
+): Unsubscribe;
+export function onPipelineSnapshot(
+  query: RealtimePipeline,
+  onNext: (snapshot: RealtimePipelineSnapshot) => void,
+  onError?: (error: FirestoreError) => void,
+  onCompletion?: () => void
+): Unsubscribe;
+export function onPipelineSnapshot(
+  query: RealtimePipeline,
+  options: SnapshotListenOptions,
+  onNext: (snapshot: RealtimePipelineSnapshot) => void,
+  onError?: (error: FirestoreError) => void,
+  onCompletion?: () => void
+): Unsubscribe;
+export function onPipelineSnapshot(
+  reference: RealtimePipeline,
+  ...args: unknown[]
+): Unsubscribe {
+  reference = getModularInstance(reference);
+
+  let options: SnapshotListenOptions = {
+    includeMetadataChanges: false,
+    source: 'default'
+  };
+  let currArg = 0;
+  if (typeof args[currArg] === 'object' && !isPartialObserver(args[currArg])) {
+    options = args[currArg] as SnapshotListenOptions;
+    currArg++;
+  }
+
+  const internalOptions = {
+    includeMetadataChanges: options.includeMetadataChanges,
+    source: options.source as ListenerDataSource
+  };
+
+  if (isPartialObserver(args[currArg])) {
+    const userObserver = args[
+      currArg
+    ] as PartialObserver<RealtimePipelineSnapshot>;
+    args[currArg] = userObserver.next?.bind(userObserver);
+    args[currArg + 1] = userObserver.error?.bind(userObserver);
+    args[currArg + 2] = userObserver.complete?.bind(userObserver);
+  }
+
+  let observer: PartialObserver<ViewSnapshot>;
+  let firestore: Firestore;
+  let internalQuery: CorePipeline;
+
+  // RealtimePipeline
+  firestore = cast(reference._db, Firestore);
+  internalQuery = toCorePipeline(reference);
+  observer = {
+    next: snapshot => {
+      if (args[currArg]) {
+        (args[currArg] as NextFn<RealtimePipelineSnapshot>)(
+          new RealtimePipelineSnapshot(reference as RealtimePipeline, snapshot)
+        );
+      }
+    },
+    error: args[currArg + 1] as ErrorFn,
+    complete: args[currArg + 2] as CompleteFn
+  };
+
+  const client = ensureFirestoreConfigured(firestore);
+  return firestoreClientListen(
+    client,
+    internalQuery,
+    internalOptions,
+    observer
+  );
+}
+
 // TODO(firestorexp): Make sure these overloads are tested via the Firestore
 // integration tests