Merge branch 'pipeline_queries_2_query_parity' into pipeline_queries_3_stable_stages

Author: daniel-sanche

google/cloud/firestore_v1/_helpers.py

@@ -120,6 +120,9 @@ def __ne__(self, other):
         else:
             return not equality_val
 
+    def __repr__(self):
+        return f"{type(self).__name__}(latitude={self.latitude}, longitude={self.longitude})"
+
 
 def verify_path(path, is_collection) -> None:
     """Verifies that a ``path`` has the correct form.

google/cloud/firestore_v1/async_client.py

@@ -417,4 +417,4 @@ def transaction(self, **kwargs) -> AsyncTransaction:
 
     @property
     def _pipeline_cls(self):
-        raise AsyncPipeline
+        return AsyncPipeline

google/cloud/firestore_v1/async_pipeline.py

@@ -13,16 +13,14 @@
 # limitations under the License.
 
 from __future__ import annotations
-import datetime
 from typing import AsyncIterable, TYPE_CHECKING
 from google.cloud.firestore_v1 import pipeline_stages as stages
-from google.cloud.firestore_v1.types.firestore import ExecutePipelineRequest
-from google.cloud.firestore_v1.async_document import AsyncDocumentReference
 from google.cloud.firestore_v1.base_pipeline import _BasePipeline
-from google.cloud.firestore_v1.pipeline_result import PipelineResult
 
 if TYPE_CHECKING:
     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
 
 
 class AsyncPipeline(_BasePipeline):
@@ -58,29 +56,24 @@ def __init__(self, client: AsyncClient, *stages: stages.Stage):
         """
         super().__init__(client, *stages)
 
-    async def execute(self) -> AsyncIterable[PipelineResult]:
-        database_name = (
-            f"projects/{self._client.project}/databases/{self._client._database}"
-        )
-        request = ExecutePipelineRequest(
-            database=database_name,
-            structured_pipeline=self._to_pb(),
-            read_time=datetime.datetime.now(),
-        )
+    async def execute(
+        self,
+        transaction: "AsyncTransaction" | None = None,
+    ) -> AsyncIterable[PipelineResult]:
+        """
+        Executes this pipeline, providing results through an Iterable
+
+        Args:
+            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).
+        """
+        request = self._prep_execute_request(transaction)
         async for response in await self._client._firestore_api.execute_pipeline(
             request
         ):
-            for doc in response.results:
-                doc_ref = (
-                    AsyncDocumentReference(doc.name, client=self._client)
-                    if doc.name
-                    else None
-                )
-                yield PipelineResult(
-                    self._client,
-                    doc.fields,
-                    doc_ref,
-                    response._pb.execution_time,
-                    doc.create_time,
-                    doc.update_tiem,
-                )
+            for result in self._execute_response_helper(response):
+                yield result

google/cloud/firestore_v1/base_client.py

@@ -61,6 +61,8 @@
 from google.cloud.firestore_v1.bulk_writer import BulkWriter, BulkWriterOptions
 from google.cloud.firestore_v1.field_path import render_field_path
 from google.cloud.firestore_v1.services.firestore import client as firestore_client
+from google.cloud.firestore_v1.pipeline_source import PipelineSource
+from google.cloud.firestore_v1.base_pipeline import _BasePipeline
 
 DEFAULT_DATABASE = "(default)"
 """str: The default database used in a :class:`~google.cloud.firestore_v1.client.Client`."""

google/cloud/firestore_v1/base_pipeline.py

@@ -13,18 +13,15 @@
 # limitations under the License.
 
 from __future__ import annotations
-from typing import Optional, Sequence
-from typing_extensions import Self
-from typing import TYPE_CHECKING
+from typing import Iterable, Sequence, TYPE_CHECKING
 from google.cloud.firestore_v1 import pipeline_stages as stages
-from google.cloud.firestore_v1.document import DocumentReference
 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 import _helpers, document
 from google.cloud.firestore_v1.pipeline_expressions import (
     Accumulator,
     Expr,
@@ -34,10 +31,13 @@
     Selectable,
     SampleOptions,
 )
+from google.cloud.firestore_v1 import _helpers
 
 if TYPE_CHECKING:
     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:
@@ -62,13 +62,14 @@ def __init__(self, client: Client | AsyncClient, *stages: stages.Stage):
         self.stages = tuple(stages)
 
     def __repr__(self):
+        cls_str = type(self).__name__
         if not self.stages:
-            return "Pipeline()"
+            return f"{cls_str}()"
         elif len(self.stages) == 1:
-            return f"Pipeline({self.stages[0]!r})"
+            return f"{cls_str}({self.stages[0]!r})"
         else:
             stages_str = ",\n  ".join([repr(s) for s in self.stages])
