Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: add support for context manager in client #987

Merged
merged 15 commits into from
Oct 4, 2021
Merged
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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 %}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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 _:
atulep marked this conversation as resolved.
Show resolved Hide resolved
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 %}
Original file line number Diff line number Diff line change
Expand Up @@ -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()
atulep marked this conversation as resolved.
Show resolved Hide resolved

try:
DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(
gapic_version=pkg_resources.get_distribution(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -477,6 +477,19 @@ class {{ service.client_name }}(metaclass={{ service.client_name }}Meta):
{{ "\n" }}
{% endfor %}

def __enter__(self):
return self
atulep marked this conversation as resolved.
Show resolved Hide resolved

def __exit__(self, type, value, traceback):
atulep marked this conversation as resolved.
Show resolved Hide resolved
"""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,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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 %}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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',
)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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',
)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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',
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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 %}
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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()



Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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',
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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',
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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()
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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()



Expand Down
Loading