allow passing string or expression into search

Author: daniel-sanche

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

@@ -394,7 +394,7 @@ def sort(self, *orders: stages.Ordering) -> "_BasePipeline":
         """
         return self._append(stages.Sort(*orders))
 
-    def search(self, options: stages.SearchOptions) -> "_BasePipeline":
+    def search(self, query_or_options: str | BooleanExpression | stages.SearchOptions) -> "_BasePipeline":
         """
         Adds a search stage to the pipeline.
 
@@ -403,7 +403,7 @@ def search(self, options: stages.SearchOptions) -> "_BasePipeline":
         Example:
             >>> from google.cloud.firestore_v1.pipeline_stages import SearchOptions
             >>> from google.cloud.firestore_v1.pipeline_expressions import And, DocumentMatches, Field, GeoPoint
-            >>> # Search for restaurants matching "waffles" within 1000m of a location
+            >>> # Search for restaurants matching either "waffles" or "pancakes" near a location
             >>> pipeline = client.pipeline().collection("restaurants").search(
             ...     SearchOptions(
             ...         query=And(
@@ -415,12 +415,13 @@ def search(self, options: stages.SearchOptions) -> "_BasePipeline":
             ... )
 
         Args:
-            options: A SearchOptions instance configuring the search.
+            options: Either a string or expression representing the search query, or
+                A `SearchOptions` instance configuring the search.
 
         Returns:
             A new Pipeline object with this stage appended to the stage list
         """
-        return self._append(stages.Search(options))
+        return self._append(stages.Search(query_or_options))
 
     def sample(self, limit_or_options: int | stages.SampleOptions) -> "_BasePipeline":
         """

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

@@ -500,8 +500,12 @@ def _pb_args(self):
 class Search(Stage):
     """Search stage."""
 
-    def __init__(self, options: SearchOptions):
+    def __init__(self, query_or_options: str | BooleanExpression | SearchOptions):
         super().__init__("search")
+        if isinstance(query_or_options, SearchOptions):
+            options = query_or_options
+        else:
+            options = SearchOptions(query=query_or_options)
         self.options = options
 
     def _pb_args(self) -> list[Value]:

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

@@ -403,6 +403,8 @@ def test_pipeline_execute_stream_equivalence():
         ("replace_with", (Field.of("n"),), stages.ReplaceWith),
         ("sort", (Field.of("n").descending(),), stages.Sort),
         ("sort", (Field.of("n").descending(), Field.of("m").ascending()), stages.Sort),
+        ("search", ("my query",), stages.Search),
+        ("search", (stages.SearchOptions(query="my query"),), stages.Search),
         ("sample", (10,), stages.Sample),
         ("sample", (stages.SampleOptions.doc_limit(10),), stages.Sample),
         ("union", (_make_pipeline(),), stages.Union),

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

@@ -827,6 +827,21 @@ def test_search_string_query_wrapping(self):
         assert options.query.name == "document_matches"
         assert options.query.params[0].value == "science"
 
+    def test_search_with_string(self):
+        stage = stages.Search("technology")
+        assert isinstance(stage.options, stages.SearchOptions)
+        assert stage.options.query.name == "document_matches"
+        assert stage.options.query.params[0].value == "technology"
+        pb_opts = stage._pb_options()
+        assert "query" in pb_opts
+
+    def test_search_with_boolean_expression(self):
+        expr = DocumentMatches("tech")
+        stage = stages.Search(expr)
+        assert isinstance(stage.options, stages.SearchOptions)
+        assert stage.options.query is expr
+        pb_opts = stage._pb_options()
+        assert "query" in pb_opts
 
 class TestSelect:
     def _make_one(self, *args, **kwargs):