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

Commit 2d8fbd5

Browse files
gcf-owl-bot[bot]gcf-owl-bot[bot]
authored
feat: add context manager support in client (#41)
- [ ] 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 9e9d4f0 commit 2d8fbd5

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

43 files changed

+573
-25
lines changed

‎google/cloud/dataflow_v1beta3/services/flex_templates_service/async_client.py

+6
Original file line numberDiff line numberDiff line change
@@ -205,6 +205,12 @@ async def launch_flex_template(
205205
# Done; return the response.
206206
return response
207207

208+
async def __aenter__(self):
209+
return self
210+
211+
async def __aexit__(self, exc_type, exc, tb):
212+
await self.transport.close()
213+
208214

209215
try:
210216
DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(

‎google/cloud/dataflow_v1beta3/services/flex_templates_service/client.py

+14-4
Original file line numberDiff line numberDiff line change
@@ -332,10 +332,7 @@ def __init__(
332332
client_cert_source_for_mtls=client_cert_source_func,
333333
quota_project_id=client_options.quota_project_id,
334334
client_info=client_info,
335-
always_use_jwt_access=(
336-
Transport == type(self).get_transport_class("grpc")
337-
or Transport == type(self).get_transport_class("grpc_asyncio")
338-
),
335+
always_use_jwt_access=True,
339336
)
340337

341338
def launch_flex_template(
@@ -382,6 +379,19 @@ def launch_flex_template(
382379
# Done; return the response.
383380
return response
384381

382+
def __enter__(self):
383+
return self
384+
385+
def __exit__(self, type, value, traceback):
386+
"""Releases underlying transport's resources.
387+
388+
.. warning::
389+
ONLY use as a context manager if the transport is NOT shared
390+
with other clients! Exiting the with block will CLOSE the transport
391+
and may cause errors in other clients!
392+
"""
393+
self.transport.close()
394+
385395

386396
try:
387397
DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(

‎google/cloud/dataflow_v1beta3/services/flex_templates_service/transports/base.py

+9
Original file line numberDiff line numberDiff line change
@@ -166,6 +166,15 @@ def _prep_wrapped_messages(self, client_info):
166166
),
167167
}
168168

169+
def close(self):
170+
"""Closes resources associated with the transport.
171+
172+
.. warning::
173+
Only call this method if the transport is NOT shared
174+
with other clients - this may cause errors in other clients!
175+
"""
176+
raise NotImplementedError()
177+
169178
@property
170179
def launch_flex_template(
171180
self,

‎google/cloud/dataflow_v1beta3/services/flex_templates_service/transports/grpc.py

+3
Original file line numberDiff line numberDiff line change
@@ -254,5 +254,8 @@ def launch_flex_template(
254254
)
255255
return self._stubs["launch_flex_template"]
256256

257+
def close(self):
258+
self.grpc_channel.close()
259+
257260

258261
__all__ = ("FlexTemplatesServiceGrpcTransport",)

‎google/cloud/dataflow_v1beta3/services/flex_templates_service/transports/grpc_asyncio.py

+3
Original file line numberDiff line numberDiff line change
@@ -258,5 +258,8 @@ def launch_flex_template(
258258
)
259259
return self._stubs["launch_flex_template"]
260260

261+
def close(self):
262+
return self.grpc_channel.close()
263+
261264

262265
__all__ = ("FlexTemplatesServiceGrpcAsyncIOTransport",)

‎google/cloud/dataflow_v1beta3/services/jobs_v1_beta3/async_client.py

+6
Original file line numberDiff line numberDiff line change
@@ -517,6 +517,12 @@ async def snapshot_job(
517517
# Done; return the response.
518518
return response
519519

520+
async def __aenter__(self):
521+
return self
522+
523+
async def __aexit__(self, exc_type, exc, tb):
524+
await self.transport.close()
525+
520526

521527
try:
522528
DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(

‎google/cloud/dataflow_v1beta3/services/jobs_v1_beta3/client.py

+14-4
Original file line numberDiff line numberDiff line change
@@ -333,10 +333,7 @@ def __init__(
333333
client_cert_source_for_mtls=client_cert_source_func,
334334
quota_project_id=client_options.quota_project_id,
335335
client_info=client_info,
336-
always_use_jwt_access=(
337-
Transport == type(self).get_transport_class("grpc")
338-
or Transport == type(self).get_transport_class("grpc_asyncio")
339-
),
336+
always_use_jwt_access=True,
340337
)
341338

342339
def create_job(
@@ -701,6 +698,19 @@ def snapshot_job(
701698
# Done; return the response.
702699
return response
703700

701+
def __enter__(self):
702+
return self
703+
704+
def __exit__(self, type, value, traceback):
705+
"""Releases underlying transport's resources.
706+
707+
.. warning::
708+
ONLY use as a context manager if the transport is NOT shared
709+
with other clients! Exiting the with block will CLOSE the transport
710+
and may cause errors in other clients!
711+
"""
712+
self.transport.close()
713+
704714

705715
try:
706716
DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(

‎google/cloud/dataflow_v1beta3/services/jobs_v1_beta3/transports/base.py

+9
Original file line numberDiff line numberDiff line change
@@ -185,6 +185,15 @@ def _prep_wrapped_messages(self, client_info):
185185
),
186186
}
187187

188+
def close(self):
189+
"""Closes resources associated with the transport.
190+
191+
.. warning::
192+
Only call this method if the transport is NOT shared
193+
with other clients - this may cause errors in other clients!
194+
"""
195+
raise NotImplementedError()
196+
188197
@property
189198
def create_job(
190199
self,

‎google/cloud/dataflow_v1beta3/services/jobs_v1_beta3/transports/grpc.py

+3
Original file line numberDiff line numberDiff line change
@@ -428,5 +428,8 @@ def snapshot_job(self) -> Callable[[jobs.SnapshotJobRequest], snapshots.Snapshot
428428
)
429429
return self._stubs["snapshot_job"]
430430

431+
def close(self):
432+
self.grpc_channel.close()
433+
431434

432435
__all__ = ("JobsV1Beta3GrpcTransport",)

‎google/cloud/dataflow_v1beta3/services/jobs_v1_beta3/transports/grpc_asyncio.py

+3
Original file line numberDiff line numberDiff line change
@@ -437,5 +437,8 @@ def snapshot_job(
437437
)
438438
return self._stubs["snapshot_job"]
439439

440+
def close(self):
441+
return self.grpc_channel.close()
442+
440443

441444
__all__ = ("JobsV1Beta3GrpcAsyncIOTransport",)

‎google/cloud/dataflow_v1beta3/services/messages_v1_beta3/async_client.py

+6
Original file line numberDiff line numberDiff line change
@@ -224,6 +224,12 @@ async def list_job_messages(
224224
# Done; return the response.
225225
return response
226226

227+
async def __aenter__(self):
228+
return self
229+
230+
async def __aexit__(self, exc_type, exc, tb):
231+
await self.transport.close()
232+
227233

228234
try:
229235
DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(

‎google/cloud/dataflow_v1beta3/services/messages_v1_beta3/client.py

+14-4
Original file line numberDiff line numberDiff line change
@@ -330,10 +330,7 @@ def __init__(
330330
client_cert_source_for_mtls=client_cert_source_func,
331331
quota_project_id=client_options.quota_project_id,
332332
client_info=client_info,
333-
always_use_jwt_access=(
334-
Transport == type(self).get_transport_class("grpc")
335-
or Transport == type(self).get_transport_class("grpc_asyncio")
336-
),
333+
always_use_jwt_access=True,
337334
)
338335

339336
def list_job_messages(
@@ -400,6 +397,19 @@ def list_job_messages(
400397
# Done; return the response.
401398
return response
402399

400+
def __enter__(self):
401+
return self
402+
403+
def __exit__(self, type, value, traceback):
404+
"""Releases underlying transport's resources.
405+
406+
.. warning::
407+
ONLY use as a context manager if the transport is NOT shared
408+
with other clients! Exiting the with block will CLOSE the transport
409+
and may cause errors in other clients!
410+
"""
411+
self.transport.close()
412+
403413

404414
try:
405415
DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(

‎google/cloud/dataflow_v1beta3/services/messages_v1_beta3/transports/base.py

+9
Original file line numberDiff line numberDiff line change
@@ -164,6 +164,15 @@ def _prep_wrapped_messages(self, client_info):
164164
),
165165
}
166166

167+
def close(self):
168+
"""Closes resources associated with the transport.
169+
170+
.. warning::
171+
Only call this method if the transport is NOT shared
172+
with other clients - this may cause errors in other clients!
173+
"""
174+
raise NotImplementedError()
175+
167176
@property
168177
def list_job_messages(
169178
self,

‎google/cloud/dataflow_v1beta3/services/messages_v1_beta3/transports/grpc.py

+3
Original file line numberDiff line numberDiff line change
@@ -260,5 +260,8 @@ def list_job_messages(
260260
)
261261
return self._stubs["list_job_messages"]
262262

263+
def close(self):
264+
self.grpc_channel.close()
265+
263266

264267
__all__ = ("MessagesV1Beta3GrpcTransport",)

‎google/cloud/dataflow_v1beta3/services/messages_v1_beta3/transports/grpc_asyncio.py

+3
Original file line numberDiff line numberDiff line change
@@ -265,5 +265,8 @@ def list_job_messages(
265265
)
266266
return self._stubs["list_job_messages"]
267267

268+
def close(self):
269+
return self.grpc_channel.close()
270+
268271

269272
__all__ = ("MessagesV1Beta3GrpcAsyncIOTransport",)

‎google/cloud/dataflow_v1beta3/services/metrics_v1_beta3/async_client.py

+6
Original file line numberDiff line numberDiff line change
@@ -333,6 +333,12 @@ async def get_stage_execution_details(
333333
# Done; return the response.
334334
return response
335335

336+
async def __aenter__(self):
337+
return self
338+
339+
async def __aexit__(self, exc_type, exc, tb):
340+
await self.transport.close()
341+
336342

337343
try:
338344
DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(

‎google/cloud/dataflow_v1beta3/services/metrics_v1_beta3/client.py

+14-4
Original file line numberDiff line numberDiff line change
@@ -331,10 +331,7 @@ def __init__(
331331
client_cert_source_for_mtls=client_cert_source_func,
332332
quota_project_id=client_options.quota_project_id,
333333
client_info=client_info,
334-
always_use_jwt_access=(
335-
Transport == type(self).get_transport_class("grpc")
336-
or Transport == type(self).get_transport_class("grpc_asyncio")
337-
),
334+
always_use_jwt_access=True,
338335
)
339336

340337
def get_job_metrics(
@@ -515,6 +512,19 @@ def get_stage_execution_details(
515512
# Done; return the response.
516513
return response
517514

515+
def __enter__(self):
516+
return self
517+
518+
def __exit__(self, type, value, traceback):
519+
"""Releases underlying transport's resources.
520+
521+
.. warning::
522+
ONLY use as a context manager if the transport is NOT shared
523+
with other clients! Exiting the with block will CLOSE the transport
524+
and may cause errors in other clients!
525+
"""
526+
self.transport.close()
527+
518528

519529
try:
520530
DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(

‎google/cloud/dataflow_v1beta3/services/metrics_v1_beta3/transports/base.py

+9
Original file line numberDiff line numberDiff line change
@@ -174,6 +174,15 @@ def _prep_wrapped_messages(self, client_info):
174174
),
175175
}
176176

177+
def close(self):
178+
"""Closes resources associated with the transport.
179+
180+
.. warning::
181+
Only call this method if the transport is NOT shared
182+
with other clients - this may cause errors in other clients!
183+
"""
184+
raise NotImplementedError()
185+
177186
@property
178187
def get_job_metrics(
179188
self,

‎google/cloud/dataflow_v1beta3/services/metrics_v1_beta3/transports/grpc.py

+3
Original file line numberDiff line numberDiff line change
@@ -321,5 +321,8 @@ def get_stage_execution_details(
321321
)
322322
return self._stubs["get_stage_execution_details"]
323323

324+
def close(self):
325+
self.grpc_channel.close()
326+
324327

325328
__all__ = ("MetricsV1Beta3GrpcTransport",)

‎google/cloud/dataflow_v1beta3/services/metrics_v1_beta3/transports/grpc_asyncio.py

+3
Original file line numberDiff line numberDiff line change
@@ -327,5 +327,8 @@ def get_stage_execution_details(
327327
)
328328
return self._stubs["get_stage_execution_details"]
329329

330+
def close(self):
331+
return self.grpc_channel.close()
332+
330333

331334
__all__ = ("MetricsV1Beta3GrpcAsyncIOTransport",)

‎google/cloud/dataflow_v1beta3/services/snapshots_v1_beta3/async_client.py

+6
Original file line numberDiff line numberDiff line change
@@ -283,6 +283,12 @@ async def list_snapshots(
283283
# Done; return the response.
284284
return response
285285

286+
async def __aenter__(self):
287+
return self
288+
289+
async def __aexit__(self, exc_type, exc, tb):
290+
await self.transport.close()
291+
286292

287293
try:
288294
DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(

‎google/cloud/dataflow_v1beta3/services/snapshots_v1_beta3/client.py

+14-4
Original file line numberDiff line numberDiff line change
@@ -331,10 +331,7 @@ def __init__(
331331
client_cert_source_for_mtls=client_cert_source_func,
332332
quota_project_id=client_options.quota_project_id,
333333
client_info=client_info,
334-
always_use_jwt_access=(
335-
Transport == type(self).get_transport_class("grpc")
336-
or Transport == type(self).get_transport_class("grpc_asyncio")
337-
),
334+
always_use_jwt_access=True,
338335
)
339336

340337
def get_snapshot(
@@ -461,6 +458,19 @@ def list_snapshots(
461458
# Done; return the response.
462459
return response
463460

461+
def __enter__(self):
462+
return self
463+
464+
def __exit__(self, type, value, traceback):
465+
"""Releases underlying transport's resources.
466+
467+
.. warning::
468+
ONLY use as a context manager if the transport is NOT shared
469+
with other clients! Exiting the with block will CLOSE the transport
470+
and may cause errors in other clients!
471+
"""
472+
self.transport.close()
473+
464474

465475
try:
466476
DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(

‎google/cloud/dataflow_v1beta3/services/snapshots_v1_beta3/transports/base.py

+9
Original file line numberDiff line numberDiff line change
@@ -170,6 +170,15 @@ def _prep_wrapped_messages(self, client_info):
170170
),
171171
}
172172

173+
def close(self):
174+
"""Closes resources associated with the transport.
175+
176+
.. warning::
177+
Only call this method if the transport is NOT shared
178+
with other clients - this may cause errors in other clients!
179+
"""
180+
raise NotImplementedError()
181+
173182
@property
174183
def get_snapshot(
175184
self,

‎google/cloud/dataflow_v1beta3/services/snapshots_v1_beta3/transports/grpc.py

+3
Original file line numberDiff line numberDiff line change
@@ -304,5 +304,8 @@ def list_snapshots(
304304
)
305305
return self._stubs["list_snapshots"]
306306

307+
def close(self):
308+
self.grpc_channel.close()
309+
307310

308311
__all__ = ("SnapshotsV1Beta3GrpcTransport",)

‎google/cloud/dataflow_v1beta3/services/snapshots_v1_beta3/transports/grpc_asyncio.py

+3
Original file line numberDiff line numberDiff line change
@@ -311,5 +311,8 @@ def list_snapshots(
311311
)
312312
return self._stubs["list_snapshots"]
313313

314+
def close(self):
315+
return self.grpc_channel.close()
316+
314317

315318
__all__ = ("SnapshotsV1Beta3GrpcAsyncIOTransport",)

0 commit comments

Comments
 (0)