feat: implement async rest mixin methods (#2181)

googleapis · Sep 25, 2024 · ed04b26 · ed04b26
1 parent bedbab8
commit ed04b26
Show file tree

Hide file tree

Showing 26 changed files with 1,247 additions and 202 deletions.
diff --git a/gapic/ads-templates/%namespace/%name/%version/%sub/services/%service/_shared_macros.j2 b/gapic/ads-templates/%namespace/%name/%version/%sub/services/%service/_shared_macros.j2
@@ -235,7 +235,7 @@ def _get_http_options():
 {% endmacro %}
 
 
-{% macro prep_wrapped_messages_async_method(service) %}
+{% macro prep_wrapped_messages_async_method(api, service) %}
 def _prep_wrapped_messages(self, client_info):
     """ Precompute the wrapped methods, overriding the base class method to use async wrappers."""
     self._wrapped_methods = {
@@ -265,6 +265,13 @@ def _prep_wrapped_messages(self, client_info):
             client_info=client_info,
         ),
         {% endfor %}{# service.methods.values() #}
+        {% for method_name in api.mixin_api_methods.keys() %}
+        self.{{ method_name|snake_case }}: self._wrap_method(
+            self.{{ method_name|snake_case }},
+            default_timeout=None,
+            client_info=client_info,
+        ),
+        {% endfor %} {# method_name in api.mixin_api_methods.keys() #}
     }
 {% endmacro %}
 

diff --git a/gapic/templates/%namespace/%name_%version/%sub/services/%service/_async_mixins.py.j2 b/gapic/templates/%namespace/%name_%version/%sub/services/%service/_async_mixins.py.j2
@@ -34,11 +34,7 @@
 
         # Wrap the RPC method; this adds retry and timeout information,
         # and friendly error handling.
-        rpc = gapic_v1.method_async.wrap_method(
-            self._client._transport.list_operations,
-            default_timeout=None,
-            client_info=DEFAULT_CLIENT_INFO,
-        )
+        rpc = self.transport._wrapped_methods[self._client._transport.list_operations]
 
         # Certain fields should be provided within the metadata header;
         # add these here.
@@ -91,11 +87,7 @@
 
         # Wrap the RPC method; this adds retry and timeout information,
         # and friendly error handling.
-        rpc = gapic_v1.method_async.wrap_method(
-            self._client._transport.get_operation,
-            default_timeout=None,
-            client_info=DEFAULT_CLIENT_INFO,
-        )
+        rpc = self.transport._wrapped_methods[self._client._transport.get_operation]
 
         # Certain fields should be provided within the metadata header;
         # add these here.
@@ -151,11 +143,7 @@
 
         # Wrap the RPC method; this adds retry and timeout information,
         # and friendly error handling.
-        rpc = gapic_v1.method_async.wrap_method(
-            self._client._transport.delete_operation,
-            default_timeout=None,
-            client_info=DEFAULT_CLIENT_INFO,
-        )
+        rpc = self.transport._wrapped_methods[self._client._transport.delete_operation]
 
         # Certain fields should be provided within the metadata header;
         # add these here.
@@ -206,11 +194,7 @@
 
         # Wrap the RPC method; this adds retry and timeout information,
         # and friendly error handling.
-        rpc = gapic_v1.method_async.wrap_method(
-            self._client._transport.cancel_operation,
-            default_timeout=None,
-            client_info=DEFAULT_CLIENT_INFO,
-        )
+        rpc = self.transport._wrapped_methods[self._client._transport.cancel_operation]
 
         # Certain fields should be provided within the metadata header;
         # add these here.
@@ -264,11 +248,7 @@
 
         # Wrap the RPC method; this adds retry and timeout information,
         # and friendly error handling.
-        rpc = gapic_v1.method_async.wrap_method(
-            self._client._transport.wait_operation,
-            default_timeout=None,
-            client_info=DEFAULT_CLIENT_INFO,
-        )
+        rpc = self.transport._wrapped_methods[self._client._transport.wait_operation]
 
         # Certain fields should be provided within the metadata header;
         # add these here.
@@ -391,11 +371,7 @@
 
         # Wrap the RPC method; this adds retry and timeout information,
         # and friendly error handling.
-        rpc = gapic_v1.method_async.wrap_method(
-            self._client._transport.set_iam_policy,
-            default_timeout=None,
-            client_info=DEFAULT_CLIENT_INFO,
-        )
+        rpc = self.transport._wrapped_methods[self._client._transport.set_iam_policy]
 
         # Certain fields should be provided within the metadata header;
         # add these here.
@@ -514,11 +490,7 @@
 
         # Wrap the RPC method; this adds retry and timeout information,
         # and friendly error handling.
-        rpc = gapic_v1.method_async.wrap_method(
-            self._client._transport.get_iam_policy,
-            default_timeout=None,
-            client_info=DEFAULT_CLIENT_INFO,
-        )
+        rpc = self.transport._wrapped_methods[self._client._transport.get_iam_policy]
 
         # Certain fields should be provided within the metadata header;
         # add these here.
@@ -575,11 +547,7 @@
 
         # Wrap the RPC method; this adds retry and timeout information,
         # and friendly error handling.
-        rpc = gapic_v1.method_async.wrap_method(
-            self._client._transport.test_iam_permissions,
-            default_timeout=None,
-            client_info=DEFAULT_CLIENT_INFO,
-        )
+        rpc = self.transport._wrapped_methods[self._client._transport.test_iam_permissions]
 
         # Certain fields should be provided within the metadata header;
         # add these here.
@@ -635,11 +603,7 @@
 
         # Wrap the RPC method; this adds retry and timeout information,
         # and friendly error handling.
-        rpc = gapic_v1.method_async.wrap_method(
-            self._client._transport.get_location,
-            default_timeout=None,
-            client_info=DEFAULT_CLIENT_INFO,
-        )
+        rpc = self.transport._wrapped_methods[self._client._transport.get_location]
 
         # Certain fields should be provided within the metadata header;
         # add these here.
@@ -691,11 +655,7 @@
 
         # Wrap the RPC method; this adds retry and timeout information,
         # and friendly error handling.
-        rpc = gapic_v1.method_async.wrap_method(
-            self._client._transport.list_locations,
-            default_timeout=None,
-            client_info=DEFAULT_CLIENT_INFO,
-        )
+        rpc = self.transport._wrapped_methods[self._client._transport.list_locations]
 
         # Certain fields should be provided within the metadata header;
         # add these here.

diff --git a/gapic/templates/%namespace/%name_%version/%sub/services/%service/_shared_macros.j2 b/gapic/templates/%namespace/%name_%version/%sub/services/%service/_shared_macros.j2
@@ -228,7 +228,7 @@ def _get_http_options():
 {% endmacro %}
 
 
-{% macro prep_wrapped_messages_async_method(service) %}
+{% macro prep_wrapped_messages_async_method(api, service) %}
 def _prep_wrapped_messages(self, client_info):
     """ Precompute the wrapped methods, overriding the base class method to use async wrappers."""
     self._wrapped_methods = {
@@ -258,6 +258,16 @@ def _prep_wrapped_messages(self, client_info):
             client_info=client_info,
         ),
         {% endfor %}{# service.methods.values() #}
+        {% for method_name in api.mixin_api_methods.keys() %}
+        {# TODO(https://github.com/googleapis/gapic-generator-python/issues/2197): Use `transport_safe_name` similar
+         # to what we do for non-mixin methods above.
+        #}
+        self.{{ method_name|snake_case }}: self._wrap_method(
+            self.{{ method_name|snake_case }},
+            default_timeout=None,
+            client_info=client_info,
+        ),
+        {% endfor %}{# method_name in api.mixin_api_methods.keys() #}
     }
 {% endmacro %}
 
@@ -310,6 +320,7 @@ class {{ async_method_name_prefix }}{{ service.name }}RestInterceptor:
             {{ async_prefix }}def post_{{ method.name|snake_case }}(self, response):
                 logging.log(f"Received response: {response}")
                 return response
+
     {% endif %}
 
     {% endfor %}
@@ -342,12 +353,11 @@ class {{ async_method_name_prefix }}{{ service.name }}RestInterceptor:
         it is returned to user code.
         """
         return response
+
     {% endif %}
     {% endif %}{# if (not is_async) or (is_async and not method.lro and not method.extended_lro and not method.paged_result_field) #}
     {% endfor %}
 
-    {# TODO(https://github.com/googleapis/gapic-generator-python/issues/2148): Remove the condition below once mixins are supported for async rest transport. #}
-    {% if not is_async %}
     {% for name, signature in api.mixin_api_signatures.items() %}
     {{ async_prefix }}def pre_{{ name|snake_case }}(
         self, request: {{signature.request_type}}, metadata: Sequence[Tuple[str, str]]
@@ -369,11 +379,15 @@ class {{ async_method_name_prefix }}{{ service.name }}RestInterceptor:
         it is returned to user code.
         """
         return response
+
     {% endfor %}
-    {% endif %}
 {% endmacro %}
 
 {% macro generate_mixin_call_method(service, api, name, sig, is_async) %}
+{# TODO(https://github.com/googleapis/gapic-generator-python/issues/2198): generate _Mixin classes
+ # and @property methods into separate macros so that _Method and _Mixin classes can be defined all
+ # together and the @property methods for each can be defined after the class definitions.
+#}
 {% set async_prefix = "async " if is_async else "" %}
 {% set async_method_name_prefix = "Async" if is_async else "" %}
 {% set async_suffix = "_async" if is_async else "" %}

diff --git a/gapic/templates/%namespace/%name_%version/%sub/services/%service/async_client.py.j2 b/gapic/templates/%namespace/%name_%version/%sub/services/%service/async_client.py.j2
@@ -513,11 +513,7 @@ class {{ service.async_client_name }}:
 
         # Wrap the RPC method; this adds retry and timeout information,
         # and friendly error handling.
-        rpc = gapic_v1.method_async.wrap_method(
-            self._client._transport.set_iam_policy,
-            default_timeout=None,
-            client_info=DEFAULT_CLIENT_INFO,
-        )
+        rpc = self.transport._wrapped_methods[self._client._transport.set_iam_policy]
 
         # Certain fields should be provided within the metadata header;
         # add these here.
@@ -632,11 +628,7 @@ class {{ service.async_client_name }}:
 
         # Wrap the RPC method; this adds retry and timeout information,
         # and friendly error handling.
-        rpc = gapic_v1.method_async.wrap_method(
-            self._client._transport.get_iam_policy,
-            default_timeout=None,
-            client_info=DEFAULT_CLIENT_INFO,
-        )
+        rpc = self.transport._wrapped_methods[self._client._transport.get_iam_policy]
 
         # Certain fields should be provided within the metadata header;
         # add these here.
@@ -689,11 +681,7 @@ class {{ service.async_client_name }}:
 
         # Wrap the RPC method; this adds retry and timeout information,
         # and friendly error handling.
-        rpc = gapic_v1.method_async.wrap_method(
-            self._client._transport.test_iam_permissions,
-            default_timeout=None,
-            client_info=DEFAULT_CLIENT_INFO,
-        )
+        rpc = self.transport._wrapped_methods[self._client._transport.test_iam_permissions]
 
         # Certain fields should be provided within the metadata header;
         # add these here.

diff --git a/.../templates/%namespace/%name_%version/%sub/services/%service/transports/grpc_asyncio.py.j2 b/.../templates/%namespace/%name_%version/%sub/services/%service/transports/grpc_asyncio.py.j2
@@ -388,7 +388,7 @@ class {{ service.grpc_asyncio_transport_name }}({{ service.name }}Transport):
         return self._stubs["test_iam_permissions"]
     {% endif %}
 
-    {{ shared_macros.prep_wrapped_messages_async_method(service)|indent(4) }}
+    {{ shared_macros.prep_wrapped_messages_async_method(api, service)|indent(4) }}
 
     {{ shared_macros.wrap_async_method_macro()|indent(4) }}
 

diff --git a/.../templates/%namespace/%name_%version/%sub/services/%service/transports/rest_asyncio.py.j2 b/.../templates/%namespace/%name_%version/%sub/services/%service/transports/rest_asyncio.py.j2
@@ -135,7 +135,7 @@ class Async{{service.name}}RestTransport(_Base{{ service.name }}RestTransport):
         self._wrap_with_kind = True
         self._prep_wrapped_messages(client_info)
 
-    {{ shared_macros.prep_wrapped_messages_async_method(service)|indent(4) }}
+    {{ shared_macros.prep_wrapped_messages_async_method(api, service)|indent(4) }}
 
     {{ shared_macros.wrap_async_method_macro()|indent(4) }}
 
@@ -225,6 +225,9 @@ class Async{{service.name}}RestTransport(_Base{{ service.name }}RestTransport):
         return self._{{method.name}}(self._session, self._host, self._interceptor)  # type: ignore
 
     {% endfor %}
+    {% for name, sig in api.mixin_api_signatures.items() %} 
+    {{ shared_macros.generate_mixin_call_method(service, api, name, sig, is_async=True) | indent(4) }}
+    {% endfor %}
 
     @property
     def kind(self) -> str:

diff --git a/gapic/templates/tests/unit/gapic/%name_%version/%sub/test_macros.j2 b/gapic/templates/tests/unit/gapic/%name_%version/%sub/test_macros.j2
@@ -1884,11 +1884,8 @@ def test_unsupported_parameter_rest_asyncio():
 {% endfor %}{# for method in service.methods.values() #}
 {% for name, sig in api.mixin_api_signatures.items() %}
 {% if 'rest' in transport %}
-{# TODO: Remove the following condition as we implement async mixins #}
-{% if not is_async %}
 {{ bad_request_mixins_test(service, api, name, sig, transport, is_async) }}
 {{ call_success_mixins_test(service, api, name, sig, transport, is_async) }}
-{% endif %}{# if not is_async #}
 {% endif %}{# if 'rest' in transport #}
 {% endfor %}
 {{ initialize_client_with_transport_test(service, transport, is_async) }}

diff --git a/tests/integration/goldens/asset/google/cloud/asset_v1/services/asset_service/async_client.py b/tests/integration/goldens/asset/google/cloud/asset_v1/services/asset_service/async_client.py
@@ -3100,11 +3100,7 @@ async def get_operation(
 
         # Wrap the RPC method; this adds retry and timeout information,
         # and friendly error handling.
-        rpc = gapic_v1.method_async.wrap_method(
-            self._client._transport.get_operation,
-            default_timeout=None,
-            client_info=DEFAULT_CLIENT_INFO,
-        )
+        rpc = self.transport._wrapped_methods[self._client._transport.get_operation]
 
         # Certain fields should be provided within the metadata header;
         # add these here.

diff --git a/...ion/goldens/asset/google/cloud/asset_v1/services/asset_service/transports/grpc_asyncio.py b/...ion/goldens/asset/google/cloud/asset_v1/services/asset_service/transports/grpc_asyncio.py
@@ -1127,6 +1127,11 @@ def _prep_wrapped_messages(self, client_info):
                 default_timeout=None,
                 client_info=client_info,
             ),
+            self.get_operation: self._wrap_method(
+                self.get_operation,
+                default_timeout=None,
+                client_info=client_info,
+            ),
         }
 
     def _wrap_method(self, func, *args, **kwargs):