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 #73

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 @@ -217,6 +217,12 @@ async def allocate_quota(
# 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/servicecontrol_v1/services/quota_controller/client.py
Expand Up @@ -332,10 +332,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 allocate_quota(
Expand Down Expand Up @@ -393,6 +390,19 @@ def allocate_quota(
# 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 @@ -162,6 +162,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 allocate_quota(
self,
Expand Down
Expand Up @@ -267,5 +267,8 @@ def allocate_quota(
)
return self._stubs["allocate_quota"]

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


__all__ = ("QuotaControllerGrpcTransport",)
Expand Up @@ -271,5 +271,8 @@ def allocate_quota(
)
return self._stubs["allocate_quota"]

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


__all__ = ("QuotaControllerGrpcAsyncIOTransport",)
Expand Up @@ -293,6 +293,12 @@ async def report(
# 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 @@ -334,10 +334,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 check(
Expand Down Expand Up @@ -463,6 +460,19 @@ def report(
# 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 @@ -176,6 +176,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 check(
self,
Expand Down
Expand Up @@ -317,5 +317,8 @@ def report(
)
return self._stubs["report"]

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


__all__ = ("ServiceControllerGrpcTransport",)
Expand Up @@ -322,5 +322,8 @@ def report(
)
return self._stubs["report"]

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


__all__ = ("ServiceControllerGrpcAsyncIOTransport",)
3 changes: 3 additions & 0 deletions google/cloud/servicecontrol_v1/types/distribution.py
Expand Up @@ -79,6 +79,7 @@ class Distribution(proto.Message):

class LinearBuckets(proto.Message):
r"""Describing buckets with constant width.

Attributes:
num_finite_buckets (int):
The number of finite buckets. With the underflow and
Expand All @@ -101,6 +102,7 @@ class LinearBuckets(proto.Message):

class ExponentialBuckets(proto.Message):
r"""Describing buckets with exponentially growing width.

Attributes:
num_finite_buckets (int):
The number of finite buckets. With the underflow and
Expand All @@ -124,6 +126,7 @@ class ExponentialBuckets(proto.Message):

class ExplicitBuckets(proto.Message):
r"""Describing buckets with arbitrary user-provided width.

Attributes:
bounds (Sequence[float]):
'bound' is a list of strictly increasing boundaries between
Expand Down
1 change: 1 addition & 0 deletions google/cloud/servicecontrol_v1/types/log_entry.py
Expand Up @@ -30,6 +30,7 @@

class LogEntry(proto.Message):
r"""An individual log entry.

Attributes:
name (str):
Required. The log to which this log entry belongs. Examples:
Expand Down
1 change: 1 addition & 0 deletions google/cloud/servicecontrol_v1/types/metric_value.py
Expand Up @@ -26,6 +26,7 @@

class MetricValue(proto.Message):
r"""Represents a single metric value.

Attributes:
labels (Sequence[google.cloud.servicecontrol_v1.types.MetricValue.LabelsEntry]):
The labels describing the metric value. See comments on
Expand Down
1 change: 1 addition & 0 deletions google/cloud/servicecontrol_v1/types/operation.py
Expand Up @@ -28,6 +28,7 @@

class Operation(proto.Message):
r"""Represents information regarding an operation.

Attributes:
operation_id (str):
Identity of the operation. This must be
Expand Down
3 changes: 3 additions & 0 deletions google/cloud/servicecontrol_v1/types/quota_controller.py
Expand Up @@ -32,6 +32,7 @@

class AllocateQuotaRequest(proto.Message):
r"""Request message for the AllocateQuota method.

Attributes:
service_name (str):
Name of the service as specified in the service
Expand All @@ -56,6 +57,7 @@ class AllocateQuotaRequest(proto.Message):

class QuotaOperation(proto.Message):
r"""Represents information regarding a quota operation.

Attributes:
operation_id (str):
Identity of the operation. This is expected to be unique
Expand Down Expand Up @@ -128,6 +130,7 @@ class QuotaMode(proto.Enum):

class AllocateQuotaResponse(proto.Message):
r"""Response message for the AllocateQuota method.

Attributes:
operation_id (str):
The same operation_id value used in the
Expand Down
6 changes: 6 additions & 0 deletions google/cloud/servicecontrol_v1/types/service_controller.py
Expand Up @@ -28,6 +28,7 @@

class CheckRequest(proto.Message):
r"""Request message for the Check method.

Attributes:
service_name (str):
The service name as specified in its service configuration.
Expand All @@ -53,6 +54,7 @@ class CheckRequest(proto.Message):

class CheckResponse(proto.Message):
r"""Response message for the Check method.

Attributes:
operation_id (str):
The same operation_id value used in the
Expand All @@ -77,6 +79,7 @@ class CheckResponse(proto.Message):

class CheckInfo(proto.Message):
r"""Contains additional information about the check operation.

Attributes:
unused_arguments (Sequence[str]):
A list of fields and label keys that are
Expand All @@ -94,6 +97,7 @@ class CheckInfo(proto.Message):

class ConsumerInfo(proto.Message):
r"""``ConsumerInfo`` provides information about the consumer.

Attributes:
project_number (int):
The Google cloud project number, e.g.
Expand Down Expand Up @@ -140,6 +144,7 @@ class ConsumerType(proto.Enum):

class ReportRequest(proto.Message):
r"""Request message for the Report method.

Attributes:
service_name (str):
The service name as specified in its service configuration.
Expand Down Expand Up @@ -178,6 +183,7 @@ class ReportRequest(proto.Message):

class ReportResponse(proto.Message):
r"""Response message for the Report method.

Attributes:
report_errors (Sequence[google.cloud.servicecontrol_v1.types.ReportResponse.ReportError]):
Partial failures, one for each ``Operation`` in the request
Expand Down
50 changes: 50 additions & 0 deletions tests/unit/gapic/servicecontrol_v1/test_quota_controller.py
Expand Up @@ -30,6 +30,7 @@
from google.api_core import gapic_v1
from google.api_core import grpc_helpers
from google.api_core import grpc_helpers_async
from google.api_core import path_template
from google.auth import credentials as ga_credentials
from google.auth.exceptions import MutualTLSChannelError
from google.cloud.servicecontrol_v1.services.quota_controller import (
Expand Down Expand Up @@ -702,6 +703,9 @@ def test_quota_controller_base_transport():
with pytest.raises(NotImplementedError):
getattr(transport, method)(request=object())

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


@requires_google_auth_gte_1_25_0
def test_quota_controller_base_transport_with_credentials_file():
Expand Down Expand Up @@ -1181,3 +1185,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 = QuotaControllerAsyncClient(
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 = QuotaControllerClient(
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 = QuotaControllerClient(
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()