Merge branch 'pipeline-preview' into pipeline_run_kokoro

Author: daniel-sanche

google/cloud/firestore_v1/async_pipeline.py

@@ -13,14 +13,20 @@
 # limitations under the License.
 
 from __future__ import annotations
-from typing import AsyncIterable, TYPE_CHECKING
+from typing import TYPE_CHECKING
 from google.cloud.firestore_v1 import pipeline_stages as stages
 from google.cloud.firestore_v1.base_pipeline import _BasePipeline
+from google.cloud.firestore_v1.pipeline_result import AsyncPipelineStream
+from google.cloud.firestore_v1.pipeline_result import PipelineSnapshot
+from google.cloud.firestore_v1.pipeline_result import PipelineResult
 
 if TYPE_CHECKING:  # pragma: NO COVER
+    import datetime
     from google.cloud.firestore_v1.async_client import AsyncClient
-    from google.cloud.firestore_v1.pipeline_result import PipelineResult
     from google.cloud.firestore_v1.async_transaction import AsyncTransaction
+    from google.cloud.firestore_v1.pipeline_expressions import Constant
+    from google.cloud.firestore_v1.types.document import Value
+    from google.cloud.firestore_v1.query_profile import PipelineExplainOptions
 
 
 class AsyncPipeline(_BasePipeline):
@@ -40,7 +46,7 @@ class AsyncPipeline(_BasePipeline):
         ...                      .collection("books")
         ...                      .where(Field.of("published").gt(1980))
         ...                      .select("title", "author")
-        ...     async for result in pipeline.execute():
+        ...     async for result in pipeline.stream():
         ...         print(result)
 
     Use `client.pipeline()` to create instances of this class.
@@ -58,39 +64,68 @@ def __init__(self, client: AsyncClient, *stages: stages.Stage):
 
     async def execute(
         self,
+        *,
         transaction: "AsyncTransaction" | None = None,
-    ) -> list[PipelineResult]:
+        read_time: datetime.datetime | None = None,
+        explain_options: PipelineExplainOptions | None = None,
+        index_mode: str | None = None,
+        additional_options: dict[str, Value | Constant] = {},
+    ) -> PipelineSnapshot[PipelineResult]:
         """
         Executes this pipeline and returns results as a list
 
         Args:
-            transaction
-                (Optional[:class:`~google.cloud.firestore_v1.transaction.Transaction`]):
+            transaction (Optional[:class:`~google.cloud.firestore_v1.transaction.Transaction`]):
                 An existing transaction that this query will run in.
                 If a ``transaction`` is used and it already has write operations
                 added, this method cannot be used (i.e. read-after-write is not
                 allowed).
+            read_time (Optional[datetime.datetime]): If set, reads documents as they were at the given
+                time. This must be a microsecond precision timestamp within the past one hour, or
+                if Point-in-Time Recovery is enabled, can additionally be a whole minute timestamp
+                within the past 7 days. For the most accurate results, use UTC timezone.
+            explain_options (Optional[:class:`~google.cloud.firestore_v1.query_profile.PipelineExplainOptions`]):
+                Options to enable query profiling for this query. When set,
+                explain_metrics will be available on the returned list.
+            index_mode (Optional[str]): Configures the pipeline to require a certain type of indexes to be present.
+                Firestore will reject the request if there is not appropiate indexes to serve the query.
+            additional_options (Optional[dict[str, Value | Constant]]): Additional options to pass to the query.
+                These options will take precedence over method argument if there is a conflict (e.g. explain_options, index_mode)
         """
-        return [result async for result in self.stream(transaction=transaction)]
+        kwargs = {k: v for k, v in locals().items() if k != "self"}
+        stream = AsyncPipelineStream(PipelineResult, self, **kwargs)
+        results = [result async for result in stream]
+        return PipelineSnapshot(results, stream)
 
