exposed options in raw_stage

Author: daniel-sanche

packages/google-cloud-firestore/google/cloud/firestore_v1/base_pipeline.py

@@ -459,7 +459,12 @@ def unnest(
         """
         return self._append(stages.Unnest(field, alias, options))
 
-    def raw_stage(self, name: str, *params: Expression) -> "_BasePipeline":
+    def raw_stage(
+        self,
+        name: str,
+        *params: Expression,
+        options: dict[str, Expression | Value] | None = None,
+    ) -> "_BasePipeline":
         """
         Adds a stage to the pipeline by specifying the stage name as an argument. This does not offer any
         type safety on the stage params and requires the caller to know the order (and optionally names)
@@ -477,11 +482,12 @@ def raw_stage(self, name: str, *params: Expression) -> "_BasePipeline":
         Args:
             name: The name of the stage.
             *params: A sequence of `Expression` objects representing the parameters for the stage.
+            options: An optional dictionary of stage options.
 
         Returns:
             A new Pipeline object with this stage appended to the stage list
         """
-        return self._append(stages.RawStage(name, *params))
+        return self._append(stages.RawStage(name, *params, options=options or {}))
 
     def offset(self, offset: int) -> "_BasePipeline":
         """

packages/google-cloud-firestore/tests/unit/v1/test_async_pipeline.py

@@ -448,3 +448,17 @@ def test_async_pipeline_aggregate_with_groups():
     assert isinstance(result_ppl.stages[0], stages.Aggregate)
     assert list(result_ppl.stages[0].groups) == [Field.of("author")]
     assert list(result_ppl.stages[0].accumulators) == [Field.of("title")]
+
+
+def test_async_pipeline_raw_stage_with_options():
+    from google.cloud.firestore_v1.base_vector_query import Field
+    from google.cloud.firestore_v1.pipeline_stages import RawStage
+
+    start_ppl = _make_async_pipeline()
+    result_ppl = start_ppl.raw_stage(
+        "stage_name", Field.of("n"), options={"key": "val"}
+    )
+    assert len(start_ppl.stages) == 0
+    assert len(result_ppl.stages) == 1
+    assert isinstance(result_ppl.stages[0], RawStage)
+    assert result_ppl.stages[0].options == {"key": "val"}

packages/google-cloud-firestore/tests/unit/v1/test_pipeline.py

@@ -437,3 +437,17 @@ def test_pipeline_aggregate_with_groups():
     assert isinstance(result_ppl.stages[0], stages.Aggregate)
     assert list(result_ppl.stages[0].groups) == [Field.of("author")]
     assert list(result_ppl.stages[0].accumulators) == [Field.of("title")]
+
+
+def test_pipeline_raw_stage_with_options():
+    from google.cloud.firestore_v1.base_vector_query import Field
+    from google.cloud.firestore_v1.pipeline_stages import RawStage
+
+    start_ppl = _make_pipeline()
+    result_ppl = start_ppl.raw_stage(
+        "stage_name", Field.of("n"), options={"key": "val"}
+    )
+    assert len(start_ppl.stages) == 0
+    assert len(result_ppl.stages) == 1
+    assert isinstance(result_ppl.stages[0], RawStage)
+    assert result_ppl.stages[0].options == {"key": "val"}