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

feat: add context manager support in client #11

Merged
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
Expand Up @@ -1198,6 +1198,12 @@ async def update_backup(
# Done; return the response.
return response

async def __aenter__(self):
return self

async def __aexit__(self, exc_type, exc, tb):
await self.transport.close()


try:
DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(
Expand Down
Expand Up @@ -392,10 +392,7 @@ def __init__(
client_cert_source_for_mtls=client_cert_source_func,
quota_project_id=client_options.quota_project_id,
client_info=client_info,
always_use_jwt_access=(
Transport == type(self).get_transport_class("grpc")
or Transport == type(self).get_transport_class("grpc_asyncio")
),
always_use_jwt_access=True,
)

def list_instances(
Expand Down Expand Up @@ -1365,6 +1362,19 @@ def update_backup(
# Done; return the response.
return response

def __enter__(self):
return self

def __exit__(self, type, value, traceback):
"""Releases underlying transport's resources.

.. warning::
ONLY use as a context manager if the transport is NOT shared
with other clients! Exiting the with block will CLOSE the transport
and may cause errors in other clients!
"""
self.transport.close()


try:
DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(
Expand Down
Expand Up @@ -233,6 +233,15 @@ def _prep_wrapped_messages(self, client_info):
),
}

def close(self):
"""Closes resources associated with the transport.

.. warning::
Only call this method if the transport is NOT shared
with other clients - this may cause errors in other clients!
"""
raise NotImplementedError()

@property
def operations_client(self) -> operations_v1.OperationsClient:
"""Return the client designed to process long-running operations."""
Expand Down
Expand Up @@ -584,5 +584,8 @@ def update_backup(
)
return self._stubs["update_backup"]

def close(self):
self.grpc_channel.close()


__all__ = ("CloudFilestoreManagerGrpcTransport",)
Expand Up @@ -598,5 +598,8 @@ def update_backup(
)
return self._stubs["update_backup"]

def close(self):
return self.grpc_channel.close()


__all__ = ("CloudFilestoreManagerGrpcAsyncIOTransport",)
16 changes: 16 additions & 0 deletions google/cloud/filestore_v1/types/cloud_filestore_service.py
Expand Up @@ -47,6 +47,7 @@

class NetworkConfig(proto.Message):
r"""Network configuration for the instance.

Attributes:
network (str):
The name of the Google Compute Engine `VPC
Expand Down Expand Up @@ -84,6 +85,7 @@ class AddressMode(proto.Enum):

class FileShareConfig(proto.Message):
r"""File share configuration for the instance.

Attributes:
name (str):
The name of the file share (must be 16
Expand Down Expand Up @@ -111,6 +113,7 @@ class FileShareConfig(proto.Message):

class NfsExportOptions(proto.Message):
r"""NFS export options specifications.

Attributes:
ip_ranges (Sequence[str]):
List of either an IPv4 addresses in the format
Expand Down Expand Up @@ -162,6 +165,7 @@ class SquashMode(proto.Enum):

class Instance(proto.Message):
r"""A Cloud Filestore instance.

Attributes:
name (str):
Output only. The resource name of the instance, in the
Expand Down Expand Up @@ -237,6 +241,7 @@ class Tier(proto.Enum):

class CreateInstanceRequest(proto.Message):
r"""CreateInstanceRequest creates an instance.

Attributes:
parent (str):
Required. The instance's project and location, in the format
Expand All @@ -259,6 +264,7 @@ class CreateInstanceRequest(proto.Message):

class GetInstanceRequest(proto.Message):
r"""GetInstanceRequest gets the state of an instance.

Attributes:
name (str):
Required. The instance resource name, in the format
Expand All @@ -270,6 +276,7 @@ class GetInstanceRequest(proto.Message):

class UpdateInstanceRequest(proto.Message):
r"""UpdateInstanceRequest updates the settings of an instance.

Attributes:
update_mask (google.protobuf.field_mask_pb2.FieldMask):
Mask of fields to update. At least one path must be supplied
Expand Down Expand Up @@ -313,6 +320,7 @@ class RestoreInstanceRequest(proto.Message):

class DeleteInstanceRequest(proto.Message):
r"""DeleteInstanceRequest deletes an instance.

Attributes:
name (str):
Required. The instance resource name, in the format
Expand All @@ -324,6 +332,7 @@ class DeleteInstanceRequest(proto.Message):

class ListInstancesRequest(proto.Message):
r"""ListInstancesRequest lists instances.

Attributes:
parent (str):
Required. The project and location for which to retrieve
Expand Down Expand Up @@ -353,6 +362,7 @@ class ListInstancesRequest(proto.Message):

class ListInstancesResponse(proto.Message):
r"""ListInstancesResponse is the result of ListInstancesRequest.

Attributes:
instances (Sequence[google.cloud.filestore_v1.types.Instance]):
A list of instances in the project for the specified
Expand Down Expand Up @@ -382,6 +392,7 @@ def raw_page(self):

class Backup(proto.Message):
r"""A Cloud Filestore backup.

Attributes:
name (str):
Output only. The resource name of the backup, in the format
Expand Down Expand Up @@ -455,6 +466,7 @@ class State(proto.Enum):

class CreateBackupRequest(proto.Message):
r"""CreateBackupRequest creates a backup.

Attributes:
parent (str):
Required. The backup's project and location, in the format
Expand All @@ -481,6 +493,7 @@ class CreateBackupRequest(proto.Message):

class DeleteBackupRequest(proto.Message):
r"""DeleteBackupRequest deletes a backup.

Attributes:
name (str):
Required. The backup resource name, in the format
Expand Down Expand Up @@ -511,6 +524,7 @@ class UpdateBackupRequest(proto.Message):

class GetBackupRequest(proto.Message):
r"""GetBackupRequest gets the state of a backup.

Attributes:
name (str):
Required. The backup resource name, in the format
Expand All @@ -522,6 +536,7 @@ class GetBackupRequest(proto.Message):

class ListBackupsRequest(proto.Message):
r"""ListBackupsRequest lists backups.

Attributes:
parent (str):
Required. The project and location for which to retrieve
Expand Down Expand Up @@ -551,6 +566,7 @@ class ListBackupsRequest(proto.Message):

class ListBackupsResponse(proto.Message):
r"""ListBackupsResponse is the result of ListBackupsRequest.

Attributes:
backups (Sequence[google.cloud.filestore_v1.types.Backup]):
A list of backups in the project for the specified location.
Expand Down
50 changes: 50 additions & 0 deletions tests/unit/gapic/filestore_v1/test_cloud_filestore_manager.py
Expand Up @@ -32,6 +32,7 @@
from google.api_core import grpc_helpers_async
from google.api_core import operation_async # type: ignore
from google.api_core import operations_v1
from google.api_core import path_template
from google.auth import credentials as ga_credentials
from google.auth.exceptions import MutualTLSChannelError
from google.cloud.common.types import operation_metadata as operation_metadata_pb2 # type: ignore
Expand Down Expand Up @@ -3331,6 +3332,9 @@ def test_cloud_filestore_manager_base_transport():
with pytest.raises(NotImplementedError):
getattr(transport, method)(request=object())

with pytest.raises(NotImplementedError):
transport.close()

# Additionally, the LRO client (a property) should
# also raise NotImplementedError
with pytest.raises(NotImplementedError):
Expand Down Expand Up @@ -3870,3 +3874,49 @@ def test_client_withDEFAULT_CLIENT_INFO():
credentials=ga_credentials.AnonymousCredentials(), client_info=client_info,
)
prep.assert_called_once_with(client_info)


@pytest.mark.asyncio
async def test_transport_close_async():
client = CloudFilestoreManagerAsyncClient(
credentials=ga_credentials.AnonymousCredentials(), transport="grpc_asyncio",
)
with mock.patch.object(
type(getattr(client.transport, "grpc_channel")), "close"
) as close:
async with client:
close.assert_not_called()
close.assert_called_once()


def test_transport_close():
transports = {
"grpc": "_grpc_channel",
}

for transport, close_name in transports.items():
client = CloudFilestoreManagerClient(
credentials=ga_credentials.AnonymousCredentials(), transport=transport
)
with mock.patch.object(
type(getattr(client.transport, close_name)), "close"
) as close:
with client:
close.assert_not_called()
close.assert_called_once()


def test_client_ctx():
transports = [
"grpc",
]
for transport in transports:
client = CloudFilestoreManagerClient(
credentials=ga_credentials.AnonymousCredentials(), transport=transport
)
# Test client calls underlying transport.
with mock.patch.object(type(client.transport), "close") as close:
close.assert_not_called()
with client:
pass
close.assert_called()