feat(firestore): literals pipeline stage (#16028)

Succeeding googleapis/python-firestore#1170 for
the monorepo migration.

Author: Linchin

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

@@ -19,7 +19,7 @@
 
 from __future__ importannotations
 
-fromtypingimportTYPE_CHECKING, Generic, TypeVar
+fromtypingimportGeneric, TypeVar, TYPE_CHECKING
 
 fromgoogle.cloud.firestore_v1importpipeline_stagesasstages
 fromgoogle.cloud.firestore_v1._helpersimportDOCUMENT_PATH_DELIMITER
@@ -32,6 +32,7 @@
 fromgoogle.cloud.firestore_v1.base_documentimportBaseDocumentReference
 fromgoogle.cloud.firestore_v1.base_queryimportBaseQuery
 fromgoogle.cloud.firestore_v1.clientimportClient
+fromgoogle.cloud.firestore_v1.pipeline_expressionsimportCONSTANT_TYPE, Expression
 
 
 PipelineType=TypeVar("PipelineType", bound=_BasePipeline)
@@ -108,3 +109,65 @@ def documents(self, *docs: "BaseDocumentReference") -> PipelineType:
  a new pipeline instance targeting the specified documents
  """
 returnself._create_pipeline(stages.Documents.of(*docs))
+
+defliterals(
+self, *documents: dict[str, Expression|CONSTANT_TYPE]
+ ) ->PipelineType:
+"""
+ Returns documents from a fixed set of predefined document objects.
+
+ 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`
+ of values of `Expression` or `CONSTANT_TYPE` types.
+ Returns:
+ A new Pipeline object with this stage appended to the stage list.
+ """
+returnself._create_pipeline(stages.Literals(*documents))

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

@@ -21,14 +21,15 @@
 
 fromabcimportABC, abstractmethod
 fromenumimportEnum
-fromtypingimportTYPE_CHECKING, Optional, Sequence
+fromtypingimportTYPE_CHECKING, Any, Mapping, Optional, Sequence
 
 fromgoogle.cloud.firestore_v1._helpersimportencode_value
 fromgoogle.cloud.firestore_v1.base_vector_queryimportDistanceMeasure
 fromgoogle.cloud.firestore_v1.pipeline_expressionsimport (
 AggregateFunction,
 AliasedExpression,
 BooleanExpression,
+CONSTANT_TYPE,
 Expression,
 Field,
 Ordering,
@@ -342,6 +343,26 @@ def _pb_args(self):
 return [Value(integer_value=self.limit)]
 
 
+classLiterals(Stage):
+"""Returns documents from a fixed set of predefined document objects."""
+
+def__init__(self, *documents: dict[str, Expression|CONSTANT_TYPE]):
+super().__init__("literals")
+self.documents: tuple[Mapping[str, Any], ...] =documents
+
+def_pb_args(self):
+args= []
+fordocinself.documents:
+encoded_doc= {}
+fork, vindoc.items():
+ifhasattr(v, "_to_pb"):
+encoded_doc[k] =v._to_pb()
+else:
+encoded_doc[k] =encode_value(v)
+args.append(Value(map_value={"fields": encoded_doc}))
+returnargs
+
+
 classOffset(Stage):
 """Skips a specified number of documents."""
 

packages/google-cloud-firestore/tests/system/pipeline_e2e/general.yaml

@@ -684,4 +684,63 @@ tests:
  - args:
  - fieldReferenceValue: awards
  - stringValue: full_replace
-name: replace_with
+name: replace_with
+ - description: literals
+pipeline:
+ - Literals:
+ - title: "The Hitchhiker's Guide to the Galaxy"
+author: "Douglas Adams"
+ - 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
+ - description: literals_with_expression_input
+pipeline:
+ - Literals:
+ - res:
+FunctionExpression.string_concat:
+ - Constant: "A"
+ - Constant: "B"
+ - Select:
+ - res
+assert_results:
+ - res: "AB"
+assert_proto:
+pipeline:
+stages:
+ - args:
+ - mapValue:
+fields:
+res:
+functionValue:
+args:
+ - stringValue: "A"
+ - stringValue: "B"
+name: string_concat
+name: literals
+ - args:
+ - mapValue:
+fields:
+res:
+fieldReferenceValue: res
+name: select

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

@@ -110,6 +110,17 @@ def test_documents(self):
 assertfirst_stage.paths[1] =="/a/2"
 assertfirst_stage.paths[2] =="/a/3"
 
+deftest_literals(self):
+fromgoogle.cloud.firestore_v1.pipeline_expressionsimportField
+
+instance=self._make_client().pipeline()
+documents= ({"field": Field.of("a")}, {"name": "joe"})
+ppl=instance.literals(*documents)
+assertisinstance(ppl, self._expected_pipeline_type)
+assertlen(ppl.stages) ==1
+first_stage=ppl.stages[0]
+assertisinstance(first_stage, stages.Literals)
+
 
 classTestPipelineSourceWithAsyncClient(TestPipelineSource):
 """

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

@@ -517,6 +517,112 @@ def test_to_pb(self):
 assertlen(result.options) ==0
 
 
+classTestLiterals:
+def_make_one(self, *args, **kwargs):
+returnstages.Literals(*args, **kwargs)
+
+deftest_ctor(self):
+val1= {"a": Field.of("a")}
+val2= {"b": 2}
+instance=self._make_one(val1, val2)
+assertinstance.documents== (val1, val2)
+assertinstance.name=="literals"
+
+deftest_ctor_extended_types(self):
+importdatetime
+fromgoogle.cloud.firestore_v1._helpersimportGeoPoint
+fromgoogle.cloud.firestore_v1.vectorimportVector
+
+doc= {
+"a": 1,
+"b": "string",
+"c": 3.14,
+"d": True,
+"e": None,
+"f": b"bytes",
+"g": datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc),
+"h": GeoPoint(1.0, 2.0),
+"i": Vector([1.0, 2.0]),
+ }
+instance=self._make_one(doc)
+assertinstance.documents== (doc,)
+assertinstance.name=="literals"
+
+deftest_ctor_w_expressions(self):
+fromgoogle.cloud.firestore_v1.pipeline_expressionsimportFunctionExpression
+
+expr=FunctionExpression("string_concat", [Constant("A"), Constant("B")])
+doc= {"res": expr}
+instance=self._make_one(doc)
+assertinstance.documents== (doc,)
+assertinstance.name=="literals"
+
+deftest_repr(self):
+fromgoogle.cloud.firestore_v1.pipeline_expressionsimportField
+
+val1= {"a": Field.of("a")}
+instance=self._make_one(val1, {"b": 2})
+repr_str=repr(instance)
+assertrepr_str=="Literals(documents=({'a': Field.of('a')}, {'b': 2}))"
+
+deftest_to_pb_constant_types(self):
+importdatetime
+fromgoogle.cloud.firestore_v1._helpersimportGeoPoint
+fromgoogle.cloud.firestore_v1.vectorimportVector
+
+doc= {
+"a": 1,
+"b": "string",
+"c": 3.14,
+"d": True,
+"e": None,
+"f": b"bytes",
+"g": datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc),
+"h": GeoPoint(1.0, 2.0),
+"i": Vector([1.0, 2.0]),
+ }
+instance=self._make_one(doc)
+result=instance._to_pb()
+assertresult.name=="literals"
+assertlen(result.args) ==1
+
+fields=result.args[0].map_value.fields
+assertfields["a"].integer_value==1
+assertfields["b"].string_value=="string"
+assertfields["c"].double_value==3.14
+assertfields["d"].boolean_valueisTrue
+assertfields["e"].null_value==0
+assertfields["f"].bytes_value==b"bytes"
+assertfields["g"].timestamp_value==datetime.datetime(
+2025, 1, 1, tzinfo=datetime.timezone.utc
+ )
+assertfields["h"].geo_point_value.latitude==1.0
+assertfields["h"].geo_point_value.longitude==2.0
+assert (
+fields["i"].map_value.fields["value"].array_value.values[0].double_value
+==1.0
+ )
+assert (
+fields["i"].map_value.fields["value"].array_value.values[1].double_value
+==2.0
+ )
+
+deftest_to_pb_w_expression(self):
+fromgoogle.cloud.firestore_v1.pipeline_expressionsimportFunctionExpression
+
+expr=FunctionExpression("string_concat", [Constant("A"), Constant("B")])
+doc= {"res": expr}
+instance=self._make_one(doc)
+result=instance._to_pb()
+assertresult.name=="literals"
+assertlen(result.args) ==1
+
+fields=result.args[0].map_value.fields
+assertfields["res"].function_value.name=="string_concat"
+assertfields["res"].function_value.args[0].string_value=="A"
+assertfields["res"].function_value.args[1].string_value=="B"
+
+
 classTestOffset:
 def_make_one(self, *args, **kwargs):
 returnstages.Offset(*args, **kwargs)