diff --git a/packages/google-cloud-firestore/google/cloud/firestore_v1/base_pipeline.py b/packages/google-cloud-firestore/google/cloud/firestore_v1/base_pipeline.py index fac7f8bc4bce..fb7bbdbbc81e 100644 --- a/packages/google-cloud-firestore/google/cloud/firestore_v1/base_pipeline.py +++ b/packages/google-cloud-firestore/google/cloud/firestore_v1/base_pipeline.py @@ -275,6 +275,7 @@ def find_nearest( stages.FindNearest(field, vector, distance_measure, options) ) + def replace_with( self, field: Selectable, diff --git a/packages/google-cloud-firestore/google/cloud/firestore_v1/pipeline_source.py b/packages/google-cloud-firestore/google/cloud/firestore_v1/pipeline_source.py index 0a7c6e7fbdab..d647f7985c9e 100644 --- a/packages/google-cloud-firestore/google/cloud/firestore_v1/pipeline_source.py +++ b/packages/google-cloud-firestore/google/cloud/firestore_v1/pipeline_source.py @@ -32,6 +32,7 @@ from google.cloud.firestore_v1.base_document import BaseDocumentReference from google.cloud.firestore_v1.base_query import BaseQuery from google.cloud.firestore_v1.client import Client + from google.cloud.firestore_v1.pipeline_expressions import Expression PipelineType = TypeVar("PipelineType", bound=_BasePipeline) @@ -108,3 +109,66 @@ def documents(self, *docs: "BaseDocumentReference") -> PipelineType: a new pipeline instance targeting the specified documents """ return self._create_pipeline(stages.Documents.of(*docs)) + + def literals(self, *documents: "Expression" | dict) -> PipelineType: + """ + Returns documents from a fixed set of predefined document objects. + + This stage is commonly used for testing other stages in isolation, + though it can also be used as inputs to join conditions. + + Example: + >>> from google.cloud.firestore_v1.pipeline_expressions import Constant + >>> documents = [ + ... {"name": "joe", "age": 10}, + ... {"name": "bob", "age": 30}, + ... {"name": "alice", "age": 40} + ... ] + >>> pipeline = client.pipeline() + ... .literals(documents) + ... .where(field("age").lessThan(35)) + + Output documents: + ```json + [ + {"name": "joe", "age": 10}, + {"name": "bob", "age": 30} + ] + ``` + + Behavior: + The `literals(...)` stage can only be used as the first stage in a pipeline (or + sub-pipeline). The order of documents returned from the `literals` matches the + order in which they are defined. + + While literal values are the most common, it is also possible to pass in + expressions, which will be evaluated and returned, making it possible to test + out different query / expression behavior without first needing to create some + test data. + + For example, the following shows how to quickly test out the `length(...)` + function on some constant test sets: + + Example: + >>> from google.cloud.firestore_v1.pipeline_expressions import Constant + >>> documents = [ + ... {"x": Constant.of("foo-bar-baz").char_length()}, + ... {"x": Constant.of("bar").char_length()} + ... ] + >>> pipeline = client.pipeline().literals(documents) + + Output documents: + ```json + [ + {"x": 11}, + {"x": 3} + ] + ``` + + Args: + *documents: One or more documents to be returned by this stage. Each can be a `dict` + or an `Expression`. + Returns: + A new Pipeline object with this stage appended to the stage list. + """ + return self._create_pipeline(stages.Literals(*documents)) diff --git a/packages/google-cloud-firestore/google/cloud/firestore_v1/pipeline_stages.py b/packages/google-cloud-firestore/google/cloud/firestore_v1/pipeline_stages.py index b00d923c673c..30d2fd8e936a 100644 --- a/packages/google-cloud-firestore/google/cloud/firestore_v1/pipeline_stages.py +++ b/packages/google-cloud-firestore/google/cloud/firestore_v1/pipeline_stages.py @@ -342,6 +342,23 @@ def _pb_args(self): return [Value(integer_value=self.limit)] +class Literals(Stage): + """Returns documents from a fixed set of predefined document objects.""" + + def __init__(self, *documents: Expression | dict): + super().__init__("literals") + self.documents = documents + + def _pb_args(self): + args = [] + for doc in self.documents: + if hasattr(doc, "_to_pb"): + args.append(doc._to_pb()) + else: + args.append(encode_value(doc)) + return args + + class Offset(Stage): """Skips a specified number of documents.""" diff --git a/packages/google-cloud-firestore/tests/system/pipeline_e2e/general.yaml b/packages/google-cloud-firestore/tests/system/pipeline_e2e/general.yaml index 46a10cd4d1af..d4aa80ba8d7b 100644 --- a/packages/google-cloud-firestore/tests/system/pipeline_e2e/general.yaml +++ b/packages/google-cloud-firestore/tests/system/pipeline_e2e/general.yaml @@ -684,4 +684,35 @@ tests: - args: - fieldReferenceValue: awards - stringValue: full_replace - name: replace_with \ No newline at end of file + name: replace_with + - description: literals + pipeline: + - Literals: + - title: "The Hitchhiker's Guide to the Galaxy" + author: "Douglas Adams" + - Constant: + value: + genre: "Science Fiction" + year: 1979 + assert_results: + - title: "The Hitchhiker's Guide to the Galaxy" + author: "Douglas Adams" + - genre: "Science Fiction" + year: 1979 + assert_proto: + pipeline: + stages: + - args: + - mapValue: + fields: + author: + stringValue: "Douglas Adams" + title: + stringValue: "The Hitchhiker's Guide to the Galaxy" + - mapValue: + fields: + genre: + stringValue: "Science Fiction" + year: + integerValue: '1979' + name: literals \ No newline at end of file diff --git a/packages/google-cloud-firestore/tests/unit/v1/test_pipeline_source.py b/packages/google-cloud-firestore/tests/unit/v1/test_pipeline_source.py index de31d47f70eb..94da65a61690 100644 --- a/packages/google-cloud-firestore/tests/unit/v1/test_pipeline_source.py +++ b/packages/google-cloud-firestore/tests/unit/v1/test_pipeline_source.py @@ -110,6 +110,17 @@ def test_documents(self): assert first_stage.paths[1] == "/a/2" assert first_stage.paths[2] == "/a/3" + def test_literals(self): + from google.cloud.firestore_v1.pipeline_expressions import Field + + instance = self._make_client().pipeline() + documents = (Field.of("a"), {"name": "joe"}) + ppl = instance.literals(*documents) + assert isinstance(ppl, self._expected_pipeline_type) + assert len(ppl.stages) == 1 + first_stage = ppl.stages[0] + assert isinstance(first_stage, stages.Literals) + class TestPipelineSourceWithAsyncClient(TestPipelineSource): """ diff --git a/packages/google-cloud-firestore/tests/unit/v1/test_pipeline_stages.py b/packages/google-cloud-firestore/tests/unit/v1/test_pipeline_stages.py index b32a6e5d3f13..82fe3e1881b2 100644 --- a/packages/google-cloud-firestore/tests/unit/v1/test_pipeline_stages.py +++ b/packages/google-cloud-firestore/tests/unit/v1/test_pipeline_stages.py @@ -517,6 +517,35 @@ def test_to_pb(self): assert len(result.options) == 0 +class TestLiterals: + def _make_one(self, *args, **kwargs): + return stages.Literals(*args, **kwargs) + + def test_ctor(self): + val1 = Constant.of({"a": 1}) + val2 = {"b": 2} + instance = self._make_one(val1, val2) + assert instance.documents == (val1, val2) + assert instance.name == "literals" + + def test_repr(self): + val1 = Constant.of({"a": 1}) + instance = self._make_one(val1, {"b": 2}) + repr_str = repr(instance) + assert repr_str == "Literals(documents=(Constant.of({'a': 1}), {'b': 2}))" + + def test_to_pb(self): + val1 = Constant.of({"a": 1}) + val2 = {"b": 2} + instance = self._make_one(val1, val2) + result = instance._to_pb() + assert result.name == "literals" + assert len(result.args) == 2 + assert result.args[0].map_value.fields["a"].integer_value == 1 + assert result.args[1].map_value.fields["b"].integer_value == 2 + assert len(result.options) == 0 + + class TestOffset: def _make_one(self, *args, **kwargs): return stages.Offset(*args, **kwargs)