feat: add support for context manager in client

gapic/ads-templates/%namespace/%name/%version/%sub/services/%service/client.py.j2

@@ -148,6 +148,18 @@ class {{ service.client_name }}(metaclass={{ service.client_name }}Meta):
         """
         return self._transport
 
+    def __enter__(self):
+        self.transport.__enter__()
+        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!
+        """
+        return self.transport.__exit__(type, value, traceback)
 
     {% for message in service.resource_messages|sort(attribute="resource_type") %}
     @staticmethod

gapic/ads-templates/%namespace/%name/%version/%sub/services/%service/transports/base.py.j2

@@ -103,6 +103,11 @@ class {{ service.name }}Transport(metaclass=abc.ABCMeta):
             {% endfor %} {# precomputed wrappers loop #}
         }
 
+    def __enter__(self):
+        raise NotImplementedError()
+
+    def __exit__(self, type, value, traceback):
+        raise NotImplementedError()
 
     {% if service.has_lro %}
 

gapic/ads-templates/%namespace/%name/%version/%sub/services/%service/transports/grpc.py.j2

@@ -188,6 +188,14 @@ class {{ service.name }}GrpcTransport({{ service.name }}Transport):
             **kwargs
         )
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, type, value, traceback):
+        """Closes underlying gRPC channel.
+        """
+        self.grpc_channel.close()
+
     @property
     def grpc_channel(self) -> grpc.Channel:
         """Return the channel designed to connect to this service.

gapic/ads-templates/tests/unit/gapic/%name_%version/%sub/test_%service.py.j2

@@ -654,6 +654,11 @@ def test_{{ service.name|snake_case }}_base_transport():
         with pytest.raises(NotImplementedError):
             getattr(transport, method)(request=object())
 
+    with pytest.raises(NotImplementedError):
+        transport.__enter__()
+    with pytest.raises(NotImplementedError):
+        transport.__exit__(None, None, None)
+
     {% if service.has_lro %}
     # Additionally, the LRO client (a property) should
     # also raise NotImplementedError
@@ -903,5 +908,28 @@ def test_client_withDEFAULT_CLIENT_INFO():
         )
         prep.assert_called_once_with(client_info)
 
+def test_grpc_transport_enter_exit():
+    client = {{ service.client_name }}(
+        credentials=ga_credentials.AnonymousCredentials(),
+        transport='grpc',
+    )
+    with mock.patch.object(type(client.transport._grpc_channel), 'close') as chan_close:
+        with client as _:
+            chan_close.assert_not_called()
+        chan_close.assert_called_once()
+
+def test_grpc_client_enter_exit():
+    client = {{ service.client_name }}(
+        credentials=ga_credentials.AnonymousCredentials(),
+        transport='grpc',
+    )
+    # Test client calls underlying transport.
+    with mock.patch.object(type(client.transport), "__enter__") as enter:
+        with mock.patch.object(type(client.transport), "__exit__") as exit:
+            enter.assert_not_called()
+            exit.assert_not_called()
+            with client as _:
+                enter.assert_called_once()
+            exit.assert_called()
 
 {% endblock %}

gapic/templates/%namespace/%name_%version/%sub/services/%service/client.py.j2

@@ -487,7 +487,9 @@ class {{ service.client_name }}(metaclass={{ service.client_name }}Meta):
     def __exit__(self, type, value, traceback):
         """Releases underlying transport's resources.
 
-        WARNING: This makes the transport unusable.
+        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!
         """
         return self.transport.__exit__(type, value, traceback)
 

gapic/templates/tests/unit/gapic/%name_%version/%sub/test_%service.py.j2

@@ -1508,6 +1508,11 @@ def test_{{ service.name|snake_case }}_base_transport():
         with pytest.raises(NotImplementedError):
             getattr(transport, method)(request=object())
 
+    with pytest.raises(NotImplementedError):
+        transport.__enter__()
+    with pytest.raises(NotImplementedError):
+        transport.__exit__(None, None, None)
+
     {% if service.has_lro %}
     # Additionally, the LRO client (a property) should
     # also raise NotImplementedError
@@ -2459,4 +2464,56 @@ async def test_test_iam_permissions_from_dict_async():
 
 {% endif %}
 
+{% if 'grpc' in opts.transport %}
+def test_grpc_transport_enter_exit():
+    client = {{ service.client_name }}(
+        credentials=ga_credentials.AnonymousCredentials(),
+        transport='grpc',
+    )
+    with mock.patch.object(type(client.transport._grpc_channel), 'close') as chan_close:
+        with client as _:
+            chan_close.assert_not_called()
+        chan_close.assert_called_once()
+
+def test_grpc_client_enter_exit():
+    client = {{ service.client_name }}(
+        credentials=ga_credentials.AnonymousCredentials(),
+        transport='grpc',
+    )
+    # Test client calls underlying transport.
+    with mock.patch.object(type(client.transport), "__enter__") as enter:
+        with mock.patch.object(type(client.transport), "__exit__") as exit:
+            enter.assert_not_called()
+            exit.assert_not_called()
+            with client as _:
+                enter.assert_called_once()
+            exit.assert_called()
+{% endif %}
+
+{% if 'rest' in opts.transport %}
+def test_rest_transport_enter_exit():
+    client = {{ service.client_name }}(
+        credentials=ga_credentials.AnonymousCredentials(),
+        transport='rest',
+    )
+    with mock.patch.object(type(client.transport._session), 'close') as sess_close:
+        with client as _:
+            sess_close.assert_not_called()
+        sess_close.assert_called_once()
+
+def test_rest_client_enter_exit():
+    client = {{ service.client_name }}(
+        credentials=ga_credentials.AnonymousCredentials(),
+        transport='rest',
+    )
+    # Test client calls underlying transport.
+    with mock.patch.object(type(client.transport), "__enter__") as enter:
+        with mock.patch.object(type(client.transport), "__exit__") as exit:
+            enter.assert_not_called()
+            exit.assert_not_called()
+            with client as _:
+                enter.assert_called_once()
+            exit.assert_called()
+{% endif %}
+
 {% endblock %}

tests/integration/goldens/asset/google/cloud/asset_v1/services/asset_service/client.py

@@ -1448,7 +1448,18 @@ def analyze_iam_policy_longrunning(self,
         # Done; return the response.
         return response
 
+    def __enter__(self):
+        self.transport.__enter__()
+        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!
+        """
+        return self.transport.__exit__(type, value, traceback)
 
 
 

tests/integration/goldens/asset/google/cloud/asset_v1/services/asset_service/transports/base.py

@@ -259,6 +259,12 @@ def _prep_wrapped_messages(self, client_info):
             ),
          }
 
+    def __enter__(self):
+        raise NotImplementedError()
+
+    def __exit__(self, type, value, traceback):
+        raise NotImplementedError()
+
     @property
     def operations_client(self) -> operations_v1.OperationsClient:
         """Return the client designed to process long-running operations."""

tests/integration/goldens/asset/google/cloud/asset_v1/services/asset_service/transports/grpc.py

@@ -224,6 +224,14 @@ def grpc_channel(self) -> grpc.Channel:
         """
         return self._grpc_channel
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, type, value, traceback):
+        """Closes underlying gRPC channel.
+        """
+        self.grpc_channel.close()
+
     @property
     def operations_client(self) -> operations_v1.OperationsClient:
         """Create the client designed to process long-running operations.

tests/integration/goldens/asset/tests/unit/gapic/asset_v1/test_asset_service.py

@@ -3550,6 +3550,11 @@ def test_asset_service_base_transport():
         with pytest.raises(NotImplementedError):
             getattr(transport, method)(request=object())
 
+    with pytest.raises(NotImplementedError):
+        transport.__enter__()
+    with pytest.raises(NotImplementedError):
+        transport.__exit__(None, None, None)
+
     # Additionally, the LRO client (a property) should
     # also raise NotImplementedError
     with pytest.raises(NotImplementedError):
@@ -4048,3 +4053,28 @@ def test_client_withDEFAULT_CLIENT_INFO():
             client_info=client_info,
         )
         prep.assert_called_once_with(client_info)
+
+
+def test_grpc_transport_enter_exit():
+    client = AssetServiceClient(
+        credentials=ga_credentials.AnonymousCredentials(),
+        transport='grpc',
+    )
+    with mock.patch.object(type(client.transport._grpc_channel), 'close') as chan_close:
+        with client as _:
+            chan_close.assert_not_called()
+        chan_close.assert_called_once()
+
+def test_grpc_client_enter_exit():
+    client = AssetServiceClient(
+        credentials=ga_credentials.AnonymousCredentials(),
+        transport='grpc',
+    )
+    # Test client calls underlying transport.
+    with mock.patch.object(type(client.transport), "__enter__") as enter:
+        with mock.patch.object(type(client.transport), "__exit__") as exit:
+            enter.assert_not_called()
+            exit.assert_not_called()
+            with client as _:
+                enter.assert_called_once()
+            exit.assert_called()

tests/integration/goldens/credentials/google/iam/credentials_v1/services/iam_credentials/client.py

@@ -807,7 +807,18 @@ def sign_jwt(self,
         # Done; return the response.
         return response
 
+    def __enter__(self):
+        self.transport.__enter__()
+        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!
+        """
+        return self.transport.__exit__(type, value, traceback)
 
 
 

tests/integration/goldens/credentials/google/iam/credentials_v1/services/iam_credentials/transports/base.py

@@ -196,6 +196,12 @@ def _prep_wrapped_messages(self, client_info):
             ),
          }
 
+    def __enter__(self):
+        raise NotImplementedError()
+
+    def __exit__(self, type, value, traceback):
+        raise NotImplementedError()
+
     @property
     def generate_access_token(self) -> Callable[
             [common.GenerateAccessTokenRequest],

tests/integration/goldens/credentials/google/iam/credentials_v1/services/iam_credentials/transports/grpc.py

@@ -229,6 +229,14 @@ def grpc_channel(self) -> grpc.Channel:
         """
         return self._grpc_channel
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, type, value, traceback):
+        """Closes underlying gRPC channel.
+        """
+        self.grpc_channel.close()
+
     @property
     def generate_access_token(self) -> Callable[
             [common.GenerateAccessTokenRequest],

tests/integration/goldens/credentials/tests/unit/gapic/credentials_v1/test_iam_credentials.py

@@ -1484,6 +1484,11 @@ def test_iam_credentials_base_transport():
         with pytest.raises(NotImplementedError):
             getattr(transport, method)(request=object())
 
+    with pytest.raises(NotImplementedError):
+        transport.__enter__()
+    with pytest.raises(NotImplementedError):
+        transport.__exit__(None, None, None)
+
 
 @requires_google_auth_gte_1_25_0
 def test_iam_credentials_base_transport_with_credentials_file():
@@ -1928,3 +1933,28 @@ def test_client_withDEFAULT_CLIENT_INFO():
             client_info=client_info,
         )
         prep.assert_called_once_with(client_info)
+
+
+def test_grpc_transport_enter_exit():
+    client = IAMCredentialsClient(
+        credentials=ga_credentials.AnonymousCredentials(),
+        transport='grpc',
+    )
+    with mock.patch.object(type(client.transport._grpc_channel), 'close') as chan_close:
+        with client as _:
+            chan_close.assert_not_called()
+        chan_close.assert_called_once()
+
+def test_grpc_client_enter_exit():
+    client = IAMCredentialsClient(
+        credentials=ga_credentials.AnonymousCredentials(),
+        transport='grpc',
+    )
+    # Test client calls underlying transport.
+    with mock.patch.object(type(client.transport), "__enter__") as enter:
+        with mock.patch.object(type(client.transport), "__exit__") as exit:
+            enter.assert_not_called()
+            exit.assert_not_called()
+            with client as _:
+                enter.assert_called_once()
+            exit.assert_called()

tests/integration/goldens/logging/google/cloud/logging_v2/services/config_service_v2/client.py

@@ -2179,7 +2179,18 @@ def update_cmek_settings(self,
         # Done; return the response.
         return response
 
+    def __enter__(self):
+        self.transport.__enter__()
+        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!
+        """
+        return self.transport.__exit__(type, value, traceback)
 
 
 

tests/integration/goldens/logging/google/cloud/logging_v2/services/config_service_v2/transports/base.py

@@ -323,6 +323,12 @@ def _prep_wrapped_messages(self, client_info):
             ),
          }
 
+    def __enter__(self):
+        raise NotImplementedError()
+
+    def __exit__(self, type, value, traceback):
+        raise NotImplementedError()
+
     @property
     def list_buckets(self) -> Callable[
             [logging_config.ListBucketsRequest],

tests/integration/goldens/logging/google/cloud/logging_v2/services/config_service_v2/transports/grpc.py

@@ -221,6 +221,14 @@ def grpc_channel(self) -> grpc.Channel:
         """
         return self._grpc_channel
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, type, value, traceback):
+        """Closes underlying gRPC channel.
+        """
+        self.grpc_channel.close()
+
     @property
     def list_buckets(self) -> Callable[
             [logging_config.ListBucketsRequest],