Skip to content
This repository has been archived by the owner on Jan 17, 2024. It is now read-only.

feat: add context manager support in client #120

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 @@ -444,6 +444,12 @@ async def update_dashboard(
# 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 @@ -362,10 +362,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 create_dashboard(
Expand Down Expand Up @@ -648,6 +645,19 @@ def update_dashboard(
# 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 @@ -178,6 +178,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 create_dashboard(
self,
Expand Down
Expand Up @@ -387,5 +387,8 @@ def update_dashboard(
)
return self._stubs["update_dashboard"]

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


__all__ = ("DashboardsServiceGrpcTransport",)
Expand Up @@ -398,5 +398,8 @@ def update_dashboard(
)
return self._stubs["update_dashboard"]

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


__all__ = ("DashboardsServiceGrpcAsyncIOTransport",)
1 change: 1 addition & 0 deletions google/cloud/monitoring_dashboard_v1/types/alertchart.py
Expand Up @@ -23,6 +23,7 @@

class AlertChart(proto.Message):
r"""A chart that displays alert policy data.

Attributes:
name (str):
Required. The resource name of the alert policy. The format
Expand Down
Expand Up @@ -33,6 +33,7 @@

class CreateDashboardRequest(proto.Message):
r"""The ``CreateDashboard`` request.

Attributes:
parent (str):
Required. The project on which to execute the request. The
Expand All @@ -59,6 +60,7 @@ class CreateDashboardRequest(proto.Message):

class ListDashboardsRequest(proto.Message):
r"""The ``ListDashboards`` request.

Attributes:
parent (str):
Required. The scope of the dashboards to list. The format
Expand All @@ -85,6 +87,7 @@ class ListDashboardsRequest(proto.Message):

class ListDashboardsResponse(proto.Message):
r"""The ``ListDashboards`` request.

Attributes:
dashboards (Sequence[google.cloud.monitoring_dashboard_v1.types.Dashboard]):
The list of requested dashboards.
Expand All @@ -107,6 +110,7 @@ def raw_page(self):

class GetDashboardRequest(proto.Message):
r"""The ``GetDashboard`` request.

Attributes:
name (str):
Required. The resource name of the Dashboard. The format is
Expand All @@ -122,6 +126,7 @@ class GetDashboardRequest(proto.Message):

class DeleteDashboardRequest(proto.Message):
r"""The ``DeleteDashboard`` request.

Attributes:
name (str):
Required. The resource name of the Dashboard. The format is:
Expand All @@ -136,6 +141,7 @@ class DeleteDashboardRequest(proto.Message):

class UpdateDashboardRequest(proto.Message):
r"""The ``UpdateDashboard`` request.

Attributes:
dashboard (google.cloud.monitoring_dashboard_v1.types.Dashboard):
Required. The dashboard that will replace the
Expand Down
2 changes: 2 additions & 0 deletions google/cloud/monitoring_dashboard_v1/types/layouts.py
Expand Up @@ -104,6 +104,7 @@ class RowLayout(proto.Message):

class Row(proto.Message):
r"""Defines the layout properties and content for a row.

Attributes:
weight (int):
The relative weight of this row. The row
Expand Down Expand Up @@ -137,6 +138,7 @@ class ColumnLayout(proto.Message):

class Column(proto.Message):
r"""Defines the layout properties and content for a column.

Attributes:
weight (int):
The relative weight of this column. The
Expand Down
1 change: 1 addition & 0 deletions google/cloud/monitoring_dashboard_v1/types/metrics.py
Expand Up @@ -176,6 +176,7 @@ class RatioPart(proto.Message):

class Threshold(proto.Message):
r"""Defines a threshold for categorizing time series values.

Attributes:
label (str):
A label for the threshold.
Expand Down
1 change: 1 addition & 0 deletions google/cloud/monitoring_dashboard_v1/types/text.py
Expand Up @@ -23,6 +23,7 @@

class Text(proto.Message):
r"""A widget that displays textual content.

Attributes:
content (str):
The text content to be displayed.
Expand Down
4 changes: 4 additions & 0 deletions google/cloud/monitoring_dashboard_v1/types/xychart.py
Expand Up @@ -26,6 +26,7 @@

class XyChart(proto.Message):
r"""A chart that displays data on a 2D (X and Y axes) plane.

Attributes:
data_sets (Sequence[google.cloud.monitoring_dashboard_v1.types.XyChart.DataSet]):
Required. The data displayed in this chart.
Expand All @@ -50,6 +51,7 @@ class XyChart(proto.Message):

class DataSet(proto.Message):
r"""Groups a time series query definition with charting options.

Attributes:
time_series_query (google.cloud.monitoring_dashboard_v1.types.TimeSeriesQuery):
Required. Fields for querying time series
Expand Down Expand Up @@ -90,6 +92,7 @@ class PlotType(proto.Enum):

class Axis(proto.Message):
r"""A chart axis.

Attributes:
label (str):
The label of the axis.
Expand Down Expand Up @@ -121,6 +124,7 @@ class Scale(proto.Enum):

class ChartOptions(proto.Message):
r"""Options to control visual rendering of a chart.

Attributes:
mode (google.cloud.monitoring_dashboard_v1.types.ChartOptions.Mode):
The chart mode.
Expand Down
50 changes: 50 additions & 0 deletions tests/unit/gapic/dashboard_v1/test_dashboards_service.py
Expand Up @@ -29,6 +29,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.monitoring_dashboard_v1.services.dashboards_service import (
Expand Down Expand Up @@ -1479,6 +1480,9 @@ def test_dashboards_service_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_dashboards_service_base_transport_with_credentials_file():
Expand Down Expand Up @@ -2016,3 +2020,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 = DashboardsServiceAsyncClient(
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 = DashboardsServiceClient(
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 = DashboardsServiceClient(
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()