Skip to content
This repository was archived by the owner on Feb 1, 2024. It is now read-only.

Commit ceec8f1

Browse files
gcf-owl-bot[bot]gcf-owl-bot[bot]
authored
feat: add context manager support in client (#173)
- [ ] 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 61a5857 commit ceec8f1

File tree

24 files changed

+266
-13
lines changed

24 files changed

+266
-13
lines changed

‎google/cloud/tasks_v2/services/cloud_tasks/async_client.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1755,6 +1755,12 @@ async def run_task(
17551755
# Done; return the response.
17561756
return response
17571757

1758+
async def __aenter__(self):
1759+
return self
1760+
1761+
async def __aexit__(self, exc_type, exc, tb):
1762+
await self.transport.close()
1763+
17581764

17591765
try:
17601766
DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(

‎google/cloud/tasks_v2/services/cloud_tasks/client.py

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -370,10 +370,7 @@ def __init__(
370370
client_cert_source_for_mtls=client_cert_source_func,
371371
quota_project_id=client_options.quota_project_id,
372372
client_info=client_info,
373-
always_use_jwt_access=(
374-
Transport == type(self).get_transport_class("grpc")
375-
or Transport == type(self).get_transport_class("grpc_asyncio")
376-
),
373+
always_use_jwt_access=True,
377374
)
378375

379376
def list_queues(
@@ -1879,6 +1876,19 @@ def run_task(
18791876
# Done; return the response.
18801877
return response
18811878

1879+
def __enter__(self):
1880+
return self
1881+
1882+
def __exit__(self, type, value, traceback):
1883+
"""Releases underlying transport's resources.
1884+
1885+
.. warning::
1886+
ONLY use as a context manager if the transport is NOT shared
1887+
with other clients! Exiting the with block will CLOSE the transport
1888+
and may cause errors in other clients!
1889+
"""
1890+
self.transport.close()
1891+
18821892

18831893
try:
18841894
DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(

‎google/cloud/tasks_v2/services/cloud_tasks/transports/base.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -305,6 +305,15 @@ def _prep_wrapped_messages(self, client_info):
305305
),
306306
}
307307

308+
def close(self):
309+
"""Closes resources associated with the transport.
310+
311+
.. warning::
312+
Only call this method if the transport is NOT shared
313+
with other clients - this may cause errors in other clients!
314+
"""
315+
raise NotImplementedError()
316+
308317
@property
309318
def list_queues(
310319
self,

‎google/cloud/tasks_v2/services/cloud_tasks/transports/grpc.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -768,5 +768,8 @@ def run_task(self) -> Callable[[cloudtasks.RunTaskRequest], task.Task]:
768768
)
769769
return self._stubs["run_task"]
770770

771+
def close(self):
772+
self.grpc_channel.close()
773+
771774

772775
__all__ = ("CloudTasksGrpcTransport",)

‎google/cloud/tasks_v2/services/cloud_tasks/transports/grpc_asyncio.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -787,5 +787,8 @@ def run_task(self) -> Callable[[cloudtasks.RunTaskRequest], Awaitable[task.Task]
787787
)
788788
return self._stubs["run_task"]
789789

790+
def close(self):
791+
return self.grpc_channel.close()
792+
790793

791794
__all__ = ("CloudTasksGrpcAsyncIOTransport",)

‎google/cloud/tasks_v2/types/task.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828

2929
class Task(proto.Message):
3030
r"""A unit of scheduled work.
31+
3132
Attributes:
3233
name (str):
3334
Optionally caller-specified in
@@ -177,6 +178,7 @@ class View(proto.Enum):
177178

178179
class Attempt(proto.Message):
179180
r"""The status of a task attempt.
181+
180182
Attributes:
181183
schedule_time (google.protobuf.timestamp_pb2.Timestamp):
182184
Output only. The time that this attempt was scheduled.

‎google/cloud/tasks_v2beta2/services/cloud_tasks/async_client.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2189,6 +2189,12 @@ async def run_task(
21892189
# Done; return the response.
21902190
return response
21912191

2192+
async def __aenter__(self):
2193+
return self
2194+
2195+
async def __aexit__(self, exc_type, exc, tb):
2196+
await self.transport.close()
2197+
21922198

21932199
try:
21942200
DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(

‎google/cloud/tasks_v2beta2/services/cloud_tasks/client.py

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -370,10 +370,7 @@ def __init__(
370370
client_cert_source_for_mtls=client_cert_source_func,
371371
quota_project_id=client_options.quota_project_id,
372372
client_info=client_info,
373-
always_use_jwt_access=(
374-
Transport == type(self).get_transport_class("grpc")
375-
or Transport == type(self).get_transport_class("grpc_asyncio")
376-
),
373+
always_use_jwt_access=True,
377374
)
378375

379376
def list_queues(
@@ -2313,6 +2310,19 @@ def run_task(
23132310
# Done; return the response.
23142311
return response
23152312

2313+
def __enter__(self):
2314+
return self
2315+
2316+
def __exit__(self, type, value, traceback):
2317+
"""Releases underlying transport's resources.
2318+
2319+
.. warning::
2320+
ONLY use as a context manager if the transport is NOT shared
2321+
with other clients! Exiting the with block will CLOSE the transport
2322+
and may cause errors in other clients!
2323+
"""
2324+
self.transport.close()
2325+
23162326

23172327
try:
23182328
DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(

‎google/cloud/tasks_v2beta2/services/cloud_tasks/transports/base.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -317,6 +317,15 @@ def _prep_wrapped_messages(self, client_info):
317317
),
318318
}
319319

320+
def close(self):
321+
"""Closes resources associated with the transport.
322+
323+
.. warning::
324+
Only call this method if the transport is NOT shared
325+
with other clients - this may cause errors in other clients!
326+
"""
327+
raise NotImplementedError()
328+
320329
@property
321330
def list_queues(
322331
self,

‎google/cloud/tasks_v2beta2/services/cloud_tasks/transports/grpc.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -927,5 +927,8 @@ def run_task(self) -> Callable[[cloudtasks.RunTaskRequest], task.Task]:
927927
)
928928
return self._stubs["run_task"]
929929

930+
def close(self):
931+
self.grpc_channel.close()
932+
930933

931934
__all__ = ("CloudTasksGrpcTransport",)

‎google/cloud/tasks_v2beta2/services/cloud_tasks/transports/grpc_asyncio.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -952,5 +952,8 @@ def run_task(self) -> Callable[[cloudtasks.RunTaskRequest], Awaitable[task.Task]
952952
)
953953
return self._stubs["run_task"]
954954

955+
def close(self):
956+
return self.grpc_channel.close()
957+
955958

956959
__all__ = ("CloudTasksGrpcAsyncIOTransport",)

‎google/cloud/tasks_v2beta2/types/queue.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -408,6 +408,7 @@ class RetryConfig(proto.Message):
408408

409409
class QueueStats(proto.Message):
410410
r"""Statistics for a queue.
411+
411412
Attributes:
412413
tasks_count (int):
413414
Output only. An estimation of the number of

‎google/cloud/tasks_v2beta2/types/target.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,8 @@ class HttpMethod(proto.Enum):
4040

4141

4242
class PullTarget(proto.Message):
43-
r"""Pull target. """
43+
r"""Pull target.
44+
"""
4445

4546

4647
class PullMessage(proto.Message):

‎google/cloud/tasks_v2beta2/types/task.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828

2929
class Task(proto.Message):
3030
r"""A unit of scheduled work.
31+
3132
Attributes:
3233
name (str):
3334
Optionally caller-specified in
@@ -129,6 +130,7 @@ class View(proto.Enum):
129130

130131
class TaskStatus(proto.Message):
131132
r"""Status of the task.
133+
132134
Attributes:
133135
attempt_dispatch_count (int):
134136
Output only. The number of attempts
@@ -169,6 +171,7 @@ class TaskStatus(proto.Message):
169171

170172
class AttemptStatus(proto.Message):
171173
r"""The status of a task attempt.
174+
172175
Attributes:
173176
schedule_time (google.protobuf.timestamp_pb2.Timestamp):
174177
Output only. The time that this attempt was scheduled.

‎google/cloud/tasks_v2beta3/services/cloud_tasks/async_client.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1757,6 +1757,12 @@ async def run_task(
17571757
# Done; return the response.
17581758
return response
17591759

1760+
async def __aenter__(self):
1761+
return self
1762+
1763+
async def __aexit__(self, exc_type, exc, tb):
1764+
await self.transport.close()
1765+
17601766

17611767
try:
17621768
DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(

‎google/cloud/tasks_v2beta3/services/cloud_tasks/client.py

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -370,10 +370,7 @@ def __init__(
370370
client_cert_source_for_mtls=client_cert_source_func,
371371
quota_project_id=client_options.quota_project_id,
372372
client_info=client_info,
373-
always_use_jwt_access=(
374-
Transport == type(self).get_transport_class("grpc")
375-
or Transport == type(self).get_transport_class("grpc_asyncio")
376-
),
373+
always_use_jwt_access=True,
377374
)
378375

379376
def list_queues(
@@ -1881,6 +1878,19 @@ def run_task(
18811878
# Done; return the response.
18821879
return response
18831880

1881+
def __enter__(self):
1882+
return self
1883+
1884+
def __exit__(self, type, value, traceback):
1885+
"""Releases underlying transport's resources.
1886+
1887+
.. warning::
1888+
ONLY use as a context manager if the transport is NOT shared
1889+
with other clients! Exiting the with block will CLOSE the transport
1890+
and may cause errors in other clients!
1891+
"""
1892+
self.transport.close()
1893+
18841894

18851895
try:
18861896
DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(

‎google/cloud/tasks_v2beta3/services/cloud_tasks/transports/base.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -305,6 +305,15 @@ def _prep_wrapped_messages(self, client_info):
305305
),
306306
}
307307

308+
def close(self):
309+
"""Closes resources associated with the transport.
310+
311+
.. warning::
312+
Only call this method if the transport is NOT shared
313+
with other clients - this may cause errors in other clients!
314+
"""
315+
raise NotImplementedError()
316+
308317
@property
309318
def list_queues(
310319
self,

‎google/cloud/tasks_v2beta3/services/cloud_tasks/transports/grpc.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -769,5 +769,8 @@ def run_task(self) -> Callable[[cloudtasks.RunTaskRequest], task.Task]:
769769
)
770770
return self._stubs["run_task"]
771771

772+
def close(self):
773+
self.grpc_channel.close()
774+
772775

773776
__all__ = ("CloudTasksGrpcTransport",)

‎google/cloud/tasks_v2beta3/services/cloud_tasks/transports/grpc_asyncio.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -788,5 +788,8 @@ def run_task(self) -> Callable[[cloudtasks.RunTaskRequest], Awaitable[task.Task]
788788
)
789789
return self._stubs["run_task"]
790790

791+
def close(self):
792+
return self.grpc_channel.close()
793+
791794

792795
__all__ = ("CloudTasksGrpcAsyncIOTransport",)

‎google/cloud/tasks_v2beta3/types/queue.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -433,6 +433,7 @@ class StackdriverLoggingConfig(proto.Message):
433433

434434
class QueueStats(proto.Message):
435435
r"""Statistics for a queue.
436+
436437
Attributes:
437438
tasks_count (int):
438439
Output only. An estimation of the number of

‎google/cloud/tasks_v2beta3/types/task.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828

2929
class Task(proto.Message):
3030
r"""A unit of scheduled work.
31+
3132
Attributes:
3233
name (str):
3334
Optionally caller-specified in
@@ -192,6 +193,7 @@ class View(proto.Enum):
192193

193194
class Attempt(proto.Message):
194195
r"""The status of a task attempt.
196+
195197
Attributes:
196198
schedule_time (google.protobuf.timestamp_pb2.Timestamp):
197199
Output only. The time that this attempt was scheduled.

‎tests/unit/gapic/tasks_v2/test_cloud_tasks.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.tasks_v2.services.cloud_tasks import CloudTasksAsyncClient
@@ -4091,6 +4092,9 @@ def test_cloud_tasks_base_transport():
40914092
with pytest.raises(NotImplementedError):
40924093
getattr(transport, method)(request=object())
40934094

4095+
with pytest.raises(NotImplementedError):
4096+
transport.close()
4097+
40944098

40954099
@requires_google_auth_gte_1_25_0
40964100
def test_cloud_tasks_base_transport_with_credentials_file():
@@ -4582,3 +4586,49 @@ def test_client_withDEFAULT_CLIENT_INFO():
45824586
credentials=ga_credentials.AnonymousCredentials(), client_info=client_info,
45834587
)
45844588
prep.assert_called_once_with(client_info)
4589+
4590+
4591+
@pytest.mark.asyncio
4592+
async def test_transport_close_async():
4593+
client = CloudTasksAsyncClient(
4594+
credentials=ga_credentials.AnonymousCredentials(), transport="grpc_asyncio",
4595+
)
4596+
with mock.patch.object(
4597+
type(getattr(client.transport, "grpc_channel")), "close"
4598+
) as close:
4599+
async with client:
4600+
close.assert_not_called()
4601+
close.assert_called_once()
4602+
4603+
4604+
def test_transport_close():
4605+
transports = {
4606+
"grpc": "_grpc_channel",
4607+
}
4608+
4609+
for transport, close_name in transports.items():
4610+
client = CloudTasksClient(
4611+
credentials=ga_credentials.AnonymousCredentials(), transport=transport
4612+
)
4613+
with mock.patch.object(
4614+
type(getattr(client.transport, close_name)), "close"
4615+
) as close:
4616+
with client:
4617+
close.assert_not_called()
4618+
close.assert_called_once()
4619+
4620+
4621+
def test_client_ctx():
4622+
transports = [
4623+
"grpc",
4624+
]
4625+
for transport in transports:
4626+
client = CloudTasksClient(
4627+
credentials=ga_credentials.AnonymousCredentials(), transport=transport
4628+
)
4629+
# Test client calls underlying transport.
4630+
with mock.patch.object(type(client.transport), "close") as close:
4631+
close.assert_not_called()
4632+
with client:
4633+
pass
4634+
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