Skip to content

Commit 776842a

Browse files
feat: add context manager support in client (#272)
- [ ] Regenerate this pull request now. chore: fix docstring for first attribute of protos committer: @busunkim96 PiperOrigin-RevId: 401271153 Source-Link: googleapis/googleapis@787f8c9 Source-Link: googleapis/googleapis-gen@81decff Copy-Tag: eyJwIjoiLmdpdGh1Yi8uT3dsQm90LnlhbWwiLCJoIjoiODFkZWNmZmU5ZmM3MjM5NmE4MTUzZTc1NmQxZDY3YTZlZWNmZDYyMCJ9
1 parent 06f97e1 commit 776842a

File tree

8 files changed

+213
-15
lines changed

8 files changed

+213
-15
lines changed

packages/google-cloud-dlp/google/cloud/dlp_v2/services/dlp_service/async_client.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3146,6 +3146,12 @@ async def finish_dlp_job(
31463146
request, retry=retry, timeout=timeout, metadata=metadata,
31473147
)
31483148

3149+
async def __aenter__(self):
3150+
return self
3151+
3152+
async def __aexit__(self, exc_type, exc, tb):
3153+
await self.transport.close()
3154+
31493155

31503156
try:
31513157
DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(

packages/google-cloud-dlp/google/cloud/dlp_v2/services/dlp_service/client.py

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -440,10 +440,7 @@ def __init__(
440440
client_cert_source_for_mtls=client_cert_source_func,
441441
quota_project_id=client_options.quota_project_id,
442442
client_info=client_info,
443-
always_use_jwt_access=(
444-
Transport == type(self).get_transport_class("grpc")
445-
or Transport == type(self).get_transport_class("grpc_asyncio")
446-
),
443+
always_use_jwt_access=True,
447444
)
448445

449446
def inspect_content(
@@ -3225,6 +3222,19 @@ def finish_dlp_job(
32253222
request, retry=retry, timeout=timeout, metadata=metadata,
32263223
)
32273224

3225+
def __enter__(self):
3226+
return self
3227+
3228+
def __exit__(self, type, value, traceback):
3229+
"""Releases underlying transport's resources.
3230+
3231+
.. warning::
3232+
ONLY use as a context manager if the transport is NOT shared
3233+
with other clients! Exiting the with block will CLOSE the transport
3234+
and may cause errors in other clients!
3235+
"""
3236+
self.transport.close()
3237+
32283238

32293239
try:
32303240
DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(

packages/google-cloud-dlp/google/cloud/dlp_v2/services/dlp_service/transports/base.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -515,6 +515,15 @@ def _prep_wrapped_messages(self, client_info):
515515
),
516516
}
517517

518+
def close(self):
519+
"""Closes resources associated with the transport.
520+
521+
.. warning::
522+
Only call this method if the transport is NOT shared
523+
with other clients - this may cause errors in other clients!
524+
"""
525+
raise NotImplementedError()
526+
518527
@property
519528
def inspect_content(
520529
self,

packages/google-cloud-dlp/google/cloud/dlp_v2/services/dlp_service/transports/grpc.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1233,5 +1233,8 @@ def finish_dlp_job(self) -> Callable[[dlp.FinishDlpJobRequest], empty_pb2.Empty]
12331233
)
12341234
return self._stubs["finish_dlp_job"]
12351235

1236+
def close(self):
1237+
self.grpc_channel.close()
1238+
12361239

12371240
__all__ = ("DlpServiceGrpcTransport",)

packages/google-cloud-dlp/google/cloud/dlp_v2/services/dlp_service/transports/grpc_asyncio.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1265,5 +1265,8 @@ def finish_dlp_job(
12651265
)
12661266
return self._stubs["finish_dlp_job"]
12671267

1268+
def close(self):
1269+
return self.grpc_channel.close()
1270+
12681271

12691272
__all__ = ("DlpServiceGrpcAsyncIOTransport",)

0 commit comments

Comments
 (0)