-            return f"Pipeline(\n  {stages_str}\n)"
+            return f"{cls_str}(\n  {stages_str}\n)"
 
     def _to_pb(self) -> StructuredPipeline_pb:
         return StructuredPipeline_pb(
@@ -81,20 +82,45 @@ def _append(self, new_stage):
         """
         return self.__class__(self._client, *self.stages, new_stage)
 
-    @staticmethod
-    def _parse_response(response_pb, client):
-        for doc in response_pb.results:
-            data = _helpers.decode_dict(doc.fields, client)
-            yield document.DocumentSnapshot(
-                None,
-                data,
-                exists=True,
-                read_time=response_pb._pb.execution_time,
-                create_time=doc.create_time,
-                update_time=doc.update_time,
+    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) -> Self:
+    def add_fields(self, *fields: Selectable) -> "_BasePipeline":
         """
         Adds new fields to outputs from previous stages.
 
@@ -124,7 +150,7 @@ def add_fields(self, *fields: Selectable) -> Self:
         """
         return self._append(stages.AddFields(*fields))
 
-    def remove_fields(self, *fields: Field | str) -> Self:
+    def remove_fields(self, *fields: Field | str) -> "_BasePipeline":
         """
         Removes fields from outputs of previous stages.
 
@@ -146,7 +172,7 @@ def remove_fields(self, *fields: Field | str) -> Self:
         """
         return self._append(stages.RemoveFields(*fields))
 
-    def select(self, *selections: str | Selectable) -> Self:
+    def select(self, *selections: str | Selectable) -> "_BasePipeline":
         """
         Selects or creates a set of fields from the outputs of previous stages.
 
@@ -179,7 +205,7 @@ def select(self, *selections: str | Selectable) -> Self:
         """
         return self._append(stages.Select(*selections))
 
-    def where(self, condition: FilterCondition) -> Self:
+    def where(self, condition: FilterCondition) -> "_BasePipeline":
         """
         Filters the documents from previous stages to only include those matching
         the specified `FilterCondition`.
@@ -223,8 +249,8 @@ def find_nearest(
         field: str | Expr,
         vector: Sequence[float] | "Vector",
         distance_measure: "DistanceMeasure",
-        options: Optional[stages.FindNearestOptions] = None,
-    ) -> Self:
+        options: stages.FindNearestOptions | None = None,
+    ) -> "_BasePipeline":
         """
         Performs vector distance (similarity) search with given parameters on the
         stage inputs.
@@ -274,7 +300,7 @@ def find_nearest(
             stages.FindNearest(field, vector, distance_measure, options)
         )
 
-    def sort(self, *orders: stages.Ordering) -> Self:
+    def sort(self, *orders: stages.Ordering) -> "_BasePipeline":
         """
         Sorts the documents from previous stages based on one or more `Ordering` criteria.
 
@@ -302,7 +328,7 @@ def sort(self, *orders: stages.Ordering) -> Self:
         """
         return self._append(stages.Sort(*orders))
 
-    def sample(self, limit_or_options: int | SampleOptions) -> Self:
+    def sample(self, limit_or_options: int | SampleOptions) -> "_BasePipeline":
         """
         Performs a pseudo-random sampling of the documents from the previous stage.
 
@@ -331,7 +357,7 @@ def sample(self, limit_or_options: int | SampleOptions) -> Self:
         """
         return self._append(stages.Sample(limit_or_options))
 
-    def union(self, other: Self) -> Self:
+    def union(self, other: Self) -> "_BasePipeline":
         """
         Performs a union of all documents from this pipeline and another pipeline,
         including duplicates.
@@ -359,7 +385,7 @@ def unnest(
         field: str | Selectable,
         alias: str | Field | None = None,
         options: Optional[stages.UnnestOptions] = None,
-    ) -> Self:
+    ) -> "_BasePipeline":
         """
         Produces a document for each element in an array field from the previous stage document.
 
@@ -417,7 +443,7 @@ def unnest(
         """
         return self._append(stages.Unnest(field, alias, options))
 
-    def generic_stage(self, name: str, *params: Expr) -> Self:
+    def generic_stage(self, name: str, *params: Expr) -> "_BasePipeline":
         """
         Adds a generic, named stage to the pipeline with specified parameters.
 
@@ -440,7 +466,7 @@ def generic_stage(self, name: str, *params: Expr) -> Self:
         """
         return self._append(stages.GenericStage(name, *params))
 
-    def offset(self, offset: int) -> Self:
+    def offset(self, offset: int) -> "_BasePipeline":
         """
         Skips the first `offset` number of documents from the results of previous stages.
 
@@ -464,7 +490,7 @@ def offset(self, offset: int) -> Self:
         """
         return self._append(stages.Offset(offset))
 
-    def limit(self, limit: int) -> Self:
+    def limit(self, limit: int) -> "_BasePipeline":
         """
         Limits the maximum number of documents returned by previous stages to `limit`.
 
@@ -492,7 +518,7 @@ def aggregate(
         self,
         *accumulators: ExprWithAlias[Accumulator],
         groups: Sequence[str | Selectable] = (),
-    ) -> Self:
+    ) -> "_BasePipeline":
         """
         Performs aggregation operations on the documents from previous stages,
         optionally grouped by specified fields or expressions.
@@ -538,7 +564,7 @@ def aggregate(
         """
         return self._append(stages.Aggregate(*accumulators, groups=groups))
 
-    def distinct(self, *fields: str | Selectable) -> Self:
+    def distinct(self, *fields: str | Selectable) -> "_BasePipeline":
         """
         Returns documents with distinct combinations of values for the specified
         fields or expressions.

google/cloud/firestore_v1/client.py

@@ -399,4 +399,4 @@ def transaction(self, **kwargs) -> Transaction:
 
     @property
     def _pipeline_cls(self):
-        raise Pipeline
+        return Pipeline

google/cloud/firestore_v1/field_path.py

@@ -16,7 +16,7 @@
 from __future__ import annotations
 import re
 from collections import abc
-from typing import Iterable, cast
+from typing import Any, Iterable, cast, MutableMapping
 
 _FIELD_PATH_MISSING_TOP = "{!r} is not contained in the data"
 _FIELD_PATH_MISSING_KEY = "{!r} is not contained in the data for the key {!r}"
@@ -170,7 +170,7 @@ def render_field_path(field_names: Iterable[str]):
 get_field_path = render_field_path  # backward-compatibility
 
 
-def get_nested_value(field_path: str, data: dict):
+def get_nested_value(field_path: str, data: MutableMapping[str, Any]):
     """Get a (potentially nested) value from a dictionary.
 
     If the data is nested, for example: