feat: add context manager support in client (#285)

gcf-owl-bot[bot] · web-flow · commit b54fb7647dea · 2021-10-08T14:34:11.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/dataproc_v1/services/autoscaling_policy_service/async_client.py b/google/cloud/dataproc_v1/services/autoscaling_policy_service/async_client.py
@@ -621,6 +621,12 @@ async def delete_autoscaling_policy(
             request, retry=retry, timeout=timeout, metadata=metadata,
         )
 
+    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/dataproc_v1/services/autoscaling_policy_service/client.py b/google/cloud/dataproc_v1/services/autoscaling_policy_service/client.py
@@ -350,10 +350,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 create_autoscaling_policy(
@@ -791,6 +788,19 @@ def delete_autoscaling_policy(
             request, retry=retry, timeout=timeout, metadata=metadata,
         )
 
+    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/dataproc_v1/services/autoscaling_policy_service/transports/base.py b/google/cloud/dataproc_v1/services/autoscaling_policy_service/transports/base.py
@@ -210,6 +210,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 create_autoscaling_policy(
         self,
diff --git a/google/cloud/dataproc_v1/services/autoscaling_policy_service/transports/grpc.py b/google/cloud/dataproc_v1/services/autoscaling_policy_service/transports/grpc.py
@@ -376,5 +376,8 @@ def delete_autoscaling_policy(
             )
         return self._stubs["delete_autoscaling_policy"]
 
+    def close(self):
+        self.grpc_channel.close()
+
 
 __all__ = ("AutoscalingPolicyServiceGrpcTransport",)
diff --git a/google/cloud/dataproc_v1/services/autoscaling_policy_service/transports/grpc_asyncio.py b/google/cloud/dataproc_v1/services/autoscaling_policy_service/transports/grpc_asyncio.py
@@ -380,5 +380,8 @@ def delete_autoscaling_policy(
             )
         return self._stubs["delete_autoscaling_policy"]
 
+    def close(self):
+        return self.grpc_channel.close()
+
 
 __all__ = ("AutoscalingPolicyServiceGrpcAsyncIOTransport",)
diff --git a/google/cloud/dataproc_v1/services/cluster_controller/async_client.py b/google/cloud/dataproc_v1/services/cluster_controller/async_client.py
@@ -1015,6 +1015,12 @@ async def diagnose_cluster(
         # 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/dataproc_v1/services/cluster_controller/client.py b/google/cloud/dataproc_v1/services/cluster_controller/client.py
@@ -369,10 +369,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 create_cluster(
@@ -1164,6 +1161,19 @@ def diagnose_cluster(
         # 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/dataproc_v1/services/cluster_controller/transports/base.py b/google/cloud/dataproc_v1/services/cluster_controller/transports/base.py
@@ -250,6 +250,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/dataproc_v1/services/cluster_controller/transports/grpc.py b/google/cloud/dataproc_v1/services/cluster_controller/transports/grpc.py
@@ -467,5 +467,8 @@ def diagnose_cluster(
             )
         return self._stubs["diagnose_cluster"]
 
+    def close(self):
+        self.grpc_channel.close()
+
 
 __all__ = ("ClusterControllerGrpcTransport",)
diff --git a/google/cloud/dataproc_v1/services/cluster_controller/transports/grpc_asyncio.py b/google/cloud/dataproc_v1/services/cluster_controller/transports/grpc_asyncio.py
@@ -478,5 +478,8 @@ def diagnose_cluster(
             )
         return self._stubs["diagnose_cluster"]
 
+    def close(self):
+        return self.grpc_channel.close()
+
 
 __all__ = ("ClusterControllerGrpcAsyncIOTransport",)
diff --git a/google/cloud/dataproc_v1/services/job_controller/async_client.py b/google/cloud/dataproc_v1/services/job_controller/async_client.py
@@ -798,6 +798,12 @@ async def delete_job(
             request, retry=retry, timeout=timeout, metadata=metadata,
         )
 
+    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/dataproc_v1/services/job_controller/client.py b/google/cloud/dataproc_v1/services/job_controller/client.py
@@ -328,10 +328,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 submit_job(
@@ -903,6 +900,19 @@ def delete_job(
             request, retry=retry, timeout=timeout, metadata=metadata,
         )
 
+    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/dataproc_v1/services/job_controller/transports/base.py b/google/cloud/dataproc_v1/services/job_controller/transports/base.py
@@ -261,6 +261,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/dataproc_v1/services/job_controller/transports/grpc.py b/google/cloud/dataproc_v1/services/job_controller/transports/grpc.py
@@ -419,5 +419,8 @@ def delete_job(self) -> Callable[[jobs.DeleteJobRequest], empty_pb2.Empty]:
             )
         return self._stubs["delete_job"]
 
+    def close(self):
+        self.grpc_channel.close()
+
 
 __all__ = ("JobControllerGrpcTransport",)
diff --git a/google/cloud/dataproc_v1/services/job_controller/transports/grpc_asyncio.py b/google/cloud/dataproc_v1/services/job_controller/transports/grpc_asyncio.py
@@ -428,5 +428,8 @@ def delete_job(
             )
         return self._stubs["delete_job"]
 
+    def close(self):
+        return self.grpc_channel.close()
+
 
 __all__ = ("JobControllerGrpcAsyncIOTransport",)
diff --git a/google/cloud/dataproc_v1/services/workflow_template_service/async_client.py b/google/cloud/dataproc_v1/services/workflow_template_service/async_client.py
@@ -948,6 +948,12 @@ async def delete_workflow_template(
             request, retry=retry, timeout=timeout, metadata=metadata,
         )
 
+    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/dataproc_v1/services/workflow_template_service/client.py b/google/cloud/dataproc_v1/services/workflow_template_service/client.py
@@ -387,10 +387,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 create_workflow_template(
@@ -1106,6 +1103,19 @@ def delete_workflow_template(
             request, retry=retry, timeout=timeout, metadata=metadata,
         )
 
+    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/dataproc_v1/services/workflow_template_service/transports/base.py b/google/cloud/dataproc_v1/services/workflow_template_service/transports/base.py
@@ -259,6 +259,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/dataproc_v1/services/workflow_template_service/transports/grpc.py b/google/cloud/dataproc_v1/services/workflow_template_service/transports/grpc.py
@@ -500,5 +500,8 @@ def delete_workflow_template(
             )
         return self._stubs["delete_workflow_template"]
 
+    def close(self):
+        self.grpc_channel.close()
+
 
 __all__ = ("WorkflowTemplateServiceGrpcTransport",)
diff --git a/google/cloud/dataproc_v1/services/workflow_template_service/transports/grpc_asyncio.py b/google/cloud/dataproc_v1/services/workflow_template_service/transports/grpc_asyncio.py
@@ -507,5 +507,8 @@ def delete_workflow_template(
             )
         return self._stubs["delete_workflow_template"]
 
+    def close(self):
+        return self.grpc_channel.close()
+
 
 __all__ = ("WorkflowTemplateServiceGrpcAsyncIOTransport",)
diff --git a/google/cloud/dataproc_v1/types/autoscaling_policies.py b/google/cloud/dataproc_v1/types/autoscaling_policies.py
@@ -84,6 +84,7 @@ class AutoscalingPolicy(proto.Message):
 
 class BasicAutoscalingAlgorithm(proto.Message):
     r"""Basic algorithm for autoscaling.
+
     Attributes:
         yarn_config (google.cloud.dataproc_v1.types.BasicYarnAutoscalingConfig):
             Required. YARN autoscaling configuration.
@@ -105,6 +106,7 @@ class BasicAutoscalingAlgorithm(proto.Message):
 
 class BasicYarnAutoscalingConfig(proto.Message):
     r"""Basic autoscaling configurations for YARN.
+
     Attributes:
         graceful_decommission_timeout (google.protobuf.duration_pb2.Duration):
             Required. Timeout for YARN graceful decommissioning of Node
@@ -216,6 +218,7 @@ class InstanceGroupAutoscalingPolicyConfig(proto.Message):
 
 class CreateAutoscalingPolicyRequest(proto.Message):
     r"""A request to create an autoscaling policy.
+
     Attributes:
         parent (str):
             Required. The "resource name" of the region or location, as
@@ -239,6 +242,7 @@ class CreateAutoscalingPolicyRequest(proto.Message):
 
 class GetAutoscalingPolicyRequest(proto.Message):
     r"""A request to fetch an autoscaling policy.
+
     Attributes:
         name (str):
             Required. The "resource name" of the autoscaling policy, as
@@ -259,6 +263,7 @@ class GetAutoscalingPolicyRequest(proto.Message):
 
 class UpdateAutoscalingPolicyRequest(proto.Message):
     r"""A request to update an autoscaling policy.
+
     Attributes:
         policy (google.cloud.dataproc_v1.types.AutoscalingPolicy):
             Required. The updated autoscaling policy.
@@ -292,6 +297,7 @@ class DeleteAutoscalingPolicyRequest(proto.Message):
 
 class ListAutoscalingPoliciesRequest(proto.Message):
     r"""A request to list autoscaling policies in a project.
+
     Attributes:
         parent (str):
             Required. The "resource name" of the region or location, as
diff --git a/google/cloud/dataproc_v1/types/clusters.py b/google/cloud/dataproc_v1/types/clusters.py
diff --git a/google/cloud/dataproc_v1/types/jobs.py b/google/cloud/dataproc_v1/types/jobs.py
diff --git a/google/cloud/dataproc_v1/types/operations.py b/google/cloud/dataproc_v1/types/operations.py
diff --git a/google/cloud/dataproc_v1/types/workflow_templates.py b/google/cloud/dataproc_v1/types/workflow_templates.py
diff --git a/tests/unit/gapic/dataproc_v1/test_autoscaling_policy_service.py b/tests/unit/gapic/dataproc_v1/test_autoscaling_policy_service.py
diff --git a/tests/unit/gapic/dataproc_v1/test_cluster_controller.py b/tests/unit/gapic/dataproc_v1/test_cluster_controller.py
diff --git a/tests/unit/gapic/dataproc_v1/test_job_controller.py b/tests/unit/gapic/dataproc_v1/test_job_controller.py
diff --git a/tests/unit/gapic/dataproc_v1/test_workflow_template_service.py b/tests/unit/gapic/dataproc_v1/test_workflow_template_service.py

Original file line number	Diff line number	Diff line change
`@@ -376,5 +376,8 @@ def delete_autoscaling_policy(`
`376`	`376`	`)`
`377`	`377`	`return self._stubs["delete_autoscaling_policy"]`
`378`	`378`
	`379`	`+ def close(self):`
	`380`	`+ self.grpc_channel.close()`
	`381`	`+`
`379`	`382`
`380`	`383`	`__all__ = ("AutoscalingPolicyServiceGrpcTransport",)`
Original file line number	Diff line number	Diff line change
`@@ -380,5 +380,8 @@ def delete_autoscaling_policy(`
`380`	`380`	`)`
`381`	`381`	`return self._stubs["delete_autoscaling_policy"]`
`382`	`382`
	`383`	`+ def close(self):`
	`384`	`+ return self.grpc_channel.close()`
	`385`	`+`
`383`	`386`
`384`	`387`	`__all__ = ("AutoscalingPolicyServiceGrpcAsyncIOTransport",)`
Original file line number	Diff line number	Diff line change
`@@ -467,5 +467,8 @@ def diagnose_cluster(`
`467`	`467`	`)`
`468`	`468`	`return self._stubs["diagnose_cluster"]`
`469`	`469`
	`470`	`+ def close(self):`
	`471`	`+ self.grpc_channel.close()`
	`472`	`+`
`470`	`473`
`471`	`474`	`__all__ = ("ClusterControllerGrpcTransport",)`
Original file line number	Diff line number	Diff line change
`@@ -478,5 +478,8 @@ def diagnose_cluster(`
`478`	`478`	`)`
`479`	`479`	`return self._stubs["diagnose_cluster"]`
`480`	`480`
	`481`	`+ def close(self):`
	`482`	`+ return self.grpc_channel.close()`
	`483`	`+`
`481`	`484`
`482`	`485`	`__all__ = ("ClusterControllerGrpcAsyncIOTransport",)`