From a110fbb8f08487c901a753fb7ee4e94f00dbf1db Mon Sep 17 00:00:00 2001 From: ohmayr Date: Tue, 10 Sep 2024 22:29:20 +0000 Subject: [PATCH 1/3] feat: wrap async rest methods with retry and error logic --- .../%service/transports/rest_asyncio.py.j2 | 28 +++- .../cloud_redis/transports/rest_asyncio.py | 156 +++++++++++++++++- 2 files changed, 180 insertions(+), 4 deletions(-) diff --git a/gapic/templates/%namespace/%name_%version/%sub/services/%service/transports/rest_asyncio.py.j2 b/gapic/templates/%namespace/%name_%version/%sub/services/%service/transports/rest_asyncio.py.j2 index 5445fad853..5e88705739 100644 --- a/gapic/templates/%namespace/%name_%version/%sub/services/%service/transports/rest_asyncio.py.j2 +++ b/gapic/templates/%namespace/%name_%version/%sub/services/%service/transports/rest_asyncio.py.j2 @@ -1,3 +1,4 @@ +{% import "%namespace/%name_%version/%sub/services/%service/_shared_macros.j2" as shared_macros %} {# TODO(https://github.com/googleapis/gapic-generator-python/issues/2121): Remove the following variable (and the condition later in this file) for async rest transport once support for it is GA. #} {% set rest_async_io_enabled = api.all_library_settings[api.naming.proto_package].python_settings.experimental_features.rest_async_io_enabled %} {% extends '_base.py.j2' %} @@ -5,16 +6,21 @@ {% block content %} try: - from google.auth.aio.transport.sessions import AsyncAuthorizedSession # type: ignore + from google.auth.aio.transport.sessions import AsyncAuthorizedSession # type: ignore except ImportError as e: # pragma: NO COVER {# TODO(https://github.com/googleapis/google-auth-library-python/pull/1577): Update the version of google-auth once the linked PR is merged. #} raise ImportError("async rest transport requires google.auth >= 2.x.x") from e from google.auth.aio import credentials as ga_credentials_async # type: ignore +from google.api_core import exceptions as core_exceptions from google.api_core import gapic_v1 +from google.api_core import retry_async as retries -from typing import Any, Optional + +from typing import Any, Callable, Tuple, Optional, Sequence, Union + +{{ shared_macros.operations_mixin_imports(api, service, opts) }} from .rest_base import _Base{{ service.name }}RestTransport @@ -79,6 +85,24 @@ class Async{{service.name}}RestTransport(_Base{{ service.name }}RestTransport): api_audience=None ) self._session = AsyncAuthorizedSession(self._credentials) + self._wrap_with_kind = True + self._prep_wrapped_messages(client_info) + + {{ shared_macros.prep_wrapped_messages_async_method(service)|indent(4) }} + + {{ shared_macros.wrap_async_method_macro()|indent(4) }} + + {% for method in service.methods.values()|sort(attribute="name") %} + + {# TODO(b/362949446) Return a callable once implemented and remove `type: ignore`. #} + @property + def {{method.transport_safe_name|snake_case}}(self) -> Callable[ + [{{method.input.ident}}], + {{method.output.ident}}]: + + return # type: ignore + + {% endfor %} @property def kind(self) -> str: diff --git a/tests/integration/goldens/redis/google/cloud/redis_v1/services/cloud_redis/transports/rest_asyncio.py b/tests/integration/goldens/redis/google/cloud/redis_v1/services/cloud_redis/transports/rest_asyncio.py index 2e5003dedd..fc8df893fa 100755 --- a/tests/integration/goldens/redis/google/cloud/redis_v1/services/cloud_redis/transports/rest_asyncio.py +++ b/tests/integration/goldens/redis/google/cloud/redis_v1/services/cloud_redis/transports/rest_asyncio.py @@ -14,15 +14,23 @@ # limitations under the License. # try: - from google.auth.aio.transport.sessions import AsyncAuthorizedSession # type: ignore + from google.auth.aio.transport.sessions import AsyncAuthorizedSession # type: ignore except ImportError as e: # pragma: NO COVER raise ImportError("async rest transport requires google.auth >= 2.x.x") from e from google.auth.aio import credentials as ga_credentials_async # type: ignore +from google.api_core import exceptions as core_exceptions from google.api_core import gapic_v1 +from google.api_core import retry_async as retries + + +from typing import Any, Callable, Tuple, Optional, Sequence, Union + + +from google.cloud.redis_v1.types import cloud_redis +from google.longrunning import operations_pb2 # type: ignore -from typing import Any, Optional from .rest_base import _BaseCloudRedisRestTransport @@ -104,6 +112,150 @@ def __init__(self, *, api_audience=None ) self._session = AsyncAuthorizedSession(self._credentials) + self._wrap_with_kind = True + self._prep_wrapped_messages(client_info) + + def _prep_wrapped_messages(self, client_info): + """ Precompute the wrapped methods, overriding the base class method to use async wrappers.""" + self._wrapped_methods = { + self.list_instances: self._wrap_method( + self.list_instances, + default_timeout=600.0, + client_info=client_info, + ), + self.get_instance: self._wrap_method( + self.get_instance, + default_timeout=600.0, + client_info=client_info, + ), + self.get_instance_auth_string: self._wrap_method( + self.get_instance_auth_string, + default_timeout=600.0, + client_info=client_info, + ), + self.create_instance: self._wrap_method( + self.create_instance, + default_timeout=600.0, + client_info=client_info, + ), + self.update_instance: self._wrap_method( + self.update_instance, + default_timeout=600.0, + client_info=client_info, + ), + self.upgrade_instance: self._wrap_method( + self.upgrade_instance, + default_timeout=600.0, + client_info=client_info, + ), + self.import_instance: self._wrap_method( + self.import_instance, + default_timeout=600.0, + client_info=client_info, + ), + self.export_instance: self._wrap_method( + self.export_instance, + default_timeout=600.0, + client_info=client_info, + ), + self.failover_instance: self._wrap_method( + self.failover_instance, + default_timeout=600.0, + client_info=client_info, + ), + self.delete_instance: self._wrap_method( + self.delete_instance, + default_timeout=600.0, + client_info=client_info, + ), + self.reschedule_maintenance: self._wrap_method( + self.reschedule_maintenance, + default_timeout=None, + client_info=client_info, + ), + } + + def _wrap_method(self, func, *args, **kwargs): + if self._wrap_with_kind: # pragma: NO COVER + kwargs["kind"] = self.kind + return gapic_v1.method_async.wrap_method(func, *args, **kwargs) + + @property + def create_instance(self) -> Callable[ + [cloud_redis.CreateInstanceRequest], + operations_pb2.Operation]: + + return # type: ignore + + @property + def delete_instance(self) -> Callable[ + [cloud_redis.DeleteInstanceRequest], + operations_pb2.Operation]: + + return # type: ignore + + @property + def export_instance(self) -> Callable[ + [cloud_redis.ExportInstanceRequest], + operations_pb2.Operation]: + + return # type: ignore + + @property + def failover_instance(self) -> Callable[ + [cloud_redis.FailoverInstanceRequest], + operations_pb2.Operation]: + + return # type: ignore + + @property + def get_instance(self) -> Callable[ + [cloud_redis.GetInstanceRequest], + cloud_redis.Instance]: + + return # type: ignore + + @property + def get_instance_auth_string(self) -> Callable[ + [cloud_redis.GetInstanceAuthStringRequest], + cloud_redis.InstanceAuthString]: + + return # type: ignore + + @property + def import_instance(self) -> Callable[ + [cloud_redis.ImportInstanceRequest], + operations_pb2.Operation]: + + return # type: ignore + + @property + def list_instances(self) -> Callable[ + [cloud_redis.ListInstancesRequest], + cloud_redis.ListInstancesResponse]: + + return # type: ignore + + @property + def reschedule_maintenance(self) -> Callable[ + [cloud_redis.RescheduleMaintenanceRequest], + operations_pb2.Operation]: + + return # type: ignore + + @property + def update_instance(self) -> Callable[ + [cloud_redis.UpdateInstanceRequest], + operations_pb2.Operation]: + + return # type: ignore + + @property + def upgrade_instance(self) -> Callable[ + [cloud_redis.UpgradeInstanceRequest], + operations_pb2.Operation]: + + return # type: ignore @property def kind(self) -> str: From 23cb93ed58022d2d1f2b777def278340eebde555 Mon Sep 17 00:00:00 2001 From: ohmayr Date: Thu, 12 Sep 2024 18:20:10 +0000 Subject: [PATCH 2/3] address PR comments --- .../%sub/services/%service/transports/rest_asyncio.py.j2 | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/gapic/templates/%namespace/%name_%version/%sub/services/%service/transports/rest_asyncio.py.j2 b/gapic/templates/%namespace/%name_%version/%sub/services/%service/transports/rest_asyncio.py.j2 index 5e88705739..a66f487f83 100644 --- a/gapic/templates/%namespace/%name_%version/%sub/services/%service/transports/rest_asyncio.py.j2 +++ b/gapic/templates/%namespace/%name_%version/%sub/services/%service/transports/rest_asyncio.py.j2 @@ -94,7 +94,8 @@ class Async{{service.name}}RestTransport(_Base{{ service.name }}RestTransport): {% for method in service.methods.values()|sort(attribute="name") %} - {# TODO(b/362949446) Return a callable once implemented and remove `type: ignore`. #} + {# TODO(b/362949446) Return a callable once the class is implemented. #} + {# TODO(https://github.com/googleapis/gapic-generator-python/issues/2154): Remove `type: ignore`. #} @property def {{method.transport_safe_name|snake_case}}(self) -> Callable[ [{{method.input.ident}}], From 354fe4c94125920c840f121b393b9a985d1afd12 Mon Sep 17 00:00:00 2001 From: ohmayr Date: Thu, 12 Sep 2024 22:17:49 +0000 Subject: [PATCH 3/3] address more PR comments --- .../services/%service/transports/rest_asyncio.py.j2 | 2 -- .../services/cloud_redis/transports/rest_asyncio.py | 12 ------------ 2 files changed, 14 deletions(-) diff --git a/gapic/templates/%namespace/%name_%version/%sub/services/%service/transports/rest_asyncio.py.j2 b/gapic/templates/%namespace/%name_%version/%sub/services/%service/transports/rest_asyncio.py.j2 index a66f487f83..39e2f24141 100644 --- a/gapic/templates/%namespace/%name_%version/%sub/services/%service/transports/rest_asyncio.py.j2 +++ b/gapic/templates/%namespace/%name_%version/%sub/services/%service/transports/rest_asyncio.py.j2 @@ -17,7 +17,6 @@ from google.api_core import exceptions as core_exceptions from google.api_core import gapic_v1 from google.api_core import retry_async as retries - from typing import Any, Callable, Tuple, Optional, Sequence, Union {{ shared_macros.operations_mixin_imports(api, service, opts) }} @@ -100,7 +99,6 @@ class Async{{service.name}}RestTransport(_Base{{ service.name }}RestTransport): def {{method.transport_safe_name|snake_case}}(self) -> Callable[ [{{method.input.ident}}], {{method.output.ident}}]: - return # type: ignore {% endfor %} diff --git a/tests/integration/goldens/redis/google/cloud/redis_v1/services/cloud_redis/transports/rest_asyncio.py b/tests/integration/goldens/redis/google/cloud/redis_v1/services/cloud_redis/transports/rest_asyncio.py index fc8df893fa..6f102a418b 100755 --- a/tests/integration/goldens/redis/google/cloud/redis_v1/services/cloud_redis/transports/rest_asyncio.py +++ b/tests/integration/goldens/redis/google/cloud/redis_v1/services/cloud_redis/transports/rest_asyncio.py @@ -24,7 +24,6 @@ from google.api_core import gapic_v1 from google.api_core import retry_async as retries - from typing import Any, Callable, Tuple, Optional, Sequence, Union @@ -184,77 +183,66 @@ def _wrap_method(self, func, *args, **kwargs): def create_instance(self) -> Callable[ [cloud_redis.CreateInstanceRequest], operations_pb2.Operation]: - return # type: ignore @property def delete_instance(self) -> Callable[ [cloud_redis.DeleteInstanceRequest], operations_pb2.Operation]: - return # type: ignore @property def export_instance(self) -> Callable[ [cloud_redis.ExportInstanceRequest], operations_pb2.Operation]: - return # type: ignore @property def failover_instance(self) -> Callable[ [cloud_redis.FailoverInstanceRequest], operations_pb2.Operation]: - return # type: ignore @property def get_instance(self) -> Callable[ [cloud_redis.GetInstanceRequest], cloud_redis.Instance]: - return # type: ignore @property def get_instance_auth_string(self) -> Callable[ [cloud_redis.GetInstanceAuthStringRequest], cloud_redis.InstanceAuthString]: - return # type: ignore @property def import_instance(self) -> Callable[ [cloud_redis.ImportInstanceRequest], operations_pb2.Operation]: - return # type: ignore @property def list_instances(self) -> Callable[ [cloud_redis.ListInstancesRequest], cloud_redis.ListInstancesResponse]: - return # type: ignore @property def reschedule_maintenance(self) -> Callable[ [cloud_redis.RescheduleMaintenanceRequest], operations_pb2.Operation]: - return # type: ignore @property def update_instance(self) -> Callable[ [cloud_redis.UpdateInstanceRequest], operations_pb2.Operation]: - return # type: ignore @property def upgrade_instance(self) -> Callable[ [cloud_redis.UpgradeInstanceRequest], operations_pb2.Operation]: - return # type: ignore @property