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

feat: add context manager support in client #74

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 @@ -1588,6 +1588,12 @@ async def delete_api_config(
# 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
18 changes: 14 additions & 4 deletions google/cloud/apigateway_v1/services/api_gateway_service/client.py
Expand Up @@ -424,10 +424,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_gateways(
Expand Down Expand Up @@ -1755,6 +1752,19 @@ def delete_api_config(
# 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 @@ -302,6 +302,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 @@ -635,5 +635,8 @@ def delete_api_config(
)
return self._stubs["delete_api_config"]

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


__all__ = ("ApiGatewayServiceGrpcTransport",)
Expand Up @@ -656,5 +656,8 @@ def delete_api_config(
)
return self._stubs["delete_api_config"]

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


__all__ = ("ApiGatewayServiceGrpcAsyncIOTransport",)
24 changes: 24 additions & 0 deletions google/cloud/apigateway_v1/types/apigateway.py
Expand Up @@ -50,6 +50,7 @@

class Api(proto.Message):
r"""An API that can be served by one or more Gateways.

Attributes:
name (str):
Output only. Resource name of the API.
Expand Down Expand Up @@ -166,6 +167,7 @@ class State(proto.Enum):

class File(proto.Message):
r"""A lightweight description of a file.

Attributes:
path (str):
The file path (full or relative path). This
Expand All @@ -180,6 +182,7 @@ class File(proto.Message):

class OpenApiDocument(proto.Message):
r"""An OpenAPI Specification Document describing an API.

Attributes:
document (google.cloud.apigateway_v1.types.ApiConfig.File):
The OpenAPI Specification document file.
Expand All @@ -189,6 +192,7 @@ class OpenApiDocument(proto.Message):

class GrpcServiceDefinition(proto.Message):
r"""A gRPC service definition.

Attributes:
file_descriptor_set (google.cloud.apigateway_v1.types.ApiConfig.File):
Input only. File descriptor set, generated by protoc.
Expand Down Expand Up @@ -285,6 +289,7 @@ class State(proto.Enum):

class ListGatewaysRequest(proto.Message):
r"""Request message for ApiGatewayService.ListGateways

Attributes:
parent (str):
Required. Parent resource of the Gateway, of the form:
Expand All @@ -308,6 +313,7 @@ class ListGatewaysRequest(proto.Message):

class ListGatewaysResponse(proto.Message):
r"""Response message for ApiGatewayService.ListGateways

Attributes:
gateways (Sequence[google.cloud.apigateway_v1.types.Gateway]):
Gateways.
Expand All @@ -328,6 +334,7 @@ def raw_page(self):

class GetGatewayRequest(proto.Message):
r"""Request message for ApiGatewayService.GetGateway

Attributes:
name (str):
Required. Resource name of the form:
Expand All @@ -339,6 +346,7 @@ class GetGatewayRequest(proto.Message):

class CreateGatewayRequest(proto.Message):
r"""Request message for ApiGatewayService.CreateGateway

Attributes:
parent (str):
Required. Parent resource of the Gateway, of the form:
Expand All @@ -358,6 +366,7 @@ class CreateGatewayRequest(proto.Message):

class UpdateGatewayRequest(proto.Message):
r"""Request message for ApiGatewayService.UpdateGateway

Attributes:
update_mask (google.protobuf.field_mask_pb2.FieldMask):
Field mask is used to specify the fields to be overwritten
Expand All @@ -378,6 +387,7 @@ class UpdateGatewayRequest(proto.Message):

class DeleteGatewayRequest(proto.Message):
r"""Request message for ApiGatewayService.DeleteGateway

Attributes:
name (str):
Required. Resource name of the form:
Expand All @@ -389,6 +399,7 @@ class DeleteGatewayRequest(proto.Message):

class ListApisRequest(proto.Message):
r"""Request message for ApiGatewayService.ListApis

Attributes:
parent (str):
Required. Parent resource of the API, of the form:
Expand All @@ -412,6 +423,7 @@ class ListApisRequest(proto.Message):

class ListApisResponse(proto.Message):
r"""Response message for ApiGatewayService.ListApis

Attributes:
apis (Sequence[google.cloud.apigateway_v1.types.Api]):
APIs.
Expand All @@ -432,6 +444,7 @@ def raw_page(self):

class GetApiRequest(proto.Message):
r"""Request message for ApiGatewayService.GetApi

Attributes:
name (str):
Required. Resource name of the form:
Expand All @@ -443,6 +456,7 @@ class GetApiRequest(proto.Message):

class CreateApiRequest(proto.Message):
r"""Request message for ApiGatewayService.CreateApi

Attributes:
parent (str):
Required. Parent resource of the API, of the form:
Expand All @@ -462,6 +476,7 @@ class CreateApiRequest(proto.Message):

class UpdateApiRequest(proto.Message):
r"""Request message for ApiGatewayService.UpdateApi

Attributes:
update_mask (google.protobuf.field_mask_pb2.FieldMask):
Field mask is used to specify the fields to be overwritten
Expand All @@ -482,6 +497,7 @@ class UpdateApiRequest(proto.Message):

class DeleteApiRequest(proto.Message):
r"""Request message for ApiGatewayService.DeleteApi

Attributes:
name (str):
Required. Resource name of the form:
Expand All @@ -493,6 +509,7 @@ class DeleteApiRequest(proto.Message):

class ListApiConfigsRequest(proto.Message):
r"""Request message for ApiGatewayService.ListApiConfigs

Attributes:
parent (str):
Required. Parent resource of the API Config, of the form:
Expand All @@ -516,6 +533,7 @@ class ListApiConfigsRequest(proto.Message):

class ListApiConfigsResponse(proto.Message):
r"""Response message for ApiGatewayService.ListApiConfigs

Attributes:
api_configs (Sequence[google.cloud.apigateway_v1.types.ApiConfig]):
API Configs.
Expand All @@ -536,6 +554,7 @@ def raw_page(self):

class GetApiConfigRequest(proto.Message):
r"""Request message for ApiGatewayService.GetApiConfig

Attributes:
name (str):
Required. Resource name of the form:
Expand All @@ -559,6 +578,7 @@ class ConfigView(proto.Enum):

class CreateApiConfigRequest(proto.Message):
r"""Request message for ApiGatewayService.CreateApiConfig

Attributes:
parent (str):
Required. Parent resource of the API Config, of the form:
Expand All @@ -578,6 +598,7 @@ class CreateApiConfigRequest(proto.Message):

class UpdateApiConfigRequest(proto.Message):
r"""Request message for ApiGatewayService.UpdateApiConfig

Attributes:
update_mask (google.protobuf.field_mask_pb2.FieldMask):
Field mask is used to specify the fields to be overwritten
Expand All @@ -598,6 +619,7 @@ class UpdateApiConfigRequest(proto.Message):

class DeleteApiConfigRequest(proto.Message):
r"""Request message for ApiGatewayService.DeleteApiConfig

Attributes:
name (str):
Required. Resource name of the form:
Expand All @@ -609,6 +631,7 @@ class DeleteApiConfigRequest(proto.Message):

class OperationMetadata(proto.Message):
r"""Represents the metadata of the long-running operation.

Attributes:
create_time (google.protobuf.timestamp_pb2.Timestamp):
Output only. The time the operation was
Expand Down Expand Up @@ -641,6 +664,7 @@ class OperationMetadata(proto.Message):

class Diagnostic(proto.Message):
r"""Diagnostic information from configuration processing.

Attributes:
location (str):
Location of the diagnostic.
Expand Down
50 changes: 50 additions & 0 deletions tests/unit/gapic/apigateway_v1/test_api_gateway_service.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.apigateway_v1.services.api_gateway_service import (
Expand Down Expand Up @@ -4248,6 +4249,9 @@ def test_api_gateway_service_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 @@ -4871,3 +4875,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 = ApiGatewayServiceAsyncClient(
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 = ApiGatewayServiceClient(
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 = ApiGatewayServiceClient(
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()