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

Commit e88a664

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

File tree

28 files changed

+347
-18
lines changed

28 files changed

+347
-18
lines changed

‎google/cloud/orchestration/airflow/service_v1/services/environments/async_client.py

+6
Original file line numberDiff line numberDiff line change
@@ -759,6 +759,12 @@ async def delete_environment(
759759
# Done; return the response.
760760
return response
761761

762+
async def __aenter__(self):
763+
return self
764+
765+
async def __aexit__(self, exc_type, exc, tb):
766+
await self.transport.close()
767+
762768

763769
try:
764770
DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(

‎google/cloud/orchestration/airflow/service_v1/services/environments/client.py

+14-4
Original file line numberDiff line numberDiff line change
@@ -348,10 +348,7 @@ def __init__(
348348
client_cert_source_for_mtls=client_cert_source_func,
349349
quota_project_id=client_options.quota_project_id,
350350
client_info=client_info,
351-
always_use_jwt_access=(
352-
Transport == type(self).get_transport_class("grpc")
353-
or Transport == type(self).get_transport_class("grpc_asyncio")
354-
),
351+
always_use_jwt_access=True,
355352
)
356353

357354
def create_environment(
@@ -950,6 +947,19 @@ def delete_environment(
950947
# Done; return the response.
951948
return response
952949

950+
def __enter__(self):
951+
return self
952+
953+
def __exit__(self, type, value, traceback):
954+
"""Releases underlying transport's resources.
955+
956+
.. warning::
957+
ONLY use as a context manager if the transport is NOT shared
958+
with other clients! Exiting the with block will CLOSE the transport
959+
and may cause errors in other clients!
960+
"""
961+
self.transport.close()
962+
953963

954964
try:
955965
DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(

‎google/cloud/orchestration/airflow/service_v1/services/environments/transports/base.py

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

176+
def close(self):
177+
"""Closes resources associated with the transport.
178+
179+
.. warning::
180+
Only call this method if the transport is NOT shared
181+
with other clients - this may cause errors in other clients!
182+
"""
183+
raise NotImplementedError()
184+
176185
@property
177186
def operations_client(self) -> operations_v1.OperationsClient:
178187
"""Return the client designed to process long-running operations."""

‎google/cloud/orchestration/airflow/service_v1/services/environments/transports/grpc.py

+3
Original file line numberDiff line numberDiff line change
@@ -374,5 +374,8 @@ def delete_environment(
374374
)
375375
return self._stubs["delete_environment"]
376376

377+
def close(self):
378+
self.grpc_channel.close()
379+
377380

378381
__all__ = ("EnvironmentsGrpcTransport",)

‎google/cloud/orchestration/airflow/service_v1/services/environments/transports/grpc_asyncio.py

+3
Original file line numberDiff line numberDiff line change
@@ -388,5 +388,8 @@ def delete_environment(
388388
)
389389
return self._stubs["delete_environment"]
390390

391+
def close(self):
392+
return self.grpc_channel.close()
393+
391394

392395
__all__ = ("EnvironmentsGrpcAsyncIOTransport",)

‎google/cloud/orchestration/airflow/service_v1/services/image_versions/async_client.py

+6
Original file line numberDiff line numberDiff line change
@@ -240,6 +240,12 @@ async def list_image_versions(
240240
# Done; return the response.
241241
return response
242242

243+
async def __aenter__(self):
244+
return self
245+
246+
async def __aexit__(self, exc_type, exc, tb):
247+
await self.transport.close()
248+
243249

244250
try:
245251
DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(

‎google/cloud/orchestration/airflow/service_v1/services/image_versions/client.py

+14-4
Original file line numberDiff line numberDiff line change
@@ -326,10 +326,7 @@ def __init__(
326326
client_cert_source_for_mtls=client_cert_source_func,
327327
quota_project_id=client_options.quota_project_id,
328328
client_info=client_info,
329-
always_use_jwt_access=(
330-
Transport == type(self).get_transport_class("grpc")
331-
or Transport == type(self).get_transport_class("grpc_asyncio")
332-
),
329+
always_use_jwt_access=True,
333330
)
334331

335332
def list_image_versions(
@@ -413,6 +410,19 @@ def list_image_versions(
413410
# Done; return the response.
414411
return response
415412

413+
def __enter__(self):
414+
return self
415+
416+
def __exit__(self, type, value, traceback):
417+
"""Releases underlying transport's resources.
418+
419+
.. warning::
420+
ONLY use as a context manager if the transport is NOT shared
421+
with other clients! Exiting the with block will CLOSE the transport
422+
and may cause errors in other clients!
423+
"""
424+
self.transport.close()
425+
416426

417427
try:
418428
DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(

‎google/cloud/orchestration/airflow/service_v1/services/image_versions/transports/base.py

+9
Original file line numberDiff line numberDiff line change
@@ -159,6 +159,15 @@ def _prep_wrapped_messages(self, client_info):
159159
),
160160
}
161161

162+
def close(self):
163+
"""Closes resources associated with the transport.
164+
165+
.. warning::
166+
Only call this method if the transport is NOT shared
167+
with other clients - this may cause errors in other clients!
168+
"""
169+
raise NotImplementedError()
170+
162171
@property
163172
def list_image_versions(
164173
self,

‎google/cloud/orchestration/airflow/service_v1/services/image_versions/transports/grpc.py

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

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

258261
__all__ = ("ImageVersionsGrpcTransport",)

‎google/cloud/orchestration/airflow/service_v1/services/image_versions/transports/grpc_asyncio.py

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

260+
def close(self):
261+
return self.grpc_channel.close()
262+
260263

261264
__all__ = ("ImageVersionsGrpcAsyncIOTransport",)

‎google/cloud/orchestration/airflow/service_v1/types/environments.py

+1
Original file line numberDiff line numberDiff line change
@@ -380,6 +380,7 @@ class WebServerNetworkAccessControl(proto.Message):
380380

381381
class AllowedIpRange(proto.Message):
382382
r"""Allowed IP range with user-provided description.
383+
383384
Attributes:
384385
value (str):
385386
IP address or range, defined using CIDR notation, of

‎google/cloud/orchestration/airflow/service_v1/types/image_versions.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -71,7 +71,7 @@ def raw_page(self):
7171

7272

7373
class ImageVersion(proto.Message):
74-
r"""ImageVersion informatio.
74+
r"""ImageVersion information
7575
7676
Attributes:
7777
image_version_id (str):

‎google/cloud/orchestration/airflow/service_v1beta1/services/environments/async_client.py

+6
Original file line numberDiff line numberDiff line change
@@ -904,6 +904,12 @@ async def check_upgrade(
904904
# Done; return the response.
905905
return response
906906

907+
async def __aenter__(self):
908+
return self
909+
910+
async def __aexit__(self, exc_type, exc, tb):
911+
await self.transport.close()
912+
907913

908914
try:
909915
DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(

‎google/cloud/orchestration/airflow/service_v1beta1/services/environments/client.py

+14-4
Original file line numberDiff line numberDiff line change
@@ -350,10 +350,7 @@ def __init__(
350350
client_cert_source_for_mtls=client_cert_source_func,
351351
quota_project_id=client_options.quota_project_id,
352352
client_info=client_info,
353-
always_use_jwt_access=(
354-
Transport == type(self).get_transport_class("grpc")
355-
or Transport == type(self).get_transport_class("grpc_asyncio")
356-
),
353+
always_use_jwt_access=True,
357354
)
358355

359356
def create_environment(
@@ -1097,6 +1094,19 @@ def check_upgrade(
10971094
# Done; return the response.
10981095
return response
10991096

1097+
def __enter__(self):
1098+
return self
1099+
1100+
def __exit__(self, type, value, traceback):
1101+
"""Releases underlying transport's resources.
1102+
1103+
.. warning::
1104+
ONLY use as a context manager if the transport is NOT shared
1105+
with other clients! Exiting the with block will CLOSE the transport
1106+
and may cause errors in other clients!
1107+
"""
1108+
self.transport.close()
1109+
11001110

11011111
try:
11021112
DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(

‎google/cloud/orchestration/airflow/service_v1beta1/services/environments/transports/base.py

+9
Original file line numberDiff line numberDiff line change
@@ -179,6 +179,15 @@ def _prep_wrapped_messages(self, client_info):
179179
),
180180
}
181181

182+
def close(self):
183+
"""Closes resources associated with the transport.
184+
185+
.. warning::
186+
Only call this method if the transport is NOT shared
187+
with other clients - this may cause errors in other clients!
188+
"""
189+
raise NotImplementedError()
190+
182191
@property
183192
def operations_client(self) -> operations_v1.OperationsClient:
184193
"""Return the client designed to process long-running operations."""

‎google/cloud/orchestration/airflow/service_v1beta1/services/environments/transports/grpc.py

+3
Original file line numberDiff line numberDiff line change
@@ -429,5 +429,8 @@ def check_upgrade(
429429
)
430430
return self._stubs["check_upgrade"]
431431

432+
def close(self):
433+
self.grpc_channel.close()
434+
432435

433436
__all__ = ("EnvironmentsGrpcTransport",)

‎google/cloud/orchestration/airflow/service_v1beta1/services/environments/transports/grpc_asyncio.py

+3
Original file line numberDiff line numberDiff line change
@@ -447,5 +447,8 @@ def check_upgrade(
447447
)
448448
return self._stubs["check_upgrade"]
449449

450+
def close(self):
451+
return self.grpc_channel.close()
452+
450453

451454
__all__ = ("EnvironmentsGrpcAsyncIOTransport",)

‎google/cloud/orchestration/airflow/service_v1beta1/services/image_versions/async_client.py

+6
Original file line numberDiff line numberDiff line change
@@ -242,6 +242,12 @@ async def list_image_versions(
242242
# Done; return the response.
243243
return response
244244

245+
async def __aenter__(self):
246+
return self
247+
248+
async def __aexit__(self, exc_type, exc, tb):
249+
await self.transport.close()
250+
245251

246252
try:
247253
DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(

‎google/cloud/orchestration/airflow/service_v1beta1/services/image_versions/client.py

+14-4
Original file line numberDiff line numberDiff line change
@@ -328,10 +328,7 @@ def __init__(
328328
client_cert_source_for_mtls=client_cert_source_func,
329329
quota_project_id=client_options.quota_project_id,
330330
client_info=client_info,
331-
always_use_jwt_access=(
332-
Transport == type(self).get_transport_class("grpc")
333-
or Transport == type(self).get_transport_class("grpc_asyncio")
334-
),
331+
always_use_jwt_access=True,
335332
)
336333

337334
def list_image_versions(
@@ -415,6 +412,19 @@ def list_image_versions(
415412
# Done; return the response.
416413
return response
417414

415+
def __enter__(self):
416+
return self
417+
418+
def __exit__(self, type, value, traceback):
419+
"""Releases underlying transport's resources.
420+
421+
.. warning::
422+
ONLY use as a context manager if the transport is NOT shared
423+
with other clients! Exiting the with block will CLOSE the transport
424+
and may cause errors in other clients!
425+
"""
426+
self.transport.close()
427+
418428

419429
try:
420430
DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(

‎google/cloud/orchestration/airflow/service_v1beta1/services/image_versions/transports/base.py

+9
Original file line numberDiff line numberDiff line change
@@ -159,6 +159,15 @@ def _prep_wrapped_messages(self, client_info):
159159
),
160160
}
161161

162+
def close(self):
163+
"""Closes resources associated with the transport.
164+
165+
.. warning::
166+
Only call this method if the transport is NOT shared
167+
with other clients - this may cause errors in other clients!
168+
"""
169+
raise NotImplementedError()
170+
162171
@property
163172
def list_image_versions(
164173
self,

‎google/cloud/orchestration/airflow/service_v1beta1/services/image_versions/transports/grpc.py

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

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

258261
__all__ = ("ImageVersionsGrpcTransport",)

‎google/cloud/orchestration/airflow/service_v1beta1/services/image_versions/transports/grpc_asyncio.py

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

260+
def close(self):
261+
return self.grpc_channel.close()
262+
260263

261264
__all__ = ("ImageVersionsGrpcAsyncIOTransport",)

‎google/cloud/orchestration/airflow/service_v1beta1/types/environments.py

+4
Original file line numberDiff line numberDiff line change
@@ -462,6 +462,7 @@ class WebServerNetworkAccessControl(proto.Message):
462462

463463
class AllowedIpRange(proto.Message):
464464
r"""Allowed IP range with user-provided description.
465+
465466
Attributes:
466467
value (str):
467468
IP address or range, defined using CIDR notation, of
@@ -964,6 +965,7 @@ class WorkloadsConfig(proto.Message):
964965

965966
class SchedulerResource(proto.Message):
966967
r"""Configuration for resources used by Airflow schedulers.
968+
967969
Attributes:
968970
cpu (float):
969971
Optional. CPU request and limit for a single
@@ -985,6 +987,7 @@ class SchedulerResource(proto.Message):
985987

986988
class WebServerResource(proto.Message):
987989
r"""Configuration for resources used by Airflow web server.
990+
988991
Attributes:
989992
cpu (float):
990993
Optional. CPU request and limit for Airflow
@@ -1003,6 +1006,7 @@ class WebServerResource(proto.Message):
10031006

10041007
class WorkerResource(proto.Message):
10051008
r"""Configuration for resources used by Airflow workers.
1009+
10061010
Attributes:
10071011
cpu (float):
10081012
Optional. CPU request and limit for a single

‎google/cloud/orchestration/airflow/service_v1beta1/types/image_versions.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -71,7 +71,7 @@ def raw_page(self):
7171

7272

7373
class ImageVersion(proto.Message):
74-
r"""Image Version informatio.
74+
r"""Image Version information
7575
7676
Attributes:
7777
image_version_id (str):

0 commit comments

Comments
 (0)