feat: add support for async rest streaming methods

Author: ohmayr

gapic/templates/%namespace/%name_%version/%sub/services/%service/transports/rest_asyncio.py.j2

@@ -17,6 +17,8 @@ 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 google.api_core import rest_helpers
+{# TODO(https://github.com/googleapis/gapic-generator-python/issues/2137): raise an import error if an older version of google.api.core is installed. #}
+from google.api_core import rest_streaming_async  # type: ignore
 
 from google.protobuf import json_format
 
@@ -112,19 +114,19 @@ class Async{{service.name}}RestTransport(_Base{{ service.name }}RestTransport):
         def __hash__(self):
             return hash("Async{{service.name}}RestTransport.{{method.name}}")
 
-        {% if method.http_options and not method.client_streaming and not method.server_streaming and not method.lro and not method.extended_lro and not method.paged_result_field %}
+        {% if method.http_options and not method.client_streaming and not method.lro and not method.extended_lro and not method.paged_result_field %}
         {% set body_spec = method.http_options[0].body %}
         {{ shared_macros.response_method(body_spec, is_async=True)|indent(8) }}
 
-        {% endif %}{#  method.http_options and not method.client_streaming and not method.server_streaming and not method.lro and not method.extended_lro and not method.paged_result_field #}
+        {% endif %}{#  method.http_options and not method.client_streaming and not method.lro and not method.extended_lro and not method.paged_result_field #}
         async def __call__(self,
                     request: {{method.input.ident}}, *,
                     retry: OptionalRetry=gapic_v1.method.DEFAULT,
                     timeout: Optional[float]=None,
                     metadata: Sequence[Tuple[str, str]]=(),
                     {# TODO(b/362949446): Update the return type as we implement this for different method types. #}
-                    ){% if not method.void %} -> {% if not method.server_streaming %}{{method.output.ident}}{% else %}None{% endif %}{% endif %}:
-         {% if method.http_options and not method.client_streaming and not method.server_streaming and not method.lro and not method.extended_lro and not method.paged_result_field %}
+                    ){% if not method.void %} -> {% if not method.server_streaming %}{{method.output.ident}}{% else %}rest_streaming_async.AsyncResponseIterator{% endif %}{% endif %}:
+         {% if method.http_options and not method.client_streaming and not method.lro and not method.extended_lro and not method.paged_result_field %}
             r"""Call the {{- ' ' -}}
             {{ (method.name|snake_case).replace('_',' ')|wrap(
                     width=70, offset=45, indent=8) }}
@@ -151,14 +153,18 @@ class Async{{service.name}}RestTransport(_Base{{ service.name }}RestTransport):
 
             {% if not method.void %}
             # Return the response
+            {% if method.server_streaming %}
+            resp = rest_streaming_async.AsyncResponseIterator(response, {{method.output.ident}})
+            {% else %}
             resp = {{method.output.ident}}()
             {% if method.output.ident.is_proto_plus_type %}
             pb_resp = {{method.output.ident}}.pb(resp)
             {% else %}
             pb_resp = resp
-            {% endif %}
+            {% endif %}{# if method.output.ident.is_proto_plus_type #}
             content = await response.read()
             json_format.Parse(content, pb_resp, ignore_unknown_fields=True)
+            {% endif %}{# if method.server_streaming #}
             return resp
 
             {% endif %}{# method.void #}
@@ -167,7 +173,7 @@ class Async{{service.name}}RestTransport(_Base{{ service.name }}RestTransport):
                 raise NotImplementedError(
                     "Method {{ method.name }} is not available over REST transport"
                 )
-        {% endif %}{#  method.http_options and not method.client_streaming and not method.server_streaming and not method.lro and not method.extended_lro and not method.paged_result_field #}
+        {% endif %}{#  method.http_options and not method.client_streaming and not method.lro and not method.extended_lro and not method.paged_result_field #}
 
     {% endfor %}
     {% for method in service.methods.values()|sort(attribute="name") %}

gapic/templates/tests/unit/gapic/%name_%version/%sub/test_%service.py.j2

@@ -20,7 +20,7 @@ except ImportError:  # pragma: NO COVER
 import grpc
 from grpc.experimental import aio
 {% if "rest" in opts.transport %}
-from collections.abc import Iterable
+from collections.abc import Iterable, AsyncIterable
 from google.protobuf import json_format
 import json
 {% endif %}
@@ -95,6 +95,11 @@ from google.iam.v1 import policy_pb2  # type: ignore
 {% endfilter %}
 {{ shared_macros.add_google_api_core_version_header_import(service.version) }}
 
+async def mock_async_gen(data, chunk_size=1):
+    for i in range(0, len(data)):  # pragma: NO COVER
+        chunk = data[i : i + chunk_size]
+        yield chunk.encode("utf-8")
+
 def client_cert_source_callback():
     return b"cert bytes", b"key bytes"
 

gapic/templates/tests/unit/gapic/%name_%version/%sub/test_macros.j2

@@ -1909,13 +1909,13 @@ def test_unsupported_parameter_rest_asyncio():
 {% endmacro %}
 
 {# is_rest_unsupported_method renders:
- # 'True' if transport is async REST.
- # 'True' if transport is sync REST and method is a client streaming method.
+ # 'True' if transport is async REST and method is one of [client_streaming, lro, extended_lro, paged_result_field].
+ # 'True' if transport is sync REST and method is a client_streaming method.
  # 'False' otherwise.
 #}
 {# NOTE: We will keep updating this method as we add support for methods in async REST. #}
 {% macro is_rest_unsupported_method(method, is_async) %}
-{%- if method.client_streaming or (is_async and (method.server_streaming or method.lro or method.extended_lro or method.paged_result_field)) -%}
+{%- if method.client_streaming or (is_async and (method.lro or method.extended_lro or method.paged_result_field)) -%}
 {{'True'}}
 {%- else -%}
 {{'False'}}
@@ -2038,7 +2038,7 @@ def test_{{transport_name}}_initialize_client():
 {# rest_method_call_success_test generates tests for rest methods
  # when they make a successful request.
  # NOTE: Currently, this macro does not support the following method
- # types: [method.server_streaming, method.lro, method.extended_lro, method.paged_result_field].
+ # types: [method.lro, method.extended_lro, method.paged_result_field].
  # As support is added for the above methods, the relevant guard can be removed from within the macro
  # TODO(https://github.com/googleapis/gapic-generator-python/issues/2142): Clean up `rest_required_tests` once support for all the methods metioned above is added here.
 #}
@@ -2051,7 +2051,7 @@ def test_{{transport_name}}_initialize_client():
 {# TODO(https://github.com/googleapis/gapic-generator-python/issues/2143): Update the guard below as support we add for each method. Remove it once we have
  # all the methods supported in async rest transport that are supported in sync rest transport.
  #}
-{% if not (method.server_streaming or method.lro or method.extended_lro or method.paged_result_field)%}
+{% if not (method.lro or method.extended_lro or method.paged_result_field)%}
 @pytest.mark.parametrize("request_type", [
   {{ method.input.ident }},
   dict,
@@ -2107,13 +2107,34 @@ def test_{{transport_name}}_initialize_client():
         {% endif %}{# method.output.ident.is_proto_plus_type #}
         json_return_value = json_format.MessageToJson(return_value)
         {% endif %}{# method.void #}
+
+        {% if method.server_streaming %}
+        json_return_value = "[{}]".format(json_return_value)
+        {% if is_async %}
+        response_value.content.return_value = mock_async_gen(json_return_value)
+        {% else %}{# not is_async #}
+        response_value.iter_content = mock.Mock(return_value=iter(json_return_value))
+        {% endif %}{# is_async #}
+        {% else %}{# not method.streaming #}
         {% if is_async %}
         response_value.read = mock.AsyncMock(return_value=json_return_value.encode('UTF-8'))
-        {% else %}{# is_async #}
+        {% else %}{# not is_async #}
         response_value.content = json_return_value.encode('UTF-8')
         {% endif %}{# is_async #}
+        {% endif %}{# method.server_streaming #}
         req.return_value = response_value
         response = {{ await_prefix }}client.{{ method_name }}(request)
+    
+    {% if method.server_streaming %}
+    {% if is_async %}
+    assert isinstance(response, AsyncIterable)
+    response = await response.__anext__()
+    {% else %}
+    assert isinstance(response, Iterable)
+    response = next(response)
+    {% endif %}
+    {% endif %}
+    
     # Establish that the response is the type that we expect.
     {% if method.void %}
     assert response is None

noxfile.py

@@ -442,7 +442,7 @@ def showcase_unit_w_rest_async(
         session.chdir(lib)
         # Note: google-api-core and google-auth are re-installed here to override the version installed in constraints.
         # TODO(https://github.com/googleapis/python-api-core/pull/694): Update the version of google-api-core once the linked PR is merged.
-        session.install('--no-cache-dir', '--force-reinstall', "google-api-core[grpc]@git+https://github.com/googleapis/python-api-core.git@7dea20d73878eca93b61bb82ae6ddf335fb3a8ca")
+        session.install('--no-cache-dir', '--force-reinstall', "google-api-core[grpc]@git+https://github.com/googleapis/python-api-core.git@16038182329055551a32acd0f9f505301be4bcc5")
         # TODO(https://github.com/googleapis/google-auth-library-python/pull/1577): Update the version of google-auth once the linked PR is merged.
         session.install('--no-cache-dir', '--force-reinstall', "google-auth@git+https://github.com/googleapis/google-auth-library-python.git@add-support-for-async-authorized-session-api")
         session.install("aiohttp")

tests/integration/goldens/asset/tests/unit/gapic/asset_v1/test_asset_service.py

@@ -23,7 +23,7 @@
 
 import grpc
 from grpc.experimental import aio
-from collections.abc import Iterable
+from collections.abc import Iterable, AsyncIterable
 from google.protobuf import json_format
 import json
 import math
@@ -71,6 +71,11 @@
 import google.auth
 
 
+async def mock_async_gen(data, chunk_size=1):
+    for i in range(0, len(data)):  # pragma: NO COVER
+        chunk = data[i : i + chunk_size]
+        yield chunk.encode("utf-8")
+
 def client_cert_source_callback():
     return b"cert bytes", b"key bytes"
 
@@ -15286,9 +15291,11 @@ def test_batch_get_assets_history_rest_call_success(request_type):
         # Convert return value to protobuf type
         return_value = asset_service.BatchGetAssetsHistoryResponse.pb(return_value)
         json_return_value = json_format.MessageToJson(return_value)
+
         response_value.content = json_return_value.encode('UTF-8')
         req.return_value = response_value
         response = client.batch_get_assets_history(request)
+
     # Establish that the response is the type that we expect.
     assert isinstance(response, asset_service.BatchGetAssetsHistoryResponse)
 
@@ -15346,9 +15353,11 @@ def test_create_feed_rest_call_success(request_type):
         # Convert return value to protobuf type
         return_value = asset_service.Feed.pb(return_value)
         json_return_value = json_format.MessageToJson(return_value)
+
         response_value.content = json_return_value.encode('UTF-8')
         req.return_value = response_value
         response = client.create_feed(request)
+
     # Establish that the response is the type that we expect.
     assert isinstance(response, asset_service.Feed)
     assert response.name == 'name_value'
@@ -15415,9 +15424,11 @@ def test_get_feed_rest_call_success(request_type):
         # Convert return value to protobuf type
         return_value = asset_service.Feed.pb(return_value)
         json_return_value = json_format.MessageToJson(return_value)
+
         response_value.content = json_return_value.encode('UTF-8')
         req.return_value = response_value
         response = client.get_feed(request)
+
     # Establish that the response is the type that we expect.
     assert isinstance(response, asset_service.Feed)
     assert response.name == 'name_value'
@@ -15479,9 +15490,11 @@ def test_list_feeds_rest_call_success(request_type):
         # Convert return value to protobuf type
         return_value = asset_service.ListFeedsResponse.pb(return_value)
         json_return_value = json_format.MessageToJson(return_value)
+
         response_value.content = json_return_value.encode('UTF-8')
         req.return_value = response_value
         response = client.list_feeds(request)
+
     # Establish that the response is the type that we expect.
     assert isinstance(response, asset_service.ListFeedsResponse)
 
@@ -15539,9 +15552,11 @@ def test_update_feed_rest_call_success(request_type):
         # Convert return value to protobuf type
         return_value = asset_service.Feed.pb(return_value)
         json_return_value = json_format.MessageToJson(return_value)
+
         response_value.content = json_return_value.encode('UTF-8')
         req.return_value = response_value
         response = client.update_feed(request)
+
     # Establish that the response is the type that we expect.
     assert isinstance(response, asset_service.Feed)
     assert response.name == 'name_value'
@@ -15600,9 +15615,11 @@ def test_delete_feed_rest_call_success(request_type):
         response_value = mock.Mock()
         response_value.status_code = 200
         json_return_value = ''
+
         response_value.content = json_return_value.encode('UTF-8')
         req.return_value = response_value
         response = client.delete_feed(request)
+
     # Establish that the response is the type that we expect.
     assert response is None
 
@@ -15698,9 +15715,11 @@ def test_analyze_iam_policy_rest_call_success(request_type):
         # Convert return value to protobuf type
         return_value = asset_service.AnalyzeIamPolicyResponse.pb(return_value)
         json_return_value = json_format.MessageToJson(return_value)
+
         response_value.content = json_return_value.encode('UTF-8')
         req.return_value = response_value
         response = client.analyze_iam_policy(request)
+
     # Establish that the response is the type that we expect.
     assert isinstance(response, asset_service.AnalyzeIamPolicyResponse)
     assert response.fully_explored is True
@@ -15775,9 +15794,11 @@ def test_analyze_move_rest_call_success(request_type):
         # Convert return value to protobuf type
         return_value = asset_service.AnalyzeMoveResponse.pb(return_value)
         json_return_value = json_format.MessageToJson(return_value)
+
         response_value.content = json_return_value.encode('UTF-8')
         req.return_value = response_value
         response = client.analyze_move(request)
+
     # Establish that the response is the type that we expect.
     assert isinstance(response, asset_service.AnalyzeMoveResponse)
 
@@ -15832,9 +15853,11 @@ def test_query_assets_rest_call_success(request_type):
         # Convert return value to protobuf type
         return_value = asset_service.QueryAssetsResponse.pb(return_value)
         json_return_value = json_format.MessageToJson(return_value)
+
         response_value.content = json_return_value.encode('UTF-8')
         req.return_value = response_value
         response = client.query_assets(request)
+
     # Establish that the response is the type that we expect.
     assert isinstance(response, asset_service.QueryAssetsResponse)
     assert response.job_reference == 'job_reference_value'
@@ -15894,9 +15917,11 @@ def test_create_saved_query_rest_call_success(request_type):
         # Convert return value to protobuf type
         return_value = asset_service.SavedQuery.pb(return_value)
         json_return_value = json_format.MessageToJson(return_value)
+
         response_value.content = json_return_value.encode('UTF-8')
         req.return_value = response_value
         response = client.create_saved_query(request)
+
     # Establish that the response is the type that we expect.
     assert isinstance(response, asset_service.SavedQuery)
     assert response.name == 'name_value'
@@ -15960,9 +15985,11 @@ def test_get_saved_query_rest_call_success(request_type):
         # Convert return value to protobuf type
         return_value = asset_service.SavedQuery.pb(return_value)
         json_return_value = json_format.MessageToJson(return_value)
+
         response_value.content = json_return_value.encode('UTF-8')
         req.return_value = response_value
         response = client.get_saved_query(request)
+
     # Establish that the response is the type that we expect.
     assert isinstance(response, asset_service.SavedQuery)
     assert response.name == 'name_value'
@@ -16047,9 +16074,11 @@ def test_update_saved_query_rest_call_success(request_type):
         # Convert return value to protobuf type
         return_value = asset_service.SavedQuery.pb(return_value)
         json_return_value = json_format.MessageToJson(return_value)
+
         response_value.content = json_return_value.encode('UTF-8')
         req.return_value = response_value
         response = client.update_saved_query(request)
+
     # Establish that the response is the type that we expect.
     assert isinstance(response, asset_service.SavedQuery)
     assert response.name == 'name_value'
@@ -16106,9 +16135,11 @@ def test_delete_saved_query_rest_call_success(request_type):
         response_value = mock.Mock()
         response_value.status_code = 200
         json_return_value = ''
+
         response_value.content = json_return_value.encode('UTF-8')
         req.return_value = response_value
         response = client.delete_saved_query(request)
+
     # Establish that the response is the type that we expect.
     assert response is None
 
@@ -16161,9 +16192,11 @@ def test_batch_get_effective_iam_policies_rest_call_success(request_type):
         # Convert return value to protobuf type
         return_value = asset_service.BatchGetEffectiveIamPoliciesResponse.pb(return_value)
         json_return_value = json_format.MessageToJson(return_value)
+
         response_value.content = json_return_value.encode('UTF-8')
         req.return_value = response_value
         response = client.batch_get_effective_iam_policies(request)
+
     # Establish that the response is the type that we expect.
     assert isinstance(response, asset_service.BatchGetEffectiveIamPoliciesResponse)