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

feat: add context manager support in client #126

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
Expand Up @@ -725,6 +725,12 @@ async def delete_access_approval_settings(
request, retry=retry, timeout=timeout, metadata=metadata,
)

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/accessapproval_v1/services/access_approval/client.py
Expand Up @@ -365,10 +365,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 list_approval_requests(
Expand Down Expand Up @@ -881,6 +878,19 @@ def delete_access_approval_settings(
request, retry=retry, timeout=timeout, metadata=metadata,
)

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
Expand Up @@ -219,6 +219,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 list_approval_requests(
self,
Expand Down
Expand Up @@ -492,5 +492,8 @@ def delete_access_approval_settings(
)
return self._stubs["delete_access_approval_settings"]

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


__all__ = ("AccessApprovalGrpcTransport",)
Expand Up @@ -498,5 +498,8 @@ def delete_access_approval_settings(
)
return self._stubs["delete_access_approval_settings"]

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


__all__ = ("AccessApprovalGrpcAsyncIOTransport",)
13 changes: 13 additions & 0 deletions google/cloud/accessapproval_v1/types/accessapproval.py
Expand Up @@ -53,6 +53,7 @@ class EnrollmentLevel(proto.Enum):

class AccessLocations(proto.Message):
r"""Home office and physical location of the principal.

Attributes:
principal_office_country (str):
The "home office" location of the principal.
Expand Down Expand Up @@ -96,6 +97,7 @@ class AccessLocations(proto.Message):

class AccessReason(proto.Message):
r"""

Attributes:
type_ (google.cloud.accessapproval_v1.types.AccessReason.Type):
Type of access justification.
Expand Down Expand Up @@ -134,6 +136,7 @@ class ApproveDecision(proto.Message):

class DismissDecision(proto.Message):
r"""A decision that has been made to dismiss an approval request.

Attributes:
dismiss_time (google.protobuf.timestamp_pb2.Timestamp):
The time at which the approval request was
Expand All @@ -147,6 +150,7 @@ class DismissDecision(proto.Message):

class ResourceProperties(proto.Message):
r"""The properties associated with the resource of the request.

Attributes:
excludes_descendants (bool):
Whether an approval will exclude the
Expand All @@ -158,6 +162,7 @@ class ResourceProperties(proto.Message):

class ApprovalRequest(proto.Message):
r"""A request for the customer to approve access to a resource.

Attributes:
name (str):
The resource name of the request. Format is
Expand Down Expand Up @@ -296,6 +301,7 @@ class AccessApprovalSettings(proto.Message):

class ListApprovalRequestsMessage(proto.Message):
r"""Request to list approval requests.

Attributes:
parent (str):
The parent resource. This may be "projects/{project_id}",
Expand Down Expand Up @@ -325,6 +331,7 @@ class ListApprovalRequestsMessage(proto.Message):

class ListApprovalRequestsResponse(proto.Message):
r"""Response to listing of ApprovalRequest objects.

Attributes:
approval_requests (Sequence[google.cloud.accessapproval_v1.types.ApprovalRequest]):
Approval request details.
Expand All @@ -345,6 +352,7 @@ def raw_page(self):

class GetApprovalRequestMessage(proto.Message):
r"""Request to get an approval request.

Attributes:
name (str):
Name of the approval request to retrieve.
Expand All @@ -355,6 +363,7 @@ class GetApprovalRequestMessage(proto.Message):

class ApproveApprovalRequestMessage(proto.Message):
r"""Request to approve an ApprovalRequest.

Attributes:
name (str):
Name of the approval request to approve.
Expand All @@ -368,6 +377,7 @@ class ApproveApprovalRequestMessage(proto.Message):

class DismissApprovalRequestMessage(proto.Message):
r"""Request to dismiss an approval request.

Attributes:
name (str):
Name of the ApprovalRequest to dismiss.
Expand All @@ -378,6 +388,7 @@ class DismissApprovalRequestMessage(proto.Message):

class GetAccessApprovalSettingsMessage(proto.Message):
r"""Request to get access approval settings.

Attributes:
name (str):
Name of the AccessApprovalSettings to
Expand All @@ -389,6 +400,7 @@ class GetAccessApprovalSettingsMessage(proto.Message):

class UpdateAccessApprovalSettingsMessage(proto.Message):
r"""Request to update access approval settings.

Attributes:
settings (google.cloud.accessapproval_v1.types.AccessApprovalSettings):
The new AccessApprovalSettings.
Expand All @@ -414,6 +426,7 @@ class UpdateAccessApprovalSettingsMessage(proto.Message):

class DeleteAccessApprovalSettingsMessage(proto.Message):
r"""Request to delete access approval settings.

Attributes:
name (str):
Name of the AccessApprovalSettings to delete.
Expand Down
50 changes: 50 additions & 0 deletions tests/unit/gapic/accessapproval_v1/test_access_approval.py
Expand Up @@ -29,6 +29,7 @@
from google.api_core import gapic_v1
from google.api_core import grpc_helpers
from google.api_core import grpc_helpers_async
from google.api_core import path_template
from google.auth import credentials as ga_credentials
from google.auth.exceptions import MutualTLSChannelError
from google.cloud.accessapproval_v1.services.access_approval import (
Expand Down Expand Up @@ -2228,6 +2229,9 @@ def test_access_approval_base_transport():
with pytest.raises(NotImplementedError):
getattr(transport, method)(request=object())

with pytest.raises(NotImplementedError):
transport.close()


@requires_google_auth_gte_1_25_0
def test_access_approval_base_transport_with_credentials_file():
Expand Down Expand Up @@ -2686,3 +2690,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 = AccessApprovalAsyncClient(
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 = AccessApprovalClient(
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 = AccessApprovalClient(
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()