feat: add context manager support in client

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(

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(

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,

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",)

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",)

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(

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(

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."""

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",)

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",)

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(

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(

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."""

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",)

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",)

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(

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(

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."""

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",)

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",)

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