diff --git a/google/cloud/firestore_v1/base_pipeline.py b/google/cloud/firestore_v1/base_pipeline.py index 153564663..0613722c5 100644 --- a/google/cloud/firestore_v1/base_pipeline.py +++ b/google/cloud/firestore_v1/base_pipeline.py @@ -273,6 +273,71 @@ def find_nearest( stages.FindNearest(field, vector, distance_measure, options) ) + def literals(self, *documents: str | Selectable) -> "_BasePipeline": + """ + 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(Constant.of(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(Constant.of(documents)) + + Output documents: + ```json + [ + {"x": 11}, + {"x": 3} + ] + ``` + + Args: + documents: A `str` or `Selectable` expression. If a `str`, it's + treated as a field path to an array of documents. + If a `Selectable`, it's usually a `Constant` + containing an array of documents (as dictionaries). + Returns: + A new Pipeline object with this stage appended to the stage list. + """ + return self._append(stages.Literals(*documents)) + def replace_with( self, field: Selectable, diff --git a/google/cloud/firestore_v1/pipeline_stages.py b/google/cloud/firestore_v1/pipeline_stages.py index 18aa27044..258a71b82 100644 --- a/google/cloud/firestore_v1/pipeline_stages.py +++ b/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: str | Selectable): + 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/tests/system/pipeline_e2e/general.yaml b/tests/system/pipeline_e2e/general.yaml index 46a10cd4d..1ce427ffe 100644 --- a/tests/system/pipeline_e2e/general.yaml +++ b/tests/system/pipeline_e2e/general.yaml @@ -684,4 +684,24 @@ 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" + assert_results: + - title: "The Hitchhiker's Guide to the Galaxy" + author: "Douglas Adams" + assert_proto: + pipeline: + stages: + - args: + - mapValue: + fields: + author: + stringValue: "Douglas Adams" + title: + stringValue: "The Hitchhiker's Guide to the Galaxy" + name: literals + \ No newline at end of file diff --git a/tests/unit/v1/test_pipeline.py b/tests/unit/v1/test_pipeline.py index 10509cafb..234610bad 100644 --- a/tests/unit/v1/test_pipeline.py +++ b/tests/unit/v1/test_pipeline.py @@ -392,6 +392,7 @@ def test_pipeline_execute_stream_equivalence(): ), ("replace_with", ("name",), stages.ReplaceWith), ("replace_with", (Field.of("n"),), stages.ReplaceWith), + ("literals", (Field.of("a"),), stages.Literals), ("sort", (Field.of("n").descending(),), stages.Sort), ("sort", (Field.of("n").descending(), Field.of("m").ascending()), stages.Sort), ("sample", (10,), stages.Sample), diff --git a/tests/unit/v1/test_pipeline_stages.py b/tests/unit/v1/test_pipeline_stages.py index a2d466f47..58277298b 100644 --- a/tests/unit/v1/test_pipeline_stages.py +++ b/tests/unit/v1/test_pipeline_stages.py @@ -516,6 +516,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 = Constant.of({"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) + repr_str = repr(instance) + assert repr_str == "Literals(documents=(Constant.of({'a': 1}),))" + + def test_to_pb(self): + val1 = Constant.of({"a": 1}) + val2 = Constant.of({"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)