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

feat: add context manager support in client #272

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/dlp_v2/services/dlp_service/async_client.py
Expand Up @@ -3146,6 +3146,12 @@ async def finish_dlp_job(
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/dlp_v2/services/dlp_service/client.py
Expand Up @@ -440,10 +440,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 inspect_content(
Expand Down Expand Up @@ -3225,6 +3222,19 @@ def finish_dlp_job(
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
9 changes: 9 additions & 0 deletions google/cloud/dlp_v2/services/dlp_service/transports/base.py
Expand Up @@ -515,6 +515,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 inspect_content(
self,
Expand Down
3 changes: 3 additions & 0 deletions google/cloud/dlp_v2/services/dlp_service/transports/grpc.py
Expand Up @@ -1233,5 +1233,8 @@ def finish_dlp_job(self) -> Callable[[dlp.FinishDlpJobRequest], empty_pb2.Empty]
)
return self._stubs["finish_dlp_job"]

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


__all__ = ("DlpServiceGrpcTransport",)
Expand Up @@ -1265,5 +1265,8 @@ def finish_dlp_job(
)
return self._stubs["finish_dlp_job"]

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


__all__ = ("DlpServiceGrpcAsyncIOTransport",)