chore: updated gapic layer for execute_query

Author: daniel-sanche

google/cloud/firestore_v1/services/firestore/async_client.py

@@ -52,6 +52,7 @@
 from google.cloud.firestore_v1.types import common
 from google.cloud.firestore_v1.types import document
 from google.cloud.firestore_v1.types import document as gf_document
+from google.cloud.firestore_v1.types import explain_stats
 from google.cloud.firestore_v1.types import firestore
 from google.cloud.firestore_v1.types import query
 from google.cloud.firestore_v1.types import query_profile
@@ -236,6 +237,9 @@ def __init__(
                 If a Callable is given, it will be called with the same set of initialization
                 arguments as used in the FirestoreTransport constructor.
                 If set to None, a transport is chosen automatically.
+                NOTE: "rest" transport functionality is currently in a
+                beta state (preview). We welcome your feedback via an
+                issue in this library's source repository.
             client_options (Optional[Union[google.api_core.client_options.ClientOptions, dict]]):
                 Custom options for the client.
 
@@ -1247,6 +1251,109 @@ async def sample_run_query():
         # Done; return the response.
         return response
 
+    def execute_pipeline(
+        self,
+        request: Optional[Union[firestore.ExecutePipelineRequest, dict]] = None,
+        *,
+        retry: OptionalRetry = gapic_v1.method.DEFAULT,
+        timeout: Union[float, object] = gapic_v1.method.DEFAULT,
+        metadata: Sequence[Tuple[str, Union[str, bytes]]] = (),
+    ) -> Awaitable[AsyncIterable[firestore.ExecutePipelineResponse]]:
+        r"""Executes a pipeline query.
+
+        .. code-block:: python
+
+            # This snippet has been automatically generated and should be regarded as a
+            # code template only.
+            # It will require modifications to work:
+            # - It may require correct/in-range values for request initialization.
+            # - It may require specifying regional endpoints when creating the service
+            #   client as shown in:
+            #   https://googleapis.dev/python/google-api-core/latest/client_options.html
+            from google.cloud import firestore_v1
+
+            async def sample_execute_pipeline():
+                # Create a client
+                client = firestore_v1.FirestoreAsyncClient()
+
+                # Initialize request argument(s)
+                structured_pipeline = firestore_v1.StructuredPipeline()
+                structured_pipeline.pipeline.stages.name = "name_value"
+
+                request = firestore_v1.ExecutePipelineRequest(
+                    structured_pipeline=structured_pipeline,
+                    transaction=b'transaction_blob',
+                    database="database_value",
+                )
+
+                # Make the request
+                stream = await client.execute_pipeline(request=request)
+
+                # Handle the response
+                async for response in stream:
+                    print(response)
+
+        Args:
+            request (Optional[Union[google.cloud.firestore_v1.types.ExecutePipelineRequest, dict]]):
+                The request object. The request for
+                [Firestore.ExecutePipeline][google.firestore.v1.Firestore.ExecutePipeline].
+            retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any,
+                should be retried.
+            timeout (float): The timeout for this request.
+            metadata (Sequence[Tuple[str, Union[str, bytes]]]): Key/value pairs which should be
+                sent along with the request as metadata. Normally, each value must be of type `str`,
+                but for metadata keys ending with the suffix `-bin`, the corresponding values must
+                be of type `bytes`.
+
+        Returns:
+            AsyncIterable[google.cloud.firestore_v1.types.ExecutePipelineResponse]:
+                The response for [Firestore.Execute][].
+        """
+        # Create or coerce a protobuf request object.
+        # - Use the request object if provided (there's no risk of modifying the input as
+        #   there are no flattened fields), or create one.
+        if not isinstance(request, firestore.ExecutePipelineRequest):
+            request = firestore.ExecutePipelineRequest(request)
+
+        # Wrap the RPC method; this adds retry and timeout information,
+        # and friendly error handling.
+        rpc = self._client._transport._wrapped_methods[
+            self._client._transport.execute_pipeline
+        ]
+
+        header_params = {}
+
+        routing_param_regex = re.compile("^projects/(?P<project_id>[^/]+)(?:/.*)?$")
+        regex_match = routing_param_regex.match(request.database)
+        if regex_match and regex_match.group("project_id"):
+            header_params["project_id"] = regex_match.group("project_id")
+
+        routing_param_regex = re.compile(
+            "^projects/[^/]+/databases/(?P<database_id>[^/]+)(?:/.*)?$"
+        )
+        regex_match = routing_param_regex.match(request.database)
+        if regex_match and regex_match.group("database_id"):
+            header_params["database_id"] = regex_match.group("database_id")
+
+        if header_params:
+            metadata = tuple(metadata) + (
+                gapic_v1.routing_header.to_grpc_metadata(header_params),
+            )
+
+        # Validate the universe domain.
+        self._client._validate_universe_domain()
+
+        # Send the request.
+        response = rpc(
+            request,
+            retry=retry,
+            timeout=timeout,
+            metadata=metadata,
+        )
+
+        # Done; return the response.
+        return response
+
     def run_aggregation_query(
         self,
         request: Optional[Union[firestore.RunAggregationQueryRequest, dict]] = None,

google/cloud/firestore_v1/services/firestore/client.py

@@ -67,6 +67,7 @@
 from google.cloud.firestore_v1.types import common
 from google.cloud.firestore_v1.types import document
 from google.cloud.firestore_v1.types import document as gf_document
+from google.cloud.firestore_v1.types import explain_stats
 from google.cloud.firestore_v1.types import firestore
 from google.cloud.firestore_v1.types import query
 from google.cloud.firestore_v1.types import query_profile
@@ -551,6 +552,9 @@ def __init__(
                 If a Callable is given, it will be called with the same set of initialization
                 arguments as used in the FirestoreTransport constructor.
                 If set to None, a transport is chosen automatically.
+                NOTE: "rest" transport functionality is currently in a
+                beta state (preview). We welcome your feedback via an
+                issue in this library's source repository.
             client_options (Optional[Union[google.api_core.client_options.ClientOptions, dict]]):
                 Custom options for the client.
 
@@ -1630,6 +1634,107 @@ def sample_run_query():
         # Done; return the response.
         return response
 
+    def execute_pipeline(
+        self,
+        request: Optional[Union[firestore.ExecutePipelineRequest, dict]] = None,
+        *,
+        retry: OptionalRetry = gapic_v1.method.DEFAULT,
+        timeout: Union[float, object] = gapic_v1.method.DEFAULT,
+        metadata: Sequence[Tuple[str, Union[str, bytes]]] = (),
+    ) -> Iterable[firestore.ExecutePipelineResponse]:
+        r"""Executes a pipeline query.
+
+        .. code-block:: python
+
+            # This snippet has been automatically generated and should be regarded as a
+            # code template only.
+            # It will require modifications to work:
+            # - It may require correct/in-range values for request initialization.
+            # - It may require specifying regional endpoints when creating the service
+            #   client as shown in:
+            #   https://googleapis.dev/python/google-api-core/latest/client_options.html
+            from google.cloud import firestore_v1
+
+            def sample_execute_pipeline():
+                # Create a client
+                client = firestore_v1.FirestoreClient()
+
+                # Initialize request argument(s)
+                structured_pipeline = firestore_v1.StructuredPipeline()
+                structured_pipeline.pipeline.stages.name = "name_value"
+
+                request = firestore_v1.ExecutePipelineRequest(
+                    structured_pipeline=structured_pipeline,
+                    transaction=b'transaction_blob',
+                    database="database_value",
+                )
+
+                # Make the request
+                stream = client.execute_pipeline(request=request)
+
+                # Handle the response
+                for response in stream:
+                    print(response)
+
+        Args:
+            request (Union[google.cloud.firestore_v1.types.ExecutePipelineRequest, dict]):
+                The request object. The request for
+                [Firestore.ExecutePipeline][google.firestore.v1.Firestore.ExecutePipeline].
+            retry (google.api_core.retry.Retry): Designation of what errors, if any,
+                should be retried.
+            timeout (float): The timeout for this request.
+            metadata (Sequence[Tuple[str, Union[str, bytes]]]): Key/value pairs which should be
+                sent along with the request as metadata. Normally, each value must be of type `str`,
+                but for metadata keys ending with the suffix `-bin`, the corresponding values must
+                be of type `bytes`.
+
+        Returns:
+            Iterable[google.cloud.firestore_v1.types.ExecutePipelineResponse]:
+                The response for [Firestore.Execute][].
+        """
+        # Create or coerce a protobuf request object.
+        # - Use the request object if provided (there's no risk of modifying the input as
+        #   there are no flattened fields), or create one.
+        if not isinstance(request, firestore.ExecutePipelineRequest):
+            request = firestore.ExecutePipelineRequest(request)
+
+        # Wrap the RPC method; this adds retry and timeout information,
+        # and friendly error handling.
+        rpc = self._transport._wrapped_methods[self._transport.execute_pipeline]
+
+        header_params = {}
+
+        routing_param_regex = re.compile("^projects/(?P<project_id>[^/]+)(?:/.*)?$")
+        regex_match = routing_param_regex.match(request.database)
+        if regex_match and regex_match.group("project_id"):
+            header_params["project_id"] = regex_match.group("project_id")
+
+        routing_param_regex = re.compile(
+            "^projects/[^/]+/databases/(?P<database_id>[^/]+)(?:/.*)?$"
+        )
+        regex_match = routing_param_regex.match(request.database)
+        if regex_match and regex_match.group("database_id"):
+            header_params["database_id"] = regex_match.group("database_id")
+
+        if header_params:
+            metadata = tuple(metadata) + (
+                gapic_v1.routing_header.to_grpc_metadata(header_params),
+            )
+
+        # Validate the universe domain.
+        self._validate_universe_domain()
+
+        # Send the request.
+        response = rpc(
+            request,
+            retry=retry,
+            timeout=timeout,
+            metadata=metadata,
+        )
+
+        # Done; return the response.
+        return response
+
     def run_aggregation_query(
         self,
         request: Optional[Union[firestore.RunAggregationQueryRequest, dict]] = None,

google/cloud/firestore_v1/services/firestore/transports/base.py

@@ -286,6 +286,11 @@ def _prep_wrapped_messages(self, client_info):
                 default_timeout=300.0,
                 client_info=client_info,
             ),
+            self.execute_pipeline: gapic_v1.method.wrap_method(
+                self.execute_pipeline,
+                default_timeout=None,
+                client_info=client_info,
+            ),
             self.run_aggregation_query: gapic_v1.method.wrap_method(
                 self.run_aggregation_query,
                 default_retry=retries.Retry(
@@ -509,6 +514,18 @@ def run_query(
     ]:
         raise NotImplementedError()
 
+    @property
+    def execute_pipeline(
+        self,
+    ) -> Callable[
+        [firestore.ExecutePipelineRequest],
+        Union[
+            firestore.ExecutePipelineResponse,
+            Awaitable[firestore.ExecutePipelineResponse],
+        ],
+    ]:
+        raise NotImplementedError()
+
     @property
     def run_aggregation_query(
         self,

google/cloud/firestore_v1/services/firestore/transports/grpc.py

@@ -571,6 +571,34 @@ def run_query(
             )
         return self._stubs["run_query"]
 
+    @property
+    def execute_pipeline(
+        self,
+    ) -> Callable[
+        [firestore.ExecutePipelineRequest], firestore.ExecutePipelineResponse
+    ]:
+        r"""Return a callable for the execute pipeline method over gRPC.
+
+        Executes a pipeline query.
+
+        Returns:
+            Callable[[~.ExecutePipelineRequest],
+                    ~.ExecutePipelineResponse]:
+                A function that, when called, will call the underlying RPC
+                on the server.
+        """
+        # Generate a "stub function" on-the-fly which will actually make
+        # the request.
+        # gRPC handles serialization and deserialization, so we just need
+        # to pass in the functions for each.
+        if "execute_pipeline" not in self._stubs:
+            self._stubs["execute_pipeline"] = self._logged_channel.unary_stream(
+                "/google.firestore.v1.Firestore/ExecutePipeline",
+                request_serializer=firestore.ExecutePipelineRequest.serialize,
+                response_deserializer=firestore.ExecutePipelineResponse.deserialize,
+            )
+        return self._stubs["execute_pipeline"]
+
     @property
     def run_aggregation_query(
         self,

google/cloud/firestore_v1/services/firestore/transports/grpc_asyncio.py

@@ -587,6 +587,34 @@ def run_query(
             )
         return self._stubs["run_query"]
 
+    @property
+    def execute_pipeline(
+        self,
+    ) -> Callable[
+        [firestore.ExecutePipelineRequest], Awaitable[firestore.ExecutePipelineResponse]
+    ]:
+        r"""Return a callable for the execute pipeline method over gRPC.
+
+        Executes a pipeline query.
+
+        Returns:
+            Callable[[~.ExecutePipelineRequest],
+                    Awaitable[~.ExecutePipelineResponse]]:
+                A function that, when called, will call the underlying RPC
+                on the server.
+        """
+        # Generate a "stub function" on-the-fly which will actually make
+        # the request.
+        # gRPC handles serialization and deserialization, so we just need
+        # to pass in the functions for each.
+        if "execute_pipeline" not in self._stubs:
+            self._stubs["execute_pipeline"] = self._logged_channel.unary_stream(
+                "/google.firestore.v1.Firestore/ExecutePipeline",
+                request_serializer=firestore.ExecutePipelineRequest.serialize,
+                response_deserializer=firestore.ExecutePipelineResponse.deserialize,
+            )
+        return self._stubs["execute_pipeline"]
+
     @property
     def run_aggregation_query(
         self,
@@ -962,6 +990,11 @@ def _prep_wrapped_messages(self, client_info):
                 default_timeout=300.0,
                 client_info=client_info,
             ),
+            self.execute_pipeline: self._wrap_method(
+                self.execute_pipeline,
+                default_timeout=None,
+                client_info=client_info,
+            ),
             self.run_aggregation_query: self._wrap_method(
                 self.run_aggregation_query,
                 default_retry=retries.AsyncRetry(