Skip to content
This repository was archived by the owner on Nov 29, 2023. It is now read-only.

Commit 4bb0fb6

Browse files
gcf-owl-bot[bot]gcf-owl-bot[bot]
authored
feat: add context manager support in client (#144)
- [ ] 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: https://github.com/googleapis/googleapis-gen/commit/81decffe9fc72396a8153e756d1d67a6eecfd620 Copy-Tag: eyJwIjoiLmdpdGh1Yi8uT3dsQm90LnlhbWwiLCJoIjoiODFkZWNmZmU5ZmM3MjM5NmE4MTUzZTc1NmQxZDY3YTZlZWNmZDYyMCJ9
1 parent 0ba2f84 commit 4bb0fb6

File tree

12 files changed

+170
-8
lines changed

12 files changed

+170
-8
lines changed

‎google/cloud/scheduler_v1/services/cloud_scheduler/async_client.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -846,6 +846,12 @@ async def run_job(
846846
# Done; return the response.
847847
return response
848848

849+
async def __aenter__(self):
850+
return self
851+
852+
async def __aexit__(self, exc_type, exc, tb):
853+
await self.transport.close()
854+
849855

850856
try:
851857
DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(

‎google/cloud/scheduler_v1/services/cloud_scheduler/client.py

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -364,10 +364,7 @@ def __init__(
364364
client_cert_source_for_mtls=client_cert_source_func,
365365
quota_project_id=client_options.quota_project_id,
366366
client_info=client_info,
367-
always_use_jwt_access=(
368-
Transport == type(self).get_transport_class("grpc")
369-
or Transport == type(self).get_transport_class("grpc_asyncio")
370-
),
367+
always_use_jwt_access=True,
371368
)
372369

373370
def list_jobs(
@@ -1014,6 +1011,19 @@ def run_job(
10141011
# Done; return the response.
10151012
return response
10161013

1014+
def __enter__(self):
1015+
return self
1016+
1017+
def __exit__(self, type, value, traceback):
1018+
"""Releases underlying transport's resources.
1019+
1020+
.. warning::
1021+
ONLY use as a context manager if the transport is NOT shared
1022+
with other clients! Exiting the with block will CLOSE the transport
1023+
and may cause errors in other clients!
1024+
"""
1025+
self.transport.close()
1026+
10171027

10181028
try:
10191029
DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(

‎google/cloud/scheduler_v1/services/cloud_scheduler/transports/base.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -217,6 +217,15 @@ def _prep_wrapped_messages(self, client_info):
217217
),
218218
}
219219

220+
def close(self):
221+
"""Closes resources associated with the transport.
222+
223+
.. warning::
224+
Only call this method if the transport is NOT shared
225+
with other clients - this may cause errors in other clients!
226+
"""
227+
raise NotImplementedError()
228+
220229
@property
221230
def list_jobs(
222231
self,

‎google/cloud/scheduler_v1/services/cloud_scheduler/transports/grpc.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -459,5 +459,8 @@ def run_job(self) -> Callable[[cloudscheduler.RunJobRequest], job.Job]:
459459
)
460460
return self._stubs["run_job"]
461461

462+
def close(self):
463+
self.grpc_channel.close()
464+
462465

463466
__all__ = ("CloudSchedulerGrpcTransport",)

‎google/cloud/scheduler_v1/services/cloud_scheduler/transports/grpc_asyncio.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -472,5 +472,8 @@ def run_job(self) -> Callable[[cloudscheduler.RunJobRequest], Awaitable[job.Job]
472472
)
473473
return self._stubs["run_job"]
474474

475+
def close(self):
476+
return self.grpc_channel.close()
477+
475478

476479
__all__ = ("CloudSchedulerGrpcAsyncIOTransport",)

‎google/cloud/scheduler_v1beta1/services/cloud_scheduler/async_client.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -867,6 +867,12 @@ async def run_job(
867867
# Done; return the response.
868868
return response
869869

870+
async def __aenter__(self):
871+
return self
872+
873+
async def __aexit__(self, exc_type, exc, tb):
874+
await self.transport.close()
875+
870876

871877
try:
872878
DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(

‎google/cloud/scheduler_v1beta1/services/cloud_scheduler/client.py

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -364,10 +364,7 @@ def __init__(
364364
client_cert_source_for_mtls=client_cert_source_func,
365365
quota_project_id=client_options.quota_project_id,
366366
client_info=client_info,
367-
always_use_jwt_access=(
368-
Transport == type(self).get_transport_class("grpc")
369-
or Transport == type(self).get_transport_class("grpc_asyncio")
370-
),
367+
always_use_jwt_access=True,
371368
)
372369

373370
def list_jobs(
@@ -1015,6 +1012,19 @@ def run_job(
10151012
# Done; return the response.
10161013
return response
10171014

1015+
def __enter__(self):
1016+
return self
1017+
1018+
def __exit__(self, type, value, traceback):
1019+
"""Releases underlying transport's resources.
1020+
1021+
.. warning::
1022+
ONLY use as a context manager if the transport is NOT shared
1023+
with other clients! Exiting the with block will CLOSE the transport
1024+
and may cause errors in other clients!
1025+
"""
1026+
self.transport.close()
1027+
10181028

10191029
try:
10201030
DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(

‎google/cloud/scheduler_v1beta1/services/cloud_scheduler/transports/base.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -241,6 +241,15 @@ def _prep_wrapped_messages(self, client_info):
241241
),
242242
}
243243

244+
def close(self):
245+
"""Closes resources associated with the transport.
246+
247+
.. warning::
248+
Only call this method if the transport is NOT shared
249+
with other clients - this may cause errors in other clients!
250+
"""
251+
raise NotImplementedError()
252+
244253
@property
245254
def list_jobs(
246255
self,

‎google/cloud/scheduler_v1beta1/services/cloud_scheduler/transports/grpc.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -459,5 +459,8 @@ def run_job(self) -> Callable[[cloudscheduler.RunJobRequest], job.Job]:
459459
)
460460
return self._stubs["run_job"]
461461

462+
def close(self):
463+
self.grpc_channel.close()
464+
462465

463466
__all__ = ("CloudSchedulerGrpcTransport",)

‎google/cloud/scheduler_v1beta1/services/cloud_scheduler/transports/grpc_asyncio.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -472,5 +472,8 @@ def run_job(self) -> Callable[[cloudscheduler.RunJobRequest], Awaitable[job.Job]
472472
)
473473
return self._stubs["run_job"]
474474

475+
def close(self):
476+
return self.grpc_channel.close()
477+
475478

476479
__all__ = ("CloudSchedulerGrpcAsyncIOTransport",)

‎tests/unit/gapic/scheduler_v1/test_cloud_scheduler.py

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
from google.api_core import gapic_v1
3030
from google.api_core import grpc_helpers
3131
from google.api_core import grpc_helpers_async
32+
from google.api_core import path_template
3233
from google.auth import credentials as ga_credentials
3334
from google.auth.exceptions import MutualTLSChannelError
3435
from google.cloud.scheduler_v1.services.cloud_scheduler import CloudSchedulerAsyncClient
@@ -2444,6 +2445,9 @@ def test_cloud_scheduler_base_transport():
24442445
with pytest.raises(NotImplementedError):
24452446
getattr(transport, method)(request=object())
24462447

2448+
with pytest.raises(NotImplementedError):
2449+
transport.close()
2450+
24472451

24482452
@requires_google_auth_gte_1_25_0
24492453
def test_cloud_scheduler_base_transport_with_credentials_file():
@@ -2946,3 +2950,49 @@ def test_client_withDEFAULT_CLIENT_INFO():
29462950
credentials=ga_credentials.AnonymousCredentials(), client_info=client_info,
29472951
)
29482952
prep.assert_called_once_with(client_info)
2953+
2954+
2955+
@pytest.mark.asyncio
2956+
async def test_transport_close_async():
2957+
client = CloudSchedulerAsyncClient(
2958+
credentials=ga_credentials.AnonymousCredentials(), transport="grpc_asyncio",
2959+
)
2960+
with mock.patch.object(
2961+
type(getattr(client.transport, "grpc_channel")), "close"
2962+
) as close:
2963+
async with client:
2964+
close.assert_not_called()
2965+
close.assert_called_once()
2966+
2967+
2968+
def test_transport_close():
2969+
transports = {
2970+
"grpc": "_grpc_channel",
2971+
}
2972+
2973+
for transport, close_name in transports.items():
2974+
client = CloudSchedulerClient(
2975+
credentials=ga_credentials.AnonymousCredentials(), transport=transport
2976+
)
2977+
with mock.patch.object(
2978+
type(getattr(client.transport, close_name)), "close"
2979+
) as close:
2980+
with client:
2981+
close.assert_not_called()
2982+
close.assert_called_once()
2983+
2984+
2985+
def test_client_ctx():
2986+
transports = [
2987+
"grpc",
2988+
]
2989+
for transport in transports:
2990+
client = CloudSchedulerClient(
2991+
credentials=ga_credentials.AnonymousCredentials(), transport=transport
2992+
)
2993+
# Test client calls underlying transport.
2994+
with mock.patch.object(type(client.transport), "close") as close:
2995+
close.assert_not_called()
2996+
with client:
2997+
pass
2998+
close.assert_called()

‎tests/unit/gapic/scheduler_v1beta1/test_cloud_scheduler.py

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
from google.api_core import gapic_v1
3030
from google.api_core import grpc_helpers
3131
from google.api_core import grpc_helpers_async
32+
from google.api_core import path_template
3233
from google.auth import credentials as ga_credentials
3334
from google.auth.exceptions import MutualTLSChannelError
3435
from google.cloud.scheduler_v1beta1.services.cloud_scheduler import (
@@ -2446,6 +2447,9 @@ def test_cloud_scheduler_base_transport():
24462447
with pytest.raises(NotImplementedError):
24472448
getattr(transport, method)(request=object())
24482449

2450+
with pytest.raises(NotImplementedError):
2451+
transport.close()
2452+
24492453

24502454
@requires_google_auth_gte_1_25_0
24512455
def test_cloud_scheduler_base_transport_with_credentials_file():
@@ -2948,3 +2952,49 @@ def test_client_withDEFAULT_CLIENT_INFO():
29482952
credentials=ga_credentials.AnonymousCredentials(), client_info=client_info,
29492953
)
29502954
prep.assert_called_once_with(client_info)
2955+
2956+
2957+
@pytest.mark.asyncio
2958+
async def test_transport_close_async():
2959+
client = CloudSchedulerAsyncClient(
2960+
credentials=ga_credentials.AnonymousCredentials(), transport="grpc_asyncio",
2961+
)
2962+
with mock.patch.object(
2963+
type(getattr(client.transport, "grpc_channel")), "close"
2964+
) as close:
2965+
async with client:
2966+
close.assert_not_called()
2967+
close.assert_called_once()
2968+
2969+
2970+
def test_transport_close():
2971+
transports = {
2972+
"grpc": "_grpc_channel",
2973+
}
2974+
2975+
for transport, close_name in transports.items():
2976+
client = CloudSchedulerClient(
2977+
credentials=ga_credentials.AnonymousCredentials(), transport=transport
2978+
)
2979+
with mock.patch.object(
2980+
type(getattr(client.transport, close_name)), "close"
2981+
) as close:
2982+
with client:
2983+
close.assert_not_called()
2984+
close.assert_called_once()
2985+
2986+
2987+
def test_client_ctx():
2988+
transports = [
2989+
"grpc",
2990+
]
2991+
for transport in transports:
2992+
client = CloudSchedulerClient(
2993+
credentials=ga_credentials.AnonymousCredentials(), transport=transport
2994+
)
2995+
# Test client calls underlying transport.
2996+
with mock.patch.object(type(client.transport), "close") as close:
2997+
close.assert_not_called()
2998+
with client:
2999+
pass
3000+
close.assert_called()

0 commit comments

Comments
 (0)
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy