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.

Author: atulep

packages/gapic-generator/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

packages/gapic-generator/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 %}
 

packages/gapic-generator/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.

packages/gapic-generator/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 %}

packages/gapic-generator/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(

packages/gapic-generator/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,

packages/gapic-generator/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 %}
 

packages/gapic-generator/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',
 )

packages/gapic-generator/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',
 )

packages/gapic-generator/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',