Skip to content

Commit a0282c9

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: googleapis/googleapis-gen@81decff Copy-Tag: eyJwIjoiLmdpdGh1Yi8uT3dsQm90LnlhbWwiLCJoIjoiODFkZWNmZmU5ZmM3MjM5NmE4MTUzZTc1NmQxZDY3YTZlZWNmZDYyMCJ9
1 parent 8fcc3df commit a0282c9

28 files changed

Lines changed: 347 additions & 18 deletions

File tree

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

Lines changed: 6 additions & 0 deletions
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(

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

Lines changed: 14 additions & 4 deletions
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(

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

Lines changed: 9 additions & 0 deletions
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."""

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

Lines changed: 3 additions & 0 deletions
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",)

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

Lines changed: 3 additions & 0 deletions
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",)

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

Lines changed: 6 additions & 0 deletions
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(

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

Lines changed: 14 additions & 4 deletions
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(

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

Lines changed: 9 additions & 0 deletions
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,

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

Lines changed: 3 additions & 0 deletions
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",)

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

Lines changed: 3 additions & 0 deletions
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",)

0 commit comments

Comments
 (0)