feat: add context manager support in client (googleapis#9)

gcf-owl-bot[bot] · web-flow · commit bdcf454b8d97 · 2021-10-08T14:26:23.000Z
- [ ] Regenerate this pull request now. chore: fix docstring for first attribute of protos committer: @busunkim96 PiperOrigin-RevId: 401271153 Source-Link: googleapis/googleapis@787f8c9 Source-Link: https://github.com/googleapis/googleapis-gen/commit/81decffe9fc72396a8153e756d1d67a6eecfd620 Copy-Tag: eyJwIjoiLmdpdGh1Yi8uT3dsQm90LnlhbWwiLCJoIjoiODFkZWNmZmU5ZmM3MjM5NmE4MTUzZTc1NmQxZDY3YTZlZWNmZDYyMCJ9
diff --git a/google/cloud/deploy_v1/services/cloud_deploy/async_client.py b/google/cloud/deploy_v1/services/cloud_deploy/async_client.py
@@ -1811,6 +1811,12 @@ async def get_config(
         # Done; return the response.
         return response
 
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb):
+        await self.transport.close()
+
 
 try:
     DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(
diff --git a/google/cloud/deploy_v1/services/cloud_deploy/client.py b/google/cloud/deploy_v1/services/cloud_deploy/client.py
@@ -474,10 +474,7 @@ def __init__(
                 client_cert_source_for_mtls=client_cert_source_func,
                 quota_project_id=client_options.quota_project_id,
                 client_info=client_info,
-                always_use_jwt_access=(
-                    Transport == type(self).get_transport_class("grpc")
-                    or Transport == type(self).get_transport_class("grpc_asyncio")
-                ),
+                always_use_jwt_access=True,
             )
 
     def list_delivery_pipelines(
@@ -2029,6 +2026,19 @@ def get_config(
         # Done; return the response.
         return response
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, type, value, traceback):
+        """Releases underlying transport's resources.
+
+        .. warning::
+            ONLY use as a context manager if the transport is NOT shared
+            with other clients! Exiting the with block will CLOSE the transport
+            and may cause errors in other clients!
+        """
+        self.transport.close()
+
 
 try:
     DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(
diff --git a/google/cloud/deploy_v1/services/cloud_deploy/transports/base.py b/google/cloud/deploy_v1/services/cloud_deploy/transports/base.py
@@ -315,6 +315,15 @@ def _prep_wrapped_messages(self, client_info):
             ),
         }
 
+    def close(self):
+        """Closes resources associated with the transport.
+
+       .. warning::
+            Only call this method if the transport is NOT shared
+            with other clients - this may cause errors in other clients!
+        """
+        raise NotImplementedError()
+
     @property
     def operations_client(self) -> operations_v1.OperationsClient:
         """Return the client designed to process long-running operations."""
diff --git a/google/cloud/deploy_v1/services/cloud_deploy/transports/grpc.py b/google/cloud/deploy_v1/services/cloud_deploy/transports/grpc.py
@@ -733,5 +733,8 @@ def get_config(
             )
         return self._stubs["get_config"]
 
+    def close(self):
+        self.grpc_channel.close()
+
 
 __all__ = ("CloudDeployGrpcTransport",)
diff --git a/google/cloud/deploy_v1/services/cloud_deploy/transports/grpc_asyncio.py b/google/cloud/deploy_v1/services/cloud_deploy/transports/grpc_asyncio.py
@@ -755,5 +755,8 @@ def get_config(
             )
         return self._stubs["get_config"]
 
+    def close(self):
+        return self.grpc_channel.close()
+
 
 __all__ = ("CloudDeployGrpcAsyncIOTransport",)
diff --git a/google/cloud/deploy_v1/types/cloud_deploy.py b/google/cloud/deploy_v1/types/cloud_deploy.py
@@ -146,6 +146,7 @@ class SerialPipeline(proto.Message):
 
 class Stage(proto.Message):
     r"""Stage specifies a location to which to deploy.
+
     Attributes:
         target_id (str):
             The target_id to which this stage points. This field refers
@@ -224,6 +225,7 @@ class PipelineCondition(proto.Message):
 
 class ListDeliveryPipelinesRequest(proto.Message):
     r"""The request object for ``ListDeliveryPipelines``.
+
     Attributes:
         parent (str):
             Required. The parent, which owns this collection of
@@ -260,6 +262,7 @@ class ListDeliveryPipelinesRequest(proto.Message):
 
 class ListDeliveryPipelinesResponse(proto.Message):
     r"""The response object from ``ListDeliveryPipelines``.
+
     Attributes:
         delivery_pipelines (Sequence[google.cloud.deploy_v1.types.DeliveryPipeline]):
             The ``DeliveryPipeline`` objects.
@@ -284,6 +287,7 @@ def raw_page(self):
 
 class GetDeliveryPipelineRequest(proto.Message):
     r"""The request object for ``GetDeliveryPipeline``
+
     Attributes:
         name (str):
             Required. Name of the ``DeliveryPipeline``. Format must be
@@ -295,6 +299,7 @@ class GetDeliveryPipelineRequest(proto.Message):
 
 class CreateDeliveryPipelineRequest(proto.Message):
     r"""The request object for ``CreateDeliveryPipeline``.
+
     Attributes:
         parent (str):
             Required. The parent collection in which the
@@ -339,6 +344,7 @@ class CreateDeliveryPipelineRequest(proto.Message):
 
 class UpdateDeliveryPipelineRequest(proto.Message):
     r"""The request object for ``UpdateDeliveryPipeline``.
+
     Attributes:
         update_mask (google.protobuf.field_mask_pb2.FieldMask):
             Required. Field mask is used to specify the fields to be
@@ -390,6 +396,7 @@ class UpdateDeliveryPipelineRequest(proto.Message):
 
 class DeleteDeliveryPipelineRequest(proto.Message):
     r"""The request object for ``DeleteDeliveryPipeline``.
+
     Attributes:
         name (str):
             Required. The name of the ``DeliveryPipeline`` to delete.
@@ -547,6 +554,7 @@ class ExecutionEnvironmentUsage(proto.Enum):
 
 class DefaultPool(proto.Message):
     r"""Execution using the default Cloud Build pool.
+
     Attributes:
         service_account (str):
             Optional. Google service account to use for execution. If
@@ -568,6 +576,7 @@ class DefaultPool(proto.Message):
 
 class PrivatePool(proto.Message):
     r"""Execution using a private Cloud Build pool.
+
     Attributes:
         worker_pool (str):
             Required. Resource name of the Cloud Build worker pool to
@@ -594,6 +603,7 @@ class PrivatePool(proto.Message):
 
 class GkeCluster(proto.Message):
     r"""Information specifying a GKE Cluster.
+
     Attributes:
         cluster (str):
             Information specifying a GKE Cluster. Format is
@@ -605,6 +615,7 @@ class GkeCluster(proto.Message):
 
 class ListTargetsRequest(proto.Message):
     r"""The request object for ``ListTargets``.
+
     Attributes:
         parent (str):
             Required. The parent, which owns this collection of targets.
@@ -641,6 +652,7 @@ class ListTargetsRequest(proto.Message):
 
 class ListTargetsResponse(proto.Message):
     r"""The response object from ``ListTargets``.
+
     Attributes:
         targets (Sequence[google.cloud.deploy_v1.types.Target]):
             The ``Target`` objects.
@@ -663,6 +675,7 @@ def raw_page(self):
 
 class GetTargetRequest(proto.Message):
     r"""The request object for ``GetTarget``.
+
     Attributes:
         name (str):
             Required. Name of the ``Target``. Format must be
@@ -674,6 +687,7 @@ class GetTargetRequest(proto.Message):
 
 class CreateTargetRequest(proto.Message):
     r"""The request object for ``CreateTarget``.
+
     Attributes:
         parent (str):
             Required. The parent collection in which the ``Target``
@@ -716,6 +730,7 @@ class CreateTargetRequest(proto.Message):
 
 class UpdateTargetRequest(proto.Message):
     r"""The request object for ``UpdateTarget``.
+
     Attributes:
         update_mask (google.protobuf.field_mask_pb2.FieldMask):
             Required. Field mask is used to specify the fields to be
@@ -764,6 +779,7 @@ class UpdateTargetRequest(proto.Message):
 
 class DeleteTargetRequest(proto.Message):
     r"""The request object for ``DeleteTarget``.
+
     Attributes:
         name (str):
             Required. The name of the ``Target`` to delete. Format
@@ -896,6 +912,7 @@ class RenderState(proto.Enum):
 
     class TargetRender(proto.Message):
         r"""Details of rendering for a single target.
+
         Attributes:
             rendering_build (str):
                 Output only. The resource name of the Cloud Build ``Build``
@@ -953,6 +970,7 @@ class TargetRenderState(proto.Enum):
 
 class BuildArtifact(proto.Message):
     r"""Description of an a image to use during Skaffold rendering.
+
     Attributes:
         image (str):
             Image name in Skaffold configuration.
@@ -969,6 +987,7 @@ class BuildArtifact(proto.Message):
 
 class TargetArtifact(proto.Message):
     r"""The artifacts produced by a target render operation.
+
     Attributes:
         artifact_uri (str):
             Output only. URI of a directory containing
@@ -990,6 +1009,7 @@ class TargetArtifact(proto.Message):
 
 class ListReleasesRequest(proto.Message):
     r"""The request object for ``ListReleases``.
+
     Attributes:
         parent (str):
             Required. The ``DeliveryPipeline`` which owns this
@@ -1025,6 +1045,7 @@ class ListReleasesRequest(proto.Message):
 
 class ListReleasesResponse(proto.Message):
     r"""The response object from ``ListReleases``.
+
     Attributes:
         releases (Sequence[google.cloud.deploy_v1.types.Release]):
             The ``Release`` objects.
@@ -1047,6 +1068,7 @@ def raw_page(self):
 
 class GetReleaseRequest(proto.Message):
     r"""The request object for ``GetRelease``.
+
     Attributes:
         name (str):
             Required. Name of the ``Release``. Format must be
@@ -1058,6 +1080,7 @@ class GetReleaseRequest(proto.Message):
 
 class CreateReleaseRequest(proto.Message):
     r"""The request object for ``CreateRelease``,
+
     Attributes:
         parent (str):
             Required. The parent collection in which the ``Release``
@@ -1209,6 +1232,7 @@ class State(proto.Enum):
 
 class ListRolloutsRequest(proto.Message):
     r"""ListRolloutsRequest is the request object used by ``ListRollouts``.
+
     Attributes:
         parent (str):
             Required. The ``Release`` which owns this collection of
@@ -1268,6 +1292,7 @@ def raw_page(self):
 
 class GetRolloutRequest(proto.Message):
     r"""GetRolloutRequest is the request object used by ``GetRollout``.
+
     Attributes:
         name (str):
             Required. Name of the ``Rollout``. Format must be
@@ -1323,6 +1348,7 @@ class CreateRolloutRequest(proto.Message):
 
 class OperationMetadata(proto.Message):
     r"""Represents the metadata of the long-running operation.
+
     Attributes:
         create_time (google.protobuf.timestamp_pb2.Timestamp):
             Output only. The time the operation was
@@ -1361,6 +1387,7 @@ class OperationMetadata(proto.Message):
 
 class ApproveRolloutRequest(proto.Message):
     r"""The request object used by ``ApproveRollout``.
+
     Attributes:
         name (str):
             Required. Name of the Rollout. Format is
@@ -1375,11 +1402,13 @@ class ApproveRolloutRequest(proto.Message):
 
 
 class ApproveRolloutResponse(proto.Message):
-    r"""The response object from ``ApproveRollout``.    """
+    r"""The response object from ``ApproveRollout``.
+    """
 
 
 class Config(proto.Message):
     r"""Service-wide configuration.
+
     Attributes:
         name (str):
             Name of the configuration.
@@ -1401,6 +1430,7 @@ class Config(proto.Message):
 
 class SkaffoldVersion(proto.Message):
     r"""Details of a supported Skaffold version.
+
     Attributes:
         version (str):
             Release version number. For example,
@@ -1416,6 +1446,7 @@ class SkaffoldVersion(proto.Message):
 
 class GetConfigRequest(proto.Message):
     r"""Request to get a configuration.
+
     Attributes:
         name (str):
             Required. Name of requested configuration.
diff --git a/tests/unit/gapic/deploy_v1/test_cloud_deploy.py b/tests/unit/gapic/deploy_v1/test_cloud_deploy.py
@@ -32,6 +32,7 @@
 from google.api_core import grpc_helpers_async
 from google.api_core import operation_async  # type: ignore
 from google.api_core import operations_v1
+from google.api_core import path_template
 from google.auth import credentials as ga_credentials
 from google.auth.exceptions import MutualTLSChannelError
 from google.cloud.deploy_v1.services.cloud_deploy import CloudDeployAsyncClient
@@ -5023,6 +5024,9 @@ def test_cloud_deploy_base_transport():
         with pytest.raises(NotImplementedError):
             getattr(transport, method)(request=object())
 
+    with pytest.raises(NotImplementedError):
+        transport.close()
+
     # Additionally, the LRO client (a property) should
     # also raise NotImplementedError
     with pytest.raises(NotImplementedError):
@@ -5704,3 +5708,49 @@ def test_client_withDEFAULT_CLIENT_INFO():
             credentials=ga_credentials.AnonymousCredentials(), client_info=client_info,
         )
         prep.assert_called_once_with(client_info)
+
+
+@pytest.mark.asyncio
+async def test_transport_close_async():
+    client = CloudDeployAsyncClient(
+        credentials=ga_credentials.AnonymousCredentials(), transport="grpc_asyncio",
+    )
+    with mock.patch.object(
+        type(getattr(client.transport, "grpc_channel")), "close"
+    ) as close:
+        async with client:
+            close.assert_not_called()
+        close.assert_called_once()
+
+
+def test_transport_close():
+    transports = {
+        "grpc": "_grpc_channel",
+    }
+
+    for transport, close_name in transports.items():
+        client = CloudDeployClient(
+            credentials=ga_credentials.AnonymousCredentials(), transport=transport
+        )
+        with mock.patch.object(
+            type(getattr(client.transport, close_name)), "close"
+        ) as close:
+            with client:
+                close.assert_not_called()
+            close.assert_called_once()
+
+
+def test_client_ctx():
+    transports = [
+        "grpc",
+    ]
+    for transport in transports:
+        client = CloudDeployClient(
+            credentials=ga_credentials.AnonymousCredentials(), transport=transport
+        )
+        # Test client calls underlying transport.
+        with mock.patch.object(type(client.transport), "close") as close:
+            close.assert_not_called()
+            with client:
+                pass
+            close.assert_called()

Original file line number	Diff line number	Diff line change
`@@ -733,5 +733,8 @@ def get_config(`
`733`	`733`	`)`
`734`	`734`	`return self._stubs["get_config"]`
`735`	`735`
	`736`	`+ def close(self):`
	`737`	`+ self.grpc_channel.close()`
	`738`	`+`
`736`	`739`
`737`	`740`	`__all__ = ("CloudDeployGrpcTransport",)`
Original file line number	Diff line number	Diff line change
`@@ -755,5 +755,8 @@ def get_config(`
`755`	`755`	`)`
`756`	`756`	`return self._stubs["get_config"]`
`757`	`757`
	`758`	`+ def close(self):`
	`759`	`+ return self.grpc_channel.close()`
	`760`	`+`
`758`	`761`
`759`	`762`	`__all__ = ("CloudDeployGrpcAsyncIOTransport",)`