feat: provide AsyncIO support for generated code (#365)

gapic/schema/wrappers.py

@@ -566,6 +566,13 @@ def __getattr__(self, name):
 
     @utils.cached_property
     def client_output(self):
+        return self._client_output(enable_asyncio=False)
+
+    @utils.cached_property
+    def client_output_async(self):
+        return self._client_output(enable_asyncio=True)
+
+    def _client_output(self, enable_asyncio: bool):
         """Return the output from the client layer.
 
         This takes into account transformations made by the outer GAPIC
@@ -584,8 +591,8 @@ def client_output(self):
         if self.lro:
             return PythonType(meta=metadata.Metadata(
                 address=metadata.Address(
-                    name='Operation',
-                    module='operation',
+                    name='AsyncOperation' if enable_asyncio else 'Operation',
+                    module='operation_async' if enable_asyncio else 'operation',
                     package=('google', 'api_core'),
                     collisions=self.lro.response_type.ident.collisions,
                 ),
@@ -603,7 +610,7 @@ def client_output(self):
         if self.paged_result_field:
             return PythonType(meta=metadata.Metadata(
                 address=metadata.Address(
-                    name=f'{self.name}Pager',
+                    name=f'{self.name}AsyncPager' if enable_asyncio else f'{self.name}Pager',
                     package=self.ident.api_naming.module_namespace + (self.ident.api_naming.versioned_module_name,) + self.ident.subpackage + (
                         'services',
                         utils.to_snake_case(self.ident.parent[-1]),
@@ -744,6 +751,8 @@ def _ref_types(self, recursive: bool) -> Sequence[Union[MessageType, EnumType]]:
         if not self.void:
             answer.append(self.client_output)
             answer.extend(self.client_output.field_types)
+            answer.append(self.client_output_async)
+            answer.extend(self.client_output_async.field_types)
 
         # If this method has LRO, it is possible (albeit unlikely) that
         # the LRO messages reside in a different module.
@@ -801,6 +810,11 @@ def client_name(self) -> str:
         """Returns the name of the generated client class"""
         return self.name + "Client"
 
+    @property
+    def async_client_name(self) -> str:
+        """Returns the name of the generated AsyncIO client class"""
+        return self.name + "AsyncClient"
+
     @property
     def transport_name(self):
         return self.name + "Transport"
@@ -809,6 +823,10 @@ def transport_name(self):
     def grpc_transport_name(self):
         return self.name + "GrpcTransport"
 
+    @property
+    def grpc_asyncio_transport_name(self):
+        return self.name + "GrpcAsyncIOTransport"
+
     @property
     def has_lro(self) -> bool:
         """Return whether the service has a long-running method."""
@@ -856,7 +874,7 @@ def names(self) -> FrozenSet[str]:
         used for imports.
         """
         # Put together a set of the service and method names.
-        answer = {self.name, self.client_name}
+        answer = {self.name, self.client_name, self.async_client_name}
         answer.update(
             utils.to_snake_case(i.name) for i in self.methods.values()
         )

gapic/templates/%namespace/%name/__init__.py.j2

@@ -12,6 +12,8 @@ from {% if api.naming.module_namespace %}{{ api.naming.module_namespace|join('.'
         if service.meta.address.subpackage == api.subpackage_view -%}
 from {% if api.naming.module_namespace %}{{ api.naming.module_namespace|join('.') }}.{% endif -%}
     {{ api.naming.versioned_module_name }}.services.{{ service.name|snake_case }}.client import {{ service.client_name }}
+from {% if api.naming.module_namespace %}{{ api.naming.module_namespace|join('.') }}.{% endif -%}
+    {{ api.naming.versioned_module_name }}.services.{{ service.name|snake_case }}.async_client import {{ service.async_client_name }}
 {% endfor -%}
 
 {#  Import messages and enums from each proto.
@@ -48,6 +50,7 @@ __all__ = (
 {% for service in api.services.values()|sort(attribute='name')
    if service.meta.address.subpackage == api.subpackage_view -%}
 '{{ service.client_name }}',
+'{{ service.async_client_name }}',
 {% endfor -%}
 {% for proto in api.protos.values()|sort(attribute='module_name')
    if proto.meta.address.subpackage == api.subpackage_view -%}

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

@@ -2,8 +2,10 @@
 
 {% block content %}
 from .client import {{ service.client_name }}
+from .async_client import {{ service.async_client_name }}
 
 __all__ = (
     '{{ service.client_name }}',
+    '{{ service.async_client_name }}',
 )
 {% endblock %}

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

@@ -0,0 +1,271 @@
+{% extends '_base.py.j2' %}
+
+{% block content %}
+from collections import OrderedDict
+import functools
+import re
+from typing import Dict, {% if service.any_server_streaming %}AsyncIterable, {% endif %}{% if service.any_client_streaming %}AsyncIterator, {% endif %}Sequence, Tuple, Type, Union
+import pkg_resources
+
+import google.api_core.client_options as ClientOptions # type: ignore
+from google.api_core import exceptions                 # type: ignore
+from google.api_core import gapic_v1                   # type: ignore
+from google.api_core import retry as retries           # type: ignore
+from google.auth import credentials                    # type: ignore
+from google.oauth2 import service_account              # type: ignore
+
+{% filter sort_lines -%}
+{% for method in service.methods.values() -%}
+{% for ref_type in method.flat_ref_types -%}
+{{ ref_type.ident.python_import }}
+{% endfor -%}
+{% endfor -%}
+{% endfilter %}
+from .transports.base import {{ service.name }}Transport
+from .transports.grpc_asyncio import {{ service.grpc_asyncio_transport_name }}
+from .client import {{ service.client_name }}
+
+
+class {{ service.async_client_name }}:
+    """{{ service.meta.doc|rst(width=72, indent=4) }}"""
+
+    _client: {{ service.client_name }}
+
+    DEFAULT_ENDPOINT = {{ service.client_name }}.DEFAULT_ENDPOINT
+    DEFAULT_MTLS_ENDPOINT = {{ service.client_name }}.DEFAULT_MTLS_ENDPOINT
+
+    {% for message in service.resource_messages -%}
+    {{ message.resource_type|snake_case }}_path = staticmethod({{ service.client_name }}.{{ message.resource_type|snake_case }}_path)
+
+    {% endfor %}
+
+    from_service_account_file = {{ service.client_name }}.from_service_account_file
+    from_service_account_json = from_service_account_file
+
+    get_transport_class = functools.partial(type({{ service.client_name }}).get_transport_class, type({{ service.client_name }}))
+
+    def __init__(self, *,
+            credentials: credentials.Credentials = None,
+            transport: Union[str, {{ service.name }}Transport] = 'grpc_asyncio',
+            client_options: ClientOptions = None,
+            ) -> None:
+        """Instantiate the {{ (service.client_name|snake_case).replace('_', ' ') }}.
+
+        Args:
+            credentials (Optional[google.auth.credentials.Credentials]): The
+                authorization credentials to attach to requests. These
+                credentials identify the application to the service; if none
+                are specified, the client will attempt to ascertain the
+                credentials from the environment.
+            transport (Union[str, ~.{{ service.name }}Transport]): The
+                transport to use. If set to None, a transport is chosen
+                automatically.
+            client_options (ClientOptions): Custom options for the client. It
+                won't take effect if a ``transport`` instance is provided.
+                (1) The ``api_endpoint`` property can be used to override the
+                default endpoint provided by the client. GOOGLE_API_USE_MTLS
+                environment variable can also be used to override the endpoint:
+                "always" (always use the default mTLS endpoint), "never" (always
+                use the default regular endpoint, this is the default value for
+                the environment variable) and "auto" (auto switch to the default
+                mTLS endpoint if client SSL credentials is present). However,
+                the ``api_endpoint`` property takes precedence if provided.
+                (2) The ``client_cert_source`` property is used to provide client
+                SSL credentials for mutual TLS transport. If not provided, the
+                default SSL credentials will be used if present.
+
+        Raises:
+            google.auth.exceptions.MutualTlsChannelError: If mutual TLS transport
+                creation failed for any reason.
+        """
+        {# NOTE(lidiz) Not using kwargs since we want the docstring and types. #}
+        self._client = {{ service.client_name }}(
+            credentials=credentials,
+            transport=transport,
+            client_options=client_options,
+        )
+
+    {% for method in service.methods.values() -%}
+    {% if not method.server_streaming %}async {% endif -%}def {{ method.name|snake_case }}(self,
+            {%- if not method.client_streaming %}
+            request: {{ method.input.ident }} = None,
+            *,
+            {% for field in method.flattened_fields.values() -%}
+            {{ field.name }}: {{ field.ident }} = None,
+            {% endfor -%}
+            {%- else %}
+            requests: AsyncIterator[{{ method.input.ident }}] = None,
+            *,
+            {% endif -%}
+            retry: retries.Retry = gapic_v1.method.DEFAULT,
+            timeout: float = None,
+            metadata: Sequence[Tuple[str, str]] = (),
+            {%- if not method.server_streaming %}
+            ) -> {{ method.client_output_async.ident }}:
+            {%- else %}
+            ) -> AsyncIterable[{{ method.client_output_async.ident }}]:
+            {%- endif %}
+        r"""{{ method.meta.doc|rst(width=72, indent=8) }}
+
+        Args:
+            {%- if not method.client_streaming %}
+            request (:class:`{{ method.input.ident.sphinx }}`):
+                The request object.{{ ' ' -}}
+                {{ method.input.meta.doc|wrap(width=72, offset=36, indent=16) }}
+            {% for key, field in method.flattened_fields.items() -%}
+            {{ field.name }} (:class:`{{ field.ident.sphinx }}`):
+                {{ field.meta.doc|rst(width=72, indent=16, nl=False) }}
+                This corresponds to the ``{{ key }}`` field
+                on the ``request`` instance; if ``request`` is provided, this
+                should not be set.
+            {% endfor -%}
+            {%- else %}
+            requests (AsyncIterator[`{{ method.input.ident.sphinx }}`]):
+                The request object AsyncIterator.{{ ' ' -}}
+                {{ method.input.meta.doc|wrap(width=72, offset=36, indent=16) }}
+            {%- endif %}
+            retry (google.api_core.retry.Retry): Designation of what errors, if any,
+                should be retried.
+            timeout (float): The timeout for this request.
+            metadata (Sequence[Tuple[str, str]]): Strings which should be
+                sent along with the request as metadata.
+        {%- if not method.void %}
+
+        Returns:
+            {%- if not method.server_streaming %}
+            {{ method.client_output_async.ident.sphinx }}:
+            {%- else %}
+            AsyncIterable[{{ method.client_output_async.ident.sphinx }}]:
+            {%- endif %}
+                {{ method.client_output_async.meta.doc|rst(width=72, indent=16) }}
+        {%- endif %}
+        """
+        {%- if not method.client_streaming %}
+        # Create or coerce a protobuf request object.
+        {% if method.flattened_fields -%}
+        # Sanity check: If we got a request object, we should *not* have
+        # gotten any keyword arguments that map to the request.
+        if request is not None and any([{{ method.flattened_fields.values()|join(', ', attribute='name') }}]):
+            raise ValueError('If the `request` argument is set, then none of '
+                             'the individual field arguments should be set.')
+
+        {% endif -%}
+        {% if method.input.ident.package != method.ident.package -%} {# request lives in a different package, so there is no proto wrapper #}
+        # The request isn't a proto-plus wrapped type,
+        # so it must be constructed via keyword expansion.
+        if isinstance(request, dict):
+            request = {{ method.input.ident }}(**request)
+        {% if method.flattened_fields -%}{# Cross-package req and flattened fields #}
+        elif not request:
+            request = {{ method.input.ident }}()
+        {% endif -%}{# Cross-package req and flattened fields #}
+        {%- else %}
+        request = {{ method.input.ident }}(request)
+        {% endif %} {# different request package #}
+
+        {#- Vanilla python protobuf wrapper types cannot _set_ repeated fields #}
+        {% if method.flattened_fields -%}
+        # If we have keyword arguments corresponding to fields on the
+        # request, apply these.
+        {% endif -%}
+        {%- for key, field in method.flattened_fields.items() if not(field.repeated and method.input.ident.package != method.ident.package) %}
+        if {{ field.name }} is not None:
+            request.{{ key }} = {{ field.name }}
+        {%- endfor %}
+        {# They can be _extended_, however -#}
+        {%- for key, field in method.flattened_fields.items() if (field.repeated and method.input.ident.package != method.ident.package) %}
+        if {{ field.name }}:
+            request.{{ key }}.extend({{ field.name }})
+        {%- endfor %}
+        {%- endif %}
+
+        # Wrap the RPC method; this adds retry and timeout information,
+        # and friendly error handling.
+        rpc = gapic_v1.method_async.wrap_method(
+            self._client._transport.{{ method.name|snake_case }},
+            {%- if method.retry %}
+            default_retry=retries.Retry(
+                {% if method.retry.initial_backoff %}initial={{ method.retry.initial_backoff }},{% endif %}
+                {% if method.retry.max_backoff %}maximum={{ method.retry.max_backoff }},{% endif %}
+                {% if method.retry.backoff_multiplier %}multiplier={{ method.retry.backoff_multiplier }},{% endif %}
+                predicate=retries.if_exception_type(
+                    {%- filter sort_lines %}
+                    {%- for ex in method.retry.retryable_exceptions %}
+                    exceptions.{{ ex.__name__ }},
+                    {%- endfor %}
+                    {%- endfilter %}
+                ),
+            ),
+            {%- endif %}
+            default_timeout={{ method.timeout }},
+            client_info=_client_info,
+        )
+        {%- if method.field_headers %}
+
+        # Certain fields should be provided within the metadata header;
+        # add these here.
+        metadata = tuple(metadata) + (
+            gapic_v1.routing_header.to_grpc_metadata((
+                {%- for field_header in method.field_headers %}
+                {%- if not method.client_streaming %}
+                ('{{ field_header }}', request.{{ field_header }}),
+                {%- endif %}
+                {%- endfor %}
+            )),
+        )
+        {%- endif %}
+
+        # Send the request.
+        {% if not method.void %}response = {% endif %}
+        {%- if not method.server_streaming %}await {% endif %}rpc(
+            {%- if not method.client_streaming %}
+            request,
+            {%- else %}
+            requests,
+            {%- endif %}
+            retry=retry,
+            timeout=timeout,
+            metadata=metadata,
+        )
+        {%- if method.lro %}
+
+        # Wrap the response in an operation future.
+        response = operation_async.from_gapic(
+            response,
+            self._client._transport.operations_client,
+            {{ method.lro.response_type.ident }},
+            metadata_type={{ method.lro.metadata_type.ident }},
+        )
+        {%- elif method.paged_result_field %}
+
+        # This method is paged; wrap the response in a pager, which provides
+        # an `__aiter__` convenience method.
+        response = {{ method.client_output_async.ident }}(
+            method=rpc,
+            request=request,
+            response=response,
+        )
+        {%- endif %}
+        {%- if not method.void %}
+
+        # Done; return the response.
+        return response
+        {%- endif %}
+        {{ '\n' }}
+    {% endfor %}
+
+
+try:
+    _client_info = gapic_v1.client_info.ClientInfo(
+        gapic_version=pkg_resources.get_distribution(
+            '{{ api.naming.warehouse_package_name }}',
+        ).version,
+    )
+except pkg_resources.DistributionNotFound:
+    _client_info = gapic_v1.client_info.ClientInfo()
+
+
+__all__ = (
+    '{{ service.async_client_name }}',
+)
+{% endblock %}