Skip to content
This repository has been archived by the owner on Jul 6, 2023. It is now read-only.

feat: add context manager support in client #44

Merged
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
6 changes: 6 additions & 0 deletions google/cloud/eventarc_v1/services/eventarc/async_client.py
Expand Up @@ -616,6 +616,12 @@ async def delete_trigger(
# 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
18 changes: 14 additions & 4 deletions google/cloud/eventarc_v1/services/eventarc/client.py
Expand Up @@ -372,10 +372,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 get_trigger(
Expand Down Expand Up @@ -828,6 +825,19 @@ def delete_trigger(
# 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(
Expand Down
9 changes: 9 additions & 0 deletions google/cloud/eventarc_v1/services/eventarc/transports/base.py
Expand Up @@ -172,6 +172,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
3 changes: 3 additions & 0 deletions google/cloud/eventarc_v1/services/eventarc/transports/grpc.py
Expand Up @@ -374,5 +374,8 @@ def delete_trigger(
)
return self._stubs["delete_trigger"]

def close(self):
self.grpc_channel.close()


__all__ = ("EventarcGrpcTransport",)
Expand Up @@ -383,5 +383,8 @@ def delete_trigger(
)
return self._stubs["delete_trigger"]

def close(self):
return self.grpc_channel.close()


__all__ = ("EventarcGrpcAsyncIOTransport",)
7 changes: 7 additions & 0 deletions google/cloud/eventarc_v1/types/eventarc.py
Expand Up @@ -36,6 +36,7 @@

class GetTriggerRequest(proto.Message):
r"""The request message for the GetTrigger method.

Attributes:
name (str):
Required. The name of the trigger to get.
Expand All @@ -46,6 +47,7 @@ class GetTriggerRequest(proto.Message):

class ListTriggersRequest(proto.Message):
r"""The request message for the ListTriggers method.

Attributes:
parent (str):
Required. The parent collection to list
Expand Down Expand Up @@ -77,6 +79,7 @@ class ListTriggersRequest(proto.Message):

class ListTriggersResponse(proto.Message):
r"""The response message for the ListTriggers method.

Attributes:
triggers (Sequence[google.cloud.eventarc_v1.types.Trigger]):
The requested triggers, up to the number specified in
Expand All @@ -102,6 +105,7 @@ def raw_page(self):

class CreateTriggerRequest(proto.Message):
r"""The request message for the CreateTrigger method.

Attributes:
parent (str):
Required. The parent collection in which to
Expand All @@ -124,6 +128,7 @@ class CreateTriggerRequest(proto.Message):

class UpdateTriggerRequest(proto.Message):
r"""The request message for the UpdateTrigger method.

Attributes:
trigger (google.cloud.eventarc_v1.types.Trigger):
The trigger to be updated.
Expand Down Expand Up @@ -151,6 +156,7 @@ class UpdateTriggerRequest(proto.Message):

class DeleteTriggerRequest(proto.Message):
r"""The request message for the DeleteTrigger method.

Attributes:
name (str):
Required. The name of the trigger to be
Expand All @@ -176,6 +182,7 @@ class DeleteTriggerRequest(proto.Message):

class OperationMetadata(proto.Message):
r"""Represents the metadata of the long-running operation.

Attributes:
create_time (google.protobuf.timestamp_pb2.Timestamp):
Output only. The time the operation was
Expand Down
4 changes: 4 additions & 0 deletions google/cloud/eventarc_v1/types/trigger.py
Expand Up @@ -33,6 +33,7 @@

class Trigger(proto.Message):
r"""A representation of the trigger resource.

Attributes:
name (str):
Required. The resource name of the trigger. Must be unique
Expand Down Expand Up @@ -124,6 +125,7 @@ class EventFilter(proto.Message):

class Destination(proto.Message):
r"""Represents a target of an invocation over HTTP.

Attributes:
cloud_run (google.cloud.eventarc_v1.types.CloudRun):
Cloud Run fully-managed service that receives
Expand Down Expand Up @@ -153,6 +155,7 @@ class Transport(proto.Message):

class CloudRun(proto.Message):
r"""Represents a Cloud Run destination.

Attributes:
service (str):
Required. The name of the Cloud Run service
Expand All @@ -179,6 +182,7 @@ class CloudRun(proto.Message):

class Pubsub(proto.Message):
r"""Represents a Pub/Sub transport.

Attributes:
topic (str):
Optional. The name of the Pub/Sub topic created and managed
Expand Down
50 changes: 50 additions & 0 deletions tests/unit/gapic/eventarc_v1/test_eventarc.py
Expand Up @@ -32,6 +32,7 @@
from google.api_core import grpc_helpers_async
from google.api_core import operation_async # type: ignore
from google.api_core import operations_v1
from google.api_core import path_template
from google.auth import credentials as ga_credentials
from google.auth.exceptions import MutualTLSChannelError
from google.cloud.eventarc_v1.services.eventarc import EventarcAsyncClient
Expand Down Expand Up @@ -1734,6 +1735,9 @@ def test_eventarc_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 @@ -2258,3 +2262,49 @@ def test_client_withDEFAULT_CLIENT_INFO():
credentials=ga_credentials.AnonymousCredentials(), client_info=client_info,
)
prep.assert_called_once_with(client_info)


@pytest.mark.asyncio
async def test_transport_close_async():
client = EventarcAsyncClient(
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 = EventarcClient(
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 = EventarcClient(
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()