From 815739b3bbe435be5b21732c84d3ada1007ab0ad Mon Sep 17 00:00:00 2001 From: "gcf-owl-bot[bot]" <78513119+gcf-owl-bot[bot]@users.noreply.github.com> Date: Thu, 7 Oct 2021 00:42:21 +0000 Subject: [PATCH] feat: add context manager support in client (#136) - [ ] Regenerate this pull request now. chore: fix docstring for first attribute of protos committer: @busunkim96 PiperOrigin-RevId: 401271153 Source-Link: https://github.com/googleapis/googleapis/commit/787f8c9a731f44e74a90b9847d48659ca9462d10 Source-Link: https://github.com/googleapis/googleapis-gen/commit/81decffe9fc72396a8153e756d1d67a6eecfd620 Copy-Tag: eyJwIjoiLmdpdGh1Yi8uT3dsQm90LnlhbWwiLCJoIjoiODFkZWNmZmU5ZmM3MjM5NmE4MTUzZTc1NmQxZDY3YTZlZWNmZDYyMCJ9 --- .../services/recommender/async_client.py | 6 +++ .../services/recommender/client.py | 18 +++++-- .../services/recommender/transports/base.py | 9 ++++ .../services/recommender/transports/grpc.py | 3 ++ .../recommender/transports/grpc_asyncio.py | 3 ++ google/cloud/recommender_v1/types/insight.py | 2 + .../recommender_v1/types/recommendation.py | 3 ++ .../types/recommender_service.py | 10 ++++ .../services/recommender/async_client.py | 6 +++ .../services/recommender/client.py | 18 +++++-- .../services/recommender/transports/base.py | 9 ++++ .../services/recommender/transports/grpc.py | 3 ++ .../recommender/transports/grpc_asyncio.py | 3 ++ .../recommender_v1beta1/types/insight.py | 2 + .../types/recommendation.py | 3 ++ .../types/recommender_service.py | 10 ++++ .../gapic/recommender_v1/test_recommender.py | 50 +++++++++++++++++++ .../recommender_v1beta1/test_recommender.py | 50 +++++++++++++++++++ 18 files changed, 200 insertions(+), 8 deletions(-) diff --git a/google/cloud/recommender_v1/services/recommender/async_client.py b/google/cloud/recommender_v1/services/recommender/async_client.py index dc5e9c0..6e7e6ff 100644 --- a/google/cloud/recommender_v1/services/recommender/async_client.py +++ b/google/cloud/recommender_v1/services/recommender/async_client.py @@ -971,6 +971,12 @@ async def mark_recommendation_failed( # 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( diff --git a/google/cloud/recommender_v1/services/recommender/client.py b/google/cloud/recommender_v1/services/recommender/client.py index 8b5d5bf..d9df1ab 100644 --- a/google/cloud/recommender_v1/services/recommender/client.py +++ b/google/cloud/recommender_v1/services/recommender/client.py @@ -410,10 +410,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_insights( @@ -1184,6 +1181,19 @@ def mark_recommendation_failed( # 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( diff --git a/google/cloud/recommender_v1/services/recommender/transports/base.py b/google/cloud/recommender_v1/services/recommender/transports/base.py index 1c302cf..6d4a9ea 100644 --- a/google/cloud/recommender_v1/services/recommender/transports/base.py +++ b/google/cloud/recommender_v1/services/recommender/transports/base.py @@ -238,6 +238,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 list_insights( self, diff --git a/google/cloud/recommender_v1/services/recommender/transports/grpc.py b/google/cloud/recommender_v1/services/recommender/transports/grpc.py index 88bcb32..39fa0d1 100644 --- a/google/cloud/recommender_v1/services/recommender/transports/grpc.py +++ b/google/cloud/recommender_v1/services/recommender/transports/grpc.py @@ -500,5 +500,8 @@ def mark_recommendation_failed( ) return self._stubs["mark_recommendation_failed"] + def close(self): + self.grpc_channel.close() + __all__ = ("RecommenderGrpcTransport",) diff --git a/google/cloud/recommender_v1/services/recommender/transports/grpc_asyncio.py b/google/cloud/recommender_v1/services/recommender/transports/grpc_asyncio.py index 7e74bcc..ec9bbb9 100644 --- a/google/cloud/recommender_v1/services/recommender/transports/grpc_asyncio.py +++ b/google/cloud/recommender_v1/services/recommender/transports/grpc_asyncio.py @@ -506,5 +506,8 @@ def mark_recommendation_failed( ) return self._stubs["mark_recommendation_failed"] + def close(self): + return self.grpc_channel.close() + __all__ = ("RecommenderGrpcAsyncIOTransport",) diff --git a/google/cloud/recommender_v1/types/insight.py b/google/cloud/recommender_v1/types/insight.py index 4a58ef3..8c09420 100644 --- a/google/cloud/recommender_v1/types/insight.py +++ b/google/cloud/recommender_v1/types/insight.py @@ -73,6 +73,7 @@ class Category(proto.Enum): class RecommendationReference(proto.Message): r"""Reference to an associated recommendation. + Attributes: recommendation (str): Recommendation resource name, e.g. @@ -102,6 +103,7 @@ class RecommendationReference(proto.Message): class InsightStateInfo(proto.Message): r"""Information related to insight state. + Attributes: state (google.cloud.recommender_v1.types.InsightStateInfo.State): Insight state. diff --git a/google/cloud/recommender_v1/types/recommendation.py b/google/cloud/recommender_v1/types/recommendation.py index e0af4be..631e2a6 100644 --- a/google/cloud/recommender_v1/types/recommendation.py +++ b/google/cloud/recommender_v1/types/recommendation.py @@ -87,6 +87,7 @@ class Recommendation(proto.Message): class InsightReference(proto.Message): r"""Reference to an associated insight. + Attributes: insight (str): Insight resource name, e.g. @@ -132,6 +133,7 @@ class RecommendationContent(proto.Message): class OperationGroup(proto.Message): r"""Group of operations that need to be performed atomically. + Attributes: operations (Sequence[google.cloud.recommender_v1.types.Operation]): List of operations across one or more @@ -301,6 +303,7 @@ class Category(proto.Enum): class RecommendationStateInfo(proto.Message): r"""Information for state. Contains state and metadata. + Attributes: state (google.cloud.recommender_v1.types.RecommendationStateInfo.State): The state of the recommendation, Eg ACTIVE, diff --git a/google/cloud/recommender_v1/types/recommender_service.py b/google/cloud/recommender_v1/types/recommender_service.py index 60da9e6..9b65552 100644 --- a/google/cloud/recommender_v1/types/recommender_service.py +++ b/google/cloud/recommender_v1/types/recommender_service.py @@ -38,6 +38,7 @@ class ListInsightsRequest(proto.Message): r"""Request for the ``ListInsights`` method. + Attributes: parent (str): Required. The container resource on which to execute the @@ -76,6 +77,7 @@ class ListInsightsRequest(proto.Message): class ListInsightsResponse(proto.Message): r"""Response to the ``ListInsights`` method. + Attributes: insights (Sequence[google.cloud.recommender_v1.types.Insight]): The set of insights for the ``parent`` resource. @@ -95,6 +97,7 @@ def raw_page(self): class GetInsightRequest(proto.Message): r"""Request to the ``GetInsight`` method. + Attributes: name (str): Required. Name of the insight. @@ -105,6 +108,7 @@ class GetInsightRequest(proto.Message): class MarkInsightAcceptedRequest(proto.Message): r"""Request for the ``MarkInsightAccepted`` method. + Attributes: name (str): Required. Name of the insight. @@ -123,6 +127,7 @@ class MarkInsightAcceptedRequest(proto.Message): class ListRecommendationsRequest(proto.Message): r"""Request for the ``ListRecommendations`` method. + Attributes: parent (str): Required. The container resource on which to execute the @@ -161,6 +166,7 @@ class ListRecommendationsRequest(proto.Message): class ListRecommendationsResponse(proto.Message): r"""Response to the ``ListRecommendations`` method. + Attributes: recommendations (Sequence[google.cloud.recommender_v1.types.Recommendation]): The set of recommendations for the ``parent`` resource. @@ -182,6 +188,7 @@ def raw_page(self): class GetRecommendationRequest(proto.Message): r"""Request to the ``GetRecommendation`` method. + Attributes: name (str): Required. Name of the recommendation. @@ -192,6 +199,7 @@ class GetRecommendationRequest(proto.Message): class MarkRecommendationClaimedRequest(proto.Message): r"""Request for the ``MarkRecommendationClaimed`` Method. + Attributes: name (str): Required. Name of the recommendation. @@ -212,6 +220,7 @@ class MarkRecommendationClaimedRequest(proto.Message): class MarkRecommendationSucceededRequest(proto.Message): r"""Request for the ``MarkRecommendationSucceeded`` Method. + Attributes: name (str): Required. Name of the recommendation. @@ -232,6 +241,7 @@ class MarkRecommendationSucceededRequest(proto.Message): class MarkRecommendationFailedRequest(proto.Message): r"""Request for the ``MarkRecommendationFailed`` Method. + Attributes: name (str): Required. Name of the recommendation. diff --git a/google/cloud/recommender_v1beta1/services/recommender/async_client.py b/google/cloud/recommender_v1beta1/services/recommender/async_client.py index 19655d2..9cd57be 100644 --- a/google/cloud/recommender_v1beta1/services/recommender/async_client.py +++ b/google/cloud/recommender_v1beta1/services/recommender/async_client.py @@ -971,6 +971,12 @@ async def mark_recommendation_failed( # 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( diff --git a/google/cloud/recommender_v1beta1/services/recommender/client.py b/google/cloud/recommender_v1beta1/services/recommender/client.py index 40f9799..490f731 100644 --- a/google/cloud/recommender_v1beta1/services/recommender/client.py +++ b/google/cloud/recommender_v1beta1/services/recommender/client.py @@ -410,10 +410,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_insights( @@ -1184,6 +1181,19 @@ def mark_recommendation_failed( # 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( diff --git a/google/cloud/recommender_v1beta1/services/recommender/transports/base.py b/google/cloud/recommender_v1beta1/services/recommender/transports/base.py index ba05958..647f6f1 100644 --- a/google/cloud/recommender_v1beta1/services/recommender/transports/base.py +++ b/google/cloud/recommender_v1beta1/services/recommender/transports/base.py @@ -238,6 +238,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 list_insights( self, diff --git a/google/cloud/recommender_v1beta1/services/recommender/transports/grpc.py b/google/cloud/recommender_v1beta1/services/recommender/transports/grpc.py index affa28d..ea135a5 100644 --- a/google/cloud/recommender_v1beta1/services/recommender/transports/grpc.py +++ b/google/cloud/recommender_v1beta1/services/recommender/transports/grpc.py @@ -500,5 +500,8 @@ def mark_recommendation_failed( ) return self._stubs["mark_recommendation_failed"] + def close(self): + self.grpc_channel.close() + __all__ = ("RecommenderGrpcTransport",) diff --git a/google/cloud/recommender_v1beta1/services/recommender/transports/grpc_asyncio.py b/google/cloud/recommender_v1beta1/services/recommender/transports/grpc_asyncio.py index abbeac6..c9f802e 100644 --- a/google/cloud/recommender_v1beta1/services/recommender/transports/grpc_asyncio.py +++ b/google/cloud/recommender_v1beta1/services/recommender/transports/grpc_asyncio.py @@ -506,5 +506,8 @@ def mark_recommendation_failed( ) return self._stubs["mark_recommendation_failed"] + def close(self): + return self.grpc_channel.close() + __all__ = ("RecommenderGrpcAsyncIOTransport",) diff --git a/google/cloud/recommender_v1beta1/types/insight.py b/google/cloud/recommender_v1beta1/types/insight.py index da4906f..2e5cf5a 100644 --- a/google/cloud/recommender_v1beta1/types/insight.py +++ b/google/cloud/recommender_v1beta1/types/insight.py @@ -74,6 +74,7 @@ class Category(proto.Enum): class RecommendationReference(proto.Message): r"""Reference to an associated recommendation. + Attributes: recommendation (str): Recommendation resource name, e.g. @@ -103,6 +104,7 @@ class RecommendationReference(proto.Message): class InsightStateInfo(proto.Message): r"""Information related to insight state. + Attributes: state (google.cloud.recommender_v1beta1.types.InsightStateInfo.State): Insight state. diff --git a/google/cloud/recommender_v1beta1/types/recommendation.py b/google/cloud/recommender_v1beta1/types/recommendation.py index 67fbf50..0ec6a65 100644 --- a/google/cloud/recommender_v1beta1/types/recommendation.py +++ b/google/cloud/recommender_v1beta1/types/recommendation.py @@ -87,6 +87,7 @@ class Recommendation(proto.Message): class InsightReference(proto.Message): r"""Reference to an associated insight. + Attributes: insight (str): Insight resource name, e.g. @@ -132,6 +133,7 @@ class RecommendationContent(proto.Message): class OperationGroup(proto.Message): r"""Group of operations that need to be performed atomically. + Attributes: operations (Sequence[google.cloud.recommender_v1beta1.types.Operation]): List of operations across one or more @@ -301,6 +303,7 @@ class Category(proto.Enum): class RecommendationStateInfo(proto.Message): r"""Information for state. Contains state and metadata. + Attributes: state (google.cloud.recommender_v1beta1.types.RecommendationStateInfo.State): The state of the recommendation, Eg ACTIVE, diff --git a/google/cloud/recommender_v1beta1/types/recommender_service.py b/google/cloud/recommender_v1beta1/types/recommender_service.py index 404a4ef..c6aeffc 100644 --- a/google/cloud/recommender_v1beta1/types/recommender_service.py +++ b/google/cloud/recommender_v1beta1/types/recommender_service.py @@ -38,6 +38,7 @@ class ListInsightsRequest(proto.Message): r"""Request for the ``ListInsights`` method. + Attributes: parent (str): Required. The container resource on which to execute the @@ -76,6 +77,7 @@ class ListInsightsRequest(proto.Message): class ListInsightsResponse(proto.Message): r"""Response to the ``ListInsights`` method. + Attributes: insights (Sequence[google.cloud.recommender_v1beta1.types.Insight]): The set of insights for the ``parent`` resource. @@ -95,6 +97,7 @@ def raw_page(self): class GetInsightRequest(proto.Message): r"""Request to the ``GetInsight`` method. + Attributes: name (str): Required. Name of the insight. @@ -105,6 +108,7 @@ class GetInsightRequest(proto.Message): class MarkInsightAcceptedRequest(proto.Message): r"""Request for the ``MarkInsightAccepted`` method. + Attributes: name (str): Required. Name of the insight. @@ -123,6 +127,7 @@ class MarkInsightAcceptedRequest(proto.Message): class ListRecommendationsRequest(proto.Message): r"""Request for the ``ListRecommendations`` method. + Attributes: parent (str): Required. The container resource on which to execute the @@ -161,6 +166,7 @@ class ListRecommendationsRequest(proto.Message): class ListRecommendationsResponse(proto.Message): r"""Response to the ``ListRecommendations`` method. + Attributes: recommendations (Sequence[google.cloud.recommender_v1beta1.types.Recommendation]): The set of recommendations for the ``parent`` resource. @@ -182,6 +188,7 @@ def raw_page(self): class GetRecommendationRequest(proto.Message): r"""Request to the ``GetRecommendation`` method. + Attributes: name (str): Required. Name of the recommendation. @@ -192,6 +199,7 @@ class GetRecommendationRequest(proto.Message): class MarkRecommendationClaimedRequest(proto.Message): r"""Request for the ``MarkRecommendationClaimed`` Method. + Attributes: name (str): Required. Name of the recommendation. @@ -212,6 +220,7 @@ class MarkRecommendationClaimedRequest(proto.Message): class MarkRecommendationSucceededRequest(proto.Message): r"""Request for the ``MarkRecommendationSucceeded`` Method. + Attributes: name (str): Required. Name of the recommendation. @@ -232,6 +241,7 @@ class MarkRecommendationSucceededRequest(proto.Message): class MarkRecommendationFailedRequest(proto.Message): r"""Request for the ``MarkRecommendationFailed`` Method. + Attributes: name (str): Required. Name of the recommendation. diff --git a/tests/unit/gapic/recommender_v1/test_recommender.py b/tests/unit/gapic/recommender_v1/test_recommender.py index 4ff21c4..f40e63b 100644 --- a/tests/unit/gapic/recommender_v1/test_recommender.py +++ b/tests/unit/gapic/recommender_v1/test_recommender.py @@ -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.recommender_v1.services.recommender import RecommenderAsyncClient @@ -2749,6 +2750,9 @@ def test_recommender_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_recommender_base_transport_with_credentials_file(): @@ -3295,3 +3299,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 = RecommenderAsyncClient( + 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 = RecommenderClient( + 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 = RecommenderClient( + 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() diff --git a/tests/unit/gapic/recommender_v1beta1/test_recommender.py b/tests/unit/gapic/recommender_v1beta1/test_recommender.py index b8905f8..c9094b9 100644 --- a/tests/unit/gapic/recommender_v1beta1/test_recommender.py +++ b/tests/unit/gapic/recommender_v1beta1/test_recommender.py @@ -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.recommender_v1beta1.services.recommender import RecommenderAsyncClient @@ -2749,6 +2750,9 @@ def test_recommender_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_recommender_base_transport_with_credentials_file(): @@ -3295,3 +3299,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 = RecommenderAsyncClient( + 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 = RecommenderClient( + 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 = RecommenderClient( + 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()