-    async def stream(
+    def stream(
         self,
+        *,
+        read_time: datetime.datetime | None = None,
         transaction: "AsyncTransaction" | None = None,
-    ) -> AsyncIterable[PipelineResult]:
+        explain_options: PipelineExplainOptions | None = None,
+        index_mode: str | None = None,
+        additional_options: dict[str, Value | Constant] = {},
+    ) -> AsyncPipelineStream[PipelineResult]:
         """
-        Process this pipeline as a stream, providing results through an Iterable
+        Process this pipeline as a stream, providing results through an AsyncIterable
 
         Args:
-            transaction
-                (Optional[:class:`~google.cloud.firestore_v1.transaction.Transaction`]):
+            transaction (Optional[:class:`~google.cloud.firestore_v1.transaction.Transaction`]):
                 An existing transaction that this query will run in.
                 If a ``transaction`` is used and it already has write operations
                 added, this method cannot be used (i.e. read-after-write is not
                 allowed).
+            read_time (Optional[datetime.datetime]): If set, reads documents as they were at the given
+                time. This must be a microsecond precision timestamp within the past one hour, or
+                if Point-in-Time Recovery is enabled, can additionally be a whole minute timestamp
+                within the past 7 days. For the most accurate results, use UTC timezone.
+            explain_options (Optional[:class:`~google.cloud.firestore_v1.query_profile.PipelineExplainOptions`]):
+                Options to enable query profiling for this query. When set,
+                explain_metrics will be available on the returned generator.
+            index_mode (Optional[str]): Configures the pipeline to require a certain type of indexes to be present.
+                Firestore will reject the request if there is not appropiate indexes to serve the query.
+            additional_options (Optional[dict[str, Value | Constant]]): Additional options to pass to the query.
+                These options will take precedence over method argument if there is a conflict (e.g. explain_options, index_mode)
         """
-        request = self._prep_execute_request(transaction)
-        async for response in await self._client._firestore_api.execute_pipeline(
-            request
-        ):
-            for result in self._execute_response_helper(response):
-                yield result
+        kwargs = {k: v for k, v in locals().items() if k != "self"}
+        return AsyncPipelineStream(PipelineResult, self, **kwargs)

google/cloud/firestore_v1/base_aggregation.py

@@ -48,6 +48,7 @@
     from google.cloud.firestore_v1.stream_generator import (
         StreamGenerator,
     )
+    from google.cloud.firestore_v1.pipeline_source import PipelineSource
 
     import datetime
 
@@ -356,19 +357,20 @@ def stream(
             A generator of the query results.
         """
 
-    def pipeline(self):
+    def _build_pipeline(self, source: "PipelineSource"):
         """
         Convert this query into a Pipeline
 
         Queries containing a `cursor` or `limit_to_last` are not currently supported
 
+        Args:
+            source: the PipelineSource to build the pipeline off of
         Raises:
-            - ValueError: raised if Query wasn't created with an associated client
             - NotImplementedError: raised if the query contains a `cursor` or `limit_to_last`
         Returns:
             a Pipeline representing the query
         """
         # use autoindexer to keep track of which field number to use for un-aliased fields
         autoindexer = itertools.count(start=1)
         exprs = [a._to_pipeline_expr(autoindexer) for a in self._aggregations]
-        return self._nested_query.pipeline().aggregate(*exprs)
+        return self._nested_query._build_pipeline(source).aggregate(*exprs)

google/cloud/firestore_v1/base_collection.py

@@ -48,6 +48,7 @@
     from google.cloud.firestore_v1.async_document import AsyncDocumentReference
     from google.cloud.firestore_v1.document import DocumentReference
     from google.cloud.firestore_v1.field_path import FieldPath
+    from google.cloud.firestore_v1.pipeline_source import PipelineSource
     from google.cloud.firestore_v1.query_profile import ExplainOptions
     from google.cloud.firestore_v1.query_results import QueryResultsList
     from google.cloud.firestore_v1.stream_generator import StreamGenerator
@@ -602,18 +603,20 @@ def find_nearest(
             distance_threshold=distance_threshold,
         )
 
-    def pipeline(self):
+    def _build_pipeline(self, source: "PipelineSource"):
         """
         Convert this query into a Pipeline
 
         Queries containing a `cursor` or `limit_to_last` are not currently supported
 
+        Args:
+            source: the PipelineSource to build the pipeline off o
         Raises:
             - NotImplementedError: raised if the query contains a `cursor` or `limit_to_last`
         Returns:
             a Pipeline representing the query
         """
-        return self._query().pipeline()
+        return self._query()._build_pipeline(source)
 
 
 def _auto_id() -> str:

google/cloud/firestore_v1/base_pipeline.py

@@ -13,15 +13,13 @@
 # limitations under the License.
 
 from __future__ import annotations
-from typing import Iterable, Sequence, TYPE_CHECKING
+from typing import Sequence, TYPE_CHECKING
 from google.cloud.firestore_v1 import pipeline_stages as stages
 from google.cloud.firestore_v1.types.pipeline import (
     StructuredPipeline as StructuredPipeline_pb,
 )
 from google.cloud.firestore_v1.vector import Vector
 from google.cloud.firestore_v1.base_vector_query import DistanceMeasure
-from google.cloud.firestore_v1.types.firestore import ExecutePipelineRequest
-from google.cloud.firestore_v1.pipeline_result import PipelineResult
 from google.cloud.firestore_v1.pipeline_expressions import (
     AggregateFunction,
     AliasedExpression,
@@ -30,13 +28,10 @@
     BooleanExpression,
     Selectable,
 )
-from google.cloud.firestore_v1 import _helpers
 
 if TYPE_CHECKING:  # pragma: NO COVER
     from google.cloud.firestore_v1.client import Client
     from google.cloud.firestore_v1.async_client import AsyncClient
-    from google.cloud.firestore_v1.types.firestore import ExecutePipelineResponse
-    from google.cloud.firestore_v1.transaction import BaseTransaction
 
 
 class _BasePipeline:
@@ -87,9 +82,10 @@ def __repr__(self):
             stages_str = ",\n  ".join([repr(s) for s in self.stages])
             return f"{cls_str}(\n  {stages_str}\n)"
 
-    def _to_pb(self) -> StructuredPipeline_pb:
+    def _to_pb(self, **options) -> StructuredPipeline_pb:
         return StructuredPipeline_pb(
-            pipeline={"stages": [s._to_pb() for s in self.stages]}
+            pipeline={"stages": [s._to_pb() for s in self.stages]},
+            options=options,
         )
 
     def _append(self, new_stage):
@@ -98,44 +94,6 @@ def _append(self, new_stage):
         """
         return self.__class__._create_with_stages(self._client, *self.stages, new_stage)
 
-    def _prep_execute_request(
-        self, transaction: BaseTransaction | None
-    ) -> ExecutePipelineRequest:
-        """
-        shared logic for creating an ExecutePipelineRequest
-        """
-        database_name = (
-            f"projects/{self._client.project}/databases/{self._client._database}"
-        )
-        transaction_id = (
-            _helpers.get_transaction_id(transaction)
-            if transaction is not None
-            else None
-        )
-        request = ExecutePipelineRequest(
-            database=database_name,
-            transaction=transaction_id,
-            structured_pipeline=self._to_pb(),
-        )
-        return request
-
-    def _execute_response_helper(
-        self, response: ExecutePipelineResponse
-    ) -> Iterable[PipelineResult]:
-        """
-        shared logic for unpacking an ExecutePipelineReponse into PipelineResults
-        """
-        for doc in response.results:
-            ref = self._client.document(doc.name) if doc.name else None
-            yield PipelineResult(
-                self._client,
-                doc.fields,
-                ref,
-                response._pb.execution_time,
-                doc._pb.create_time if doc.create_time else None,
-                doc._pb.update_time if doc.update_time else None,
-            )
-
     def add_fields(self, *fields: Selectable) -> "_BasePipeline":
         """
         Adds new fields to outputs from previous stages.

google/cloud/firestore_v1/base_query.py

@@ -67,6 +67,7 @@
     from google.cloud.firestore_v1.query_profile import ExplainOptions
     from google.cloud.firestore_v1.query_results import QueryResultsList
     from google.cloud.firestore_v1.stream_generator import StreamGenerator
+    from google.cloud.firestore_v1.pipeline_source import PipelineSource
 
     import datetime
 
@@ -1129,24 +1130,23 @@ def recursive(self: QueryType) -> QueryType:
 
         return copied
 
-    def pipeline(self):
+    def _build_pipeline(self, source: "PipelineSource"):
         """
         Convert this query into a Pipeline
 
         Queries containing a `cursor` or `limit_to_last` are not currently supported
 
+        Args:
+            source: the PipelineSource to build the pipeline off of
         Raises:
-            - ValueError: raised if Query wasn't created with an associated client
             - NotImplementedError: raised if the query contains a `cursor` or `limit_to_last`
         Returns:
             a Pipeline representing the query
         """
-        if not self._client:
-            raise ValueError("Query does not have an associated client")
         if self._all_descendants:
-            ppl = self._client.pipeline().collection_group(self._parent.id)
+            ppl = source.collection_group(self._parent.id)
         else:
-            ppl = self._client.pipeline().collection(self._parent._path)
+            ppl = source.collection(self._parent._path)
 
         # Filters
         for filter_ in self._field_filters: