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

Commit c1fbdff

Browse files
feat: add context manager support in client (#11)
- [ ] 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 f43834b commit c1fbdff

File tree

7 files changed

+101
-4
lines changed

7 files changed

+101
-4
lines changed

google/cloud/filestore_v1/services/cloud_filestore_manager/async_client.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1198,6 +1198,12 @@ async def update_backup(
11981198
# Done; return the response.
11991199
return response
12001200

1201+
async def __aenter__(self):
1202+
return self
1203+
1204+
async def __aexit__(self, exc_type, exc, tb):
1205+
await self.transport.close()
1206+
12011207

12021208
try:
12031209
DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(

google/cloud/filestore_v1/services/cloud_filestore_manager/client.py

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -392,10 +392,7 @@ def __init__(
392392
client_cert_source_for_mtls=client_cert_source_func,
393393
quota_project_id=client_options.quota_project_id,
394394
client_info=client_info,
395-
always_use_jwt_access=(
396-
Transport == type(self).get_transport_class("grpc")
397-
or Transport == type(self).get_transport_class("grpc_asyncio")
398-
),
395+
always_use_jwt_access=True,
399396
)
400397

401398
def list_instances(
@@ -1365,6 +1362,19 @@ def update_backup(
13651362
# Done; return the response.
13661363
return response
13671364

1365+
def __enter__(self):
1366+
return self
1367+
1368+
def __exit__(self, type, value, traceback):
1369+
"""Releases underlying transport's resources.
1370+
1371+
.. warning::
1372+
ONLY use as a context manager if the transport is NOT shared
1373+
with other clients! Exiting the with block will CLOSE the transport
1374+
and may cause errors in other clients!
1375+
"""
1376+
self.transport.close()
1377+
13681378

13691379
try:
13701380
DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(

google/cloud/filestore_v1/services/cloud_filestore_manager/transports/base.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -233,6 +233,15 @@ def _prep_wrapped_messages(self, client_info):
233233
),
234234
}
235235

236+
def close(self):
237+
"""Closes resources associated with the transport.
238+
239+
.. warning::
240+
Only call this method if the transport is NOT shared
241+
with other clients - this may cause errors in other clients!
242+
"""
243+
raise NotImplementedError()
244+
236245
@property
237246
def operations_client(self) -> operations_v1.OperationsClient:
238247
"""Return the client designed to process long-running operations."""

google/cloud/filestore_v1/services/cloud_filestore_manager/transports/grpc.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -584,5 +584,8 @@ def update_backup(
584584
)
585585
return self._stubs["update_backup"]
586586

587+
def close(self):
588+
self.grpc_channel.close()
589+
587590

588591
__all__ = ("CloudFilestoreManagerGrpcTransport",)

google/cloud/filestore_v1/services/cloud_filestore_manager/transports/grpc_asyncio.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -598,5 +598,8 @@ def update_backup(
598598
)
599599
return self._stubs["update_backup"]
600600

601+
def close(self):
602+
return self.grpc_channel.close()
603+
601604

602605
__all__ = ("CloudFilestoreManagerGrpcAsyncIOTransport",)

google/cloud/filestore_v1/types/cloud_filestore_service.py

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@
4747

4848
class NetworkConfig(proto.Message):
4949
r"""Network configuration for the instance.
50+
5051
Attributes:
5152
network (str):
5253
The name of the Google Compute Engine `VPC
@@ -84,6 +85,7 @@ class AddressMode(proto.Enum):
8485

8586
class FileShareConfig(proto.Message):
8687
r"""File share configuration for the instance.
88+
8789
Attributes:
8890
name (str):
8991
The name of the file share (must be 16
@@ -111,6 +113,7 @@ class FileShareConfig(proto.Message):
111113

112114
class NfsExportOptions(proto.Message):
113115
r"""NFS export options specifications.
116+
114117
Attributes:
115118
ip_ranges (Sequence[str]):
116119
List of either an IPv4 addresses in the format
@@ -162,6 +165,7 @@ class SquashMode(proto.Enum):
162165

163166
class Instance(proto.Message):
164167
r"""A Cloud Filestore instance.
168+
165169
Attributes:
166170
name (str):
167171
Output only. The resource name of the instance, in the
@@ -237,6 +241,7 @@ class Tier(proto.Enum):
237241

238242
class CreateInstanceRequest(proto.Message):
239243
r"""CreateInstanceRequest creates an instance.
244+
240245
Attributes:
241246
parent (str):
242247
Required. The instance's project and location, in the format
@@ -259,6 +264,7 @@ class CreateInstanceRequest(proto.Message):
259264

260265
class GetInstanceRequest(proto.Message):
261266
r"""GetInstanceRequest gets the state of an instance.
267+
262268
Attributes:
263269
name (str):
264270
Required. The instance resource name, in the format
@@ -270,6 +276,7 @@ class GetInstanceRequest(proto.Message):
270276

271277
class UpdateInstanceRequest(proto.Message):
272278
r"""UpdateInstanceRequest updates the settings of an instance.
279+
273280
Attributes:
274281
update_mask (google.protobuf.field_mask_pb2.FieldMask):
275282
Mask of fields to update. At least one path must be supplied
@@ -313,6 +320,7 @@ class RestoreInstanceRequest(proto.Message):
313320

314321
class DeleteInstanceRequest(proto.Message):
315322
r"""DeleteInstanceRequest deletes an instance.
323+
316324
Attributes:
317325
name (str):
318326
Required. The instance resource name, in the format
@@ -324,6 +332,7 @@ class DeleteInstanceRequest(proto.Message):
324332

325333
class ListInstancesRequest(proto.Message):
326334
r"""ListInstancesRequest lists instances.
335+
327336
Attributes:
328337
parent (str):
329338
Required. The project and location for which to retrieve
@@ -353,6 +362,7 @@ class ListInstancesRequest(proto.Message):
353362

354363
class ListInstancesResponse(proto.Message):
355364
r"""ListInstancesResponse is the result of ListInstancesRequest.
365+
356366
Attributes:
357367
instances (Sequence[google.cloud.filestore_v1.types.Instance]):
358368
A list of instances in the project for the specified
@@ -382,6 +392,7 @@ def raw_page(self):
382392

383393
class Backup(proto.Message):
384394
r"""A Cloud Filestore backup.
395+
385396
Attributes:
386397
name (str):
387398
Output only. The resource name of the backup, in the format
@@ -455,6 +466,7 @@ class State(proto.Enum):
455466

456467
class CreateBackupRequest(proto.Message):
457468
r"""CreateBackupRequest creates a backup.
469+
458470
Attributes:
459471
parent (str):
460472
Required. The backup's project and location, in the format
@@ -481,6 +493,7 @@ class CreateBackupRequest(proto.Message):
481493

482494
class DeleteBackupRequest(proto.Message):
483495
r"""DeleteBackupRequest deletes a backup.
496+
484497
Attributes:
485498
name (str):
486499
Required. The backup resource name, in the format
@@ -511,6 +524,7 @@ class UpdateBackupRequest(proto.Message):
511524

512525
class GetBackupRequest(proto.Message):
513526
r"""GetBackupRequest gets the state of a backup.
527+
514528
Attributes:
515529
name (str):
516530
Required. The backup resource name, in the format
@@ -522,6 +536,7 @@ class GetBackupRequest(proto.Message):
522536

523537
class ListBackupsRequest(proto.Message):
524538
r"""ListBackupsRequest lists backups.
539+
525540
Attributes:
526541
parent (str):
527542
Required. The project and location for which to retrieve
@@ -551,6 +566,7 @@ class ListBackupsRequest(proto.Message):
551566

552567
class ListBackupsResponse(proto.Message):
553568
r"""ListBackupsResponse is the result of ListBackupsRequest.
569+
554570
Attributes:
555571
backups (Sequence[google.cloud.filestore_v1.types.Backup]):
556572
A list of backups in the project for the specified location.

tests/unit/gapic/filestore_v1/test_cloud_filestore_manager.py

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
from google.api_core import grpc_helpers_async
3333
from google.api_core import operation_async # type: ignore
3434
from google.api_core import operations_v1
35+
from google.api_core import path_template
3536
from google.auth import credentials as ga_credentials
3637
from google.auth.exceptions import MutualTLSChannelError
3738
from google.cloud.common.types import operation_metadata as operation_metadata_pb2 # type: ignore
@@ -3331,6 +3332,9 @@ def test_cloud_filestore_manager_base_transport():
33313332
with pytest.raises(NotImplementedError):
33323333
getattr(transport, method)(request=object())
33333334

3335+
with pytest.raises(NotImplementedError):
3336+
transport.close()
3337+
33343338
# Additionally, the LRO client (a property) should
33353339
# also raise NotImplementedError
33363340
with pytest.raises(NotImplementedError):
@@ -3870,3 +3874,49 @@ def test_client_withDEFAULT_CLIENT_INFO():
38703874
credentials=ga_credentials.AnonymousCredentials(), client_info=client_info,
38713875
)
38723876
prep.assert_called_once_with(client_info)
3877+
3878+
3879+
@pytest.mark.asyncio
3880+
async def test_transport_close_async():
3881+
client = CloudFilestoreManagerAsyncClient(
3882+
credentials=ga_credentials.AnonymousCredentials(), transport="grpc_asyncio",
3883+
)
3884+
with mock.patch.object(
3885+
type(getattr(client.transport, "grpc_channel")), "close"
3886+
) as close:
3887+
async with client:
3888+
close.assert_not_called()
3889+
close.assert_called_once()
3890+
3891+
3892+
def test_transport_close():
3893+
transports = {
3894+
"grpc": "_grpc_channel",
3895+
}
3896+
3897+
for transport, close_name in transports.items():
3898+
client = CloudFilestoreManagerClient(
3899+
credentials=ga_credentials.AnonymousCredentials(), transport=transport
3900+
)
3901+
with mock.patch.object(
3902+
type(getattr(client.transport, close_name)), "close"
3903+
) as close:
3904+
with client:
3905+
close.assert_not_called()
3906+
close.assert_called_once()
3907+
3908+
3909+
def test_client_ctx():
3910+
transports = [
3911+
"grpc",
3912+
]
3913+
for transport in transports:
3914+
client = CloudFilestoreManagerClient(
3915+
credentials=ga_credentials.AnonymousCredentials(), transport=transport
3916+
)
3917+
# Test client calls underlying transport.
3918+
with mock.patch.object(type(client.transport), "close") as close:
3919+
close.assert_not_called()
3920+
with client:
3921+
pass
3922+
close.assert_called()

0 commit comments

Comments
 (0)