Pipeline tests part 1 (#14885)

Co-authored-by: wu-hui <wu.hui.github@gmail.com>
Co-authored-by: Nick Cooke <36927374+ncooke3@users.noreply.github.com>

Author: 3 people

Firestore/Source/API/FIRPipelineBridge.mm

@@ -20,6 +20,7 @@
 
 #include <memory>
 
+#import "Firestore/Source/API/FIRCollectionReference+Internal.h"
 #import "Firestore/Source/API/FIRDocumentReference+Internal.h"
 #import "Firestore/Source/API/FIRFieldPath+Internal.h"
 #import "Firestore/Source/API/FIRFirestore+Internal.h"
@@ -39,6 +40,7 @@
 #include "Firestore/core/src/api/pipeline_result.h"
 #include "Firestore/core/src/api/pipeline_snapshot.h"
 #include "Firestore/core/src/api/stages.h"
+#include "Firestore/core/src/util/comparison.h"
 #include "Firestore/core/src/util/error_apple.h"
 #include "Firestore/core/src/util/status.h"
 #include "Firestore/core/src/util/string_apple.h"
@@ -57,12 +59,12 @@
 using firebase::firestore::api::Field;
 using firebase::firestore::api::FindNearestStage;
 using firebase::firestore::api::FunctionExpr;
-using firebase::firestore::api::GenericStage;
 using firebase::firestore::api::LimitStage;
 using firebase::firestore::api::MakeFIRTimestamp;
 using firebase::firestore::api::OffsetStage;
 using firebase::firestore::api::Ordering;
 using firebase::firestore::api::Pipeline;
+using firebase::firestore::api::RawStage;
 using firebase::firestore::api::RemoveFieldsStage;
 using firebase::firestore::api::ReplaceWith;
 using firebase::firestore::api::Sample;
@@ -73,13 +75,21 @@
 using firebase::firestore::api::Where;
 using firebase::firestore::model::FieldPath;
 using firebase::firestore::nanopb::SharedMessage;
+using firebase::firestore::util::ComparisonResult;
 using firebase::firestore::util::MakeCallback;
 using firebase::firestore::util::MakeNSString;
 using firebase::firestore::util::MakeString;
 using firebase::firestore::util::ThrowInvalidArgument;
 
 NS_ASSUME_NONNULL_BEGIN
 
+inline std::string EnsureLeadingSlash(const std::string &path) {
+  if (!path.empty() && path[0] == '/') {
+    return path;
+  }
+  return "/" + path;
+}
+
 @implementation FIRExprBridge
 @end
 
@@ -216,10 +226,19 @@ @implementation FIRCollectionSourceStageBridge {
   std::shared_ptr<CollectionSource> collection_source;
 }
 
-- (id)initWithPath:(NSString *)path {
+- (id)initWithRef:(FIRCollectionReference *)ref firestore:(FIRFirestore *)db {
   self = [super init];
   if (self) {
-    collection_source = std::make_shared<CollectionSource>(MakeString(path));
+    if (ref.firestore.databaseID.CompareTo(db.databaseID) != ComparisonResult::Same) {
+      ThrowInvalidArgument(
+          "Invalid CollectionReference. The project ID (\"%s\") or the database (\"%s\") does not "
+          "match "
+          "the project ID (\"%s\") and database (\"%s\") of the target database of this Pipeline.",
+          ref.firestore.databaseID.project_id(), ref.firestore.databaseID.database_id(),
+          db.databaseID.project_id(), db.databaseID.project_id());
+    }
+    collection_source =
+        std::make_shared<CollectionSource>(EnsureLeadingSlash(MakeString(ref.path)));
   }
   return self;
 }
@@ -270,12 +289,21 @@ @implementation FIRDocumentsSourceStageBridge {
   std::shared_ptr<DocumentsSource> cpp_document_source;
 }
 
-- (id)initWithDocuments:(NSArray<NSString *> *)documents {
+- (id)initWithDocuments:(NSArray<FIRDocumentReference *> *)documents firestore:(FIRFirestore *)db {
   self = [super init];
   if (self) {
     std::vector<std::string> cpp_documents;
-    for (NSString *doc in documents) {
-      cpp_documents.push_back(MakeString(doc));
+    for (FIRDocumentReference *doc in documents) {
+      if (doc.firestore.databaseID.CompareTo(db.databaseID) != ComparisonResult::Same) {
+        ThrowInvalidArgument("Invalid DocumentReference. The project ID (\"%s\") or the database "
+                             "(\"%s\") does not match "
+                             "the project ID (\"%s\") and database (\"%s\") of the target database "
+                             "of this Pipeline.",
+                             doc.firestore.databaseID.project_id(),
+                             doc.firestore.databaseID.database_id(), db.databaseID.project_id(),
+                             db.databaseID.project_id());
+      }
+      cpp_documents.push_back(EnsureLeadingSlash(MakeString(doc.path)));
     }
     cpp_document_source = std::make_shared<DocumentsSource>(std::move(cpp_documents));
   }
@@ -754,12 +782,12 @@ - (id)initWithField:(FIRExprBridge *)field indexField:(NSString *_Nullable)index
 
 @end
 
-@implementation FIRGenericStageBridge {
+@implementation FIRRawStageBridge {
   NSString *_name;
   NSArray<FIRExprBridge *> *_params;
   NSDictionary<NSString *, FIRExprBridge *> *_Nullable _options;
   Boolean isUserDataRead;
-  std::shared_ptr<GenericStage> cpp_generic_stage;
+  std::shared_ptr<RawStage> cpp_generic_stage;
 }
 
 - (id)initWithName:(NSString *)name
@@ -787,8 +815,8 @@ - (id)initWithName:(NSString *)name
         cpp_options[MakeString(key)] = [_options[key] cppExprWithReader:reader];
       }
     }
-    cpp_generic_stage = std::make_shared<GenericStage>(MakeString(_name), std::move(cpp_params),
-                                                       std::move(cpp_options));
+    cpp_generic_stage = std::make_shared<RawStage>(MakeString(_name), std::move(cpp_params),
+                                                   std::move(cpp_options));
   }
 
   isUserDataRead = YES;

Firestore/Source/Public/FirebaseFirestore/FIRPipelineBridge.h

@@ -70,7 +70,7 @@ NS_SWIFT_SENDABLE
 NS_SWIFT_NAME(CollectionSourceStageBridge)
 @interface FIRCollectionSourceStageBridge : FIRStageBridge
 
-- (id)initWithPath:(NSString *)path;
+- (id)initWithRef:(FIRCollectionReference *)ref firestore:(FIRFirestore *)db;
 
 @end
 
@@ -94,7 +94,7 @@ NS_SWIFT_SENDABLE
 NS_SWIFT_NAME(DocumentsSourceStageBridge)
 @interface FIRDocumentsSourceStageBridge : FIRStageBridge
 
-- (id)initWithDocuments:(NSArray<NSString *> *)documents;
+- (id)initWithDocuments:(NSArray<FIRDocumentReference *> *)documents firestore:(FIRFirestore *)db;
 
 @end
 
@@ -195,8 +195,8 @@ NS_SWIFT_NAME(UnnestStageBridge)
 @end
 
 NS_SWIFT_SENDABLE
-NS_SWIFT_NAME(GenericStageBridge)
-@interface FIRGenericStageBridge : FIRStageBridge
+NS_SWIFT_NAME(RawStageBridge)
+@interface FIRRawStageBridge : FIRStageBridge
 - (id)initWithName:(NSString *)name
             params:(NSArray<FIRExprBridge *> *)params
            options:(NSDictionary<NSString *, FIRExprBridge *> *_Nullable)options;

Firestore/Swift/Source/SwiftAPI/Pipeline/Pipeline.swift

@@ -699,7 +699,7 @@ public struct Pipeline: @unchecked Sendable {
   /// ```swift
   /// // let pipeline: Pipeline = ...
   /// // Example: Assuming a hypothetical backend stage "customFilterV2".
-  /// let genericPipeline = pipeline.genericStage(
+  /// let genericPipeline = pipeline.rawStage(
   ///   name: "customFilterV2",
   ///   params: [Field("userScore"), 80], // Ordered parameters.
   ///   options: ["mode": "strict", "logLevel": 2]  // Optional named parameters.
@@ -712,10 +712,10 @@ public struct Pipeline: @unchecked Sendable {
   ///   - params: An array of ordered, `Sendable` parameters for the stage.
   ///   - options: Optional dictionary of named, `Sendable` parameters.
   /// - Returns: A new `Pipeline` object with this stage appended.
-  public func genericStage(name: String, params: [Sendable],
-                           options: [String: Sendable]? = nil) -> Pipeline {
+  public func rawStage(name: String, params: [Sendable],
+                       options: [String: Sendable]? = nil) -> Pipeline {
     return Pipeline(
-      stages: stages + [GenericStage(name: name, params: params, options: options)],
+      stages: stages + [RawStage(name: name, params: params, options: options)],
       db: db
     )
   }

Firestore/Swift/Source/SwiftAPI/Pipeline/PipelineSnapshot.swift

@@ -25,7 +25,7 @@ public struct PipelineSnapshot: Sendable {
   public let pipeline: Pipeline
 
   /// An array of all the results in the `PipelineSnapshot`.
-  let results_cache: [PipelineResult]
+  public let results: [PipelineResult]
 
   /// The time at which the pipeline producing this result was executed.
   public let executionTime: Timestamp
@@ -36,10 +36,6 @@ public struct PipelineSnapshot: Sendable {
     self.bridge = bridge
     self.pipeline = pipeline
     executionTime = self.bridge.execution_time
-    results_cache = self.bridge.results.map { PipelineResult($0) }
-  }
-
-  public func results() -> [PipelineResult] {
-    return results_cache
+    results = self.bridge.results.map { PipelineResult($0) }
   }
 }

Firestore/Swift/Source/SwiftAPI/Pipeline/PipelineSource.swift

@@ -21,8 +21,12 @@ public struct PipelineSource: @unchecked Sendable {
   }
 
   public func collection(_ path: String) -> Pipeline {
-    let normalizedPath = path.hasPrefix("/") ? path : "/" + path
-    return Pipeline(stages: [CollectionSource(collection: normalizedPath)], db: db)
+    return Pipeline(stages: [CollectionSource(collection: db.collection(path), db: db)], db: db)
+  }
+
+  public func collection(_ ref: CollectionReference) -> Pipeline {
+    let collectionStage = CollectionSource(collection: ref, db: db)
+    return Pipeline(stages: [collectionStage], db: db)
   }
 
   public func collectionGroup(_ collectionId: String) -> Pipeline {
@@ -37,13 +41,13 @@ public struct PipelineSource: @unchecked Sendable {
   }
 
   public func documents(_ docs: [DocumentReference]) -> Pipeline {
-    let paths = docs.map { $0.path.hasPrefix("/") ? $0.path : "/" + $0.path }
-    return Pipeline(stages: [DocumentsSource(paths: paths)], db: db)
+    return Pipeline(stages: [DocumentsSource(docs: docs, db: db)], db: db)
   }
 
   public func documents(_ paths: [String]) -> Pipeline {
-    let normalizedPaths = paths.map { $0.hasPrefix("/") ? $0 : "/" + $0 }
-    return Pipeline(stages: [DocumentsSource(paths: normalizedPaths)], db: db)
+    let docs = paths.map { db.document($0) }
+    let documentsStage = DocumentsSource(docs: docs, db: db)
+    return Pipeline(stages: [documentsStage], db: db)
   }
 
   public func create(from query: Query) -> Pipeline {

Firestore/Swift/Source/SwiftAPI/Stages.swift

@@ -33,11 +33,13 @@ class CollectionSource: Stage {
   let name: String = "collection"
 
   let bridge: StageBridge
-  private var collection: String
+  private var collection: CollectionReference
+  private let db: Firestore
 
-  init(collection: String) {
+  init(collection: CollectionReference, db: Firestore) {
     self.collection = collection
-    bridge = CollectionSourceStageBridge(path: collection)
+    self.db = db
+    bridge = CollectionSourceStageBridge(ref: collection, firestore: db)
   }
 }
 
@@ -70,12 +72,14 @@ class DatabaseSource: Stage {
 class DocumentsSource: Stage {
   let name: String = "documents"
   let bridge: StageBridge
-  private var references: [String]
+  private var docs: [DocumentReference]
+  private let db: Firestore
 
   // Initialize with an array of String paths
-  init(paths: [String]) {
-    references = paths
-    bridge = DocumentsSourceStageBridge(documents: paths)
+  init(docs: [DocumentReference], db: Firestore) {
+    self.docs = docs
+    self.db = db
+    bridge = DocumentsSourceStageBridge(documents: docs, firestore: db)
   }
 }
 
@@ -337,7 +341,7 @@ class Unnest: Stage {
 }
 
 @available(iOS 13, tvOS 13, macOS 10.15, macCatalyst 13, watchOS 7, *)
-class GenericStage: Stage {
+class RawStage: Stage {
   let name: String
   let bridge: StageBridge
   private var params: [Sendable]
@@ -349,6 +353,6 @@ class GenericStage: Stage {
     self.options = options
     let bridgeParams = params.map { Helper.sendableToExpr($0).toBridge() }
     let bridgeOptions = options?.mapValues { Helper.sendableToExpr($0).toBridge() }
-    bridge = GenericStageBridge(name: name, params: bridgeParams, options: bridgeOptions)
+    bridge = RawStageBridge(name: name, params: bridgeParams, options: bridgeOptions)
   }
 }

Firestore/Swift/Tests/Integration/PipelineApiTests.swift

@@ -263,22 +263,22 @@ final class PipelineTests: FSTIntegrationTestCase {
     // ... }
   }
 
-  func testGenericStage() async throws {
+  func testRawStage() async throws {
     // Assume we don't have a built-in "where" stage, the customer could still
-    // add this stage by calling genericStage, passing the name of the stage "where",
+    // add this stage by calling rawStage, passing the name of the stage "where",
     // and providing positional argument values.
     _ = db.pipeline().collection("books")
-      .genericStage(name: "where",
-                    params: [Field("published").lt(1900)])
+      .rawStage(name: "where",
+                params: [Field("published").lt(1900)])
       .select("title", "author")
 
     // In cases where the stage also supports named argument values, then these can be
     // provided with a third argument that maps the argument name to value.
     // Note that these named arguments are always optional in the stage definition.
     _ = db.pipeline().collection("books")
-      .genericStage(name: "where",
-                    params: [Field("published").lt(1900)],
-                    options: ["someOptionalParamName": "the argument value for this param"])
+      .rawStage(name: "where",
+                params: [Field("published").lt(1900)],
+                options: ["someOptionalParamName": "the argument value for this param"])
       .select("title", "author")
   }