feat: add support for context manager in client (#987)

* feat: add support for context manager in client.

* chore: remove extra whitespace.

* chore: adds autogenerated unit tests.

* chore: adds stronger warning.

* chore: fixes tests.

* chore: adds auto-generated tests for ads.

* chore: updates golden files.

* chore: refactor.

* chore: refactor.

* feat: adds close() to transport and ctx to async client.

* feat: adds close method and removes ctx from transport in ads.

* chore: adds warning infobox to docstring.

* chore: updates integration tests.

* chore: fixes typo.

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):
+        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()
 
     {% 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,14 @@ class {{ service.name }}Transport(metaclass=abc.ABCMeta):
             {% endfor %} {# precomputed wrappers loop #}
         }
 
+    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()
 
     {% if service.has_lro %}
 

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

@@ -188,6 +188,9 @@ class {{ service.name }}GrpcTransport({{ service.name }}Transport):
             **kwargs
         )
 
+    def close(self):
+        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,9 @@ def test_{{ service.name|snake_case }}_base_transport():
         with pytest.raises(NotImplementedError):
             getattr(transport, method)(request=object())
 
+    with pytest.raises(NotImplementedError):
+        transport.close()
+
     {% if service.has_lro %}
     # Additionally, the LRO client (a property) should
     # also raise NotImplementedError
@@ -903,5 +906,26 @@ def test_client_withDEFAULT_CLIENT_INFO():
         )
         prep.assert_called_once_with(client_info)
 
+def test_grpc_transport_close():
+    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_ctx():
+    client = {{ service.client_name }}(
+        credentials=ga_credentials.AnonymousCredentials(),
+        transport='grpc',
+    )
+    # Test client calls underlying transport.
+    with mock.patch.object(type(client.transport), "close") as close:
+        close.assert_not_called()
+        with client as _:
+            pass
+        close.assert_called()
 
 {% endblock %}

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

@@ -599,6 +599,12 @@ class {{ service.async_client_name }}:
         return response
     {% endif %}
 
+    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(
         gapic_version=pkg_resources.get_distribution(

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

@@ -477,6 +477,19 @@ class {{ service.client_name }}(metaclass={{ service.client_name }}Meta):
         {{ "\n" }}
     {% endfor %}
 
+    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()
+
     {% if opts.add_iam_methods %}
     def set_iam_policy(
         self,

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

@@ -174,6 +174,14 @@ class {{ service.name }}Transport(abc.ABC):
             {% endfor %} {# precomputed wrappers loop #}
         }
 
+    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()
 
     {% if service.has_lro %}
 

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

@@ -224,6 +224,7 @@ class {{ service.name }}GrpcTransport({{ service.name }}Transport):
         """Return the channel designed to connect to this service.
         """
         return self._grpc_channel
+
     {% if service.has_lro %}
 
     @property
@@ -355,6 +356,9 @@ class {{ service.name }}GrpcTransport({{ service.name }}Transport):
         return self._stubs["test_iam_permissions"]
     {% endif %}
 
+    def close(self):
+        self.grpc_channel.close()
+
 __all__ = (
     '{{ service.name }}GrpcTransport',
 )

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

@@ -359,6 +359,10 @@ class {{ service.grpc_asyncio_transport_name }}({{ service.name }}Transport):
         return self._stubs["test_iam_permissions"]
     {% endif %}
 
+    def close(self):
+        return self.grpc_channel.close()
+
+
 __all__ = (
     '{{ service.name }}GrpcAsyncIOTransport',
 )

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

@@ -283,6 +283,9 @@ class {{service.name}}RestTransport({{service.name}}Transport):
         return self._{{method.name | snake_case}}
     {%- endfor %}
 
+    def close(self):
+        self._session.close()
+
 
 __all__=(
     '{{ service.name }}RestTransport',

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

@@ -1516,6 +1516,9 @@ def test_{{ service.name|snake_case }}_base_transport():
         with pytest.raises(NotImplementedError):
             getattr(transport, method)(request=object())
 
+    with pytest.raises(NotImplementedError):
+        transport.close()
+
     {% if service.has_lro %}
     # Additionally, the LRO client (a property) should
     # also raise NotImplementedError
@@ -2487,4 +2490,56 @@ async def test_test_iam_permissions_from_dict_async():
 
 {% endif %}
 
+@pytest.mark.asyncio
+async def test_transport_close_async():
+    client = {{ service.async_client_name }}(
+        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 = {
+        {% if 'rest' in opts.transport %}
+        "rest": "_session",
+        {% endif %}
+        {% if 'grpc' in opts.transport %}
+        "grpc": "_grpc_channel",
+        {% endif %}
+    }
+
+    for transport, close_name in transports.items():
+        client = {{ service.client_name }}(
+            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 = [
+        {% if 'rest' in opts.transport %}
+        'rest',
+        {% endif %}
+        {% if 'grpc' in opts.transport %}
+        'grpc',
+        {% endif %}
+    ]
+    for transport in transports:
+        client = {{ service.client_name }}(
+            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()
+
 {% endblock %}

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

@@ -1292,9 +1292,11 @@ async def analyze_iam_policy_longrunning(self,
         # 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(

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

@@ -1445,7 +1445,18 @@ def analyze_iam_policy_longrunning(self,
         # 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()
 
 
 

tests/integration/goldens/asset/google/cloud/asset_v1/services/asset_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."""

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

@@ -592,6 +592,8 @@ def analyze_iam_policy_longrunning(self) -> Callable[
             )
         return self._stubs['analyze_iam_policy_longrunning']
 
+    def close(self):
+        self.grpc_channel.close()
 
 __all__ = (
     'AssetServiceGrpcTransport',

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

@@ -596,6 +596,9 @@ def analyze_iam_policy_longrunning(self) -> Callable[
             )
         return self._stubs['analyze_iam_policy_longrunning']
 
+    def close(self):
+        return self.grpc_channel.close()
+
 
 __all__ = (
     'AssetServiceGrpcAsyncIOTransport',

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

@@ -3551,6 +3551,9 @@ def test_asset_service_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):
@@ -4049,3 +4052,46 @@ def test_client_withDEFAULT_CLIENT_INFO():
             client_info=client_info,
         )
         prep.assert_called_once_with(client_info)
+
+
+@pytest.mark.asyncio
+async def test_transport_close_async():
+    client = AssetServiceAsyncClient(
+        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 = AssetServiceClient(
+            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 = AssetServiceClient(
+            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()

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

@@ -644,9 +644,11 @@ async def sign_jwt(self,
         # 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(

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

@@ -804,7 +804,18 @@ def sign_jwt(self,
         # 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()