feat: leverage auth session and implement close method

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

@@ -66,7 +66,13 @@ from .transports.grpc_asyncio import {{ service.grpc_asyncio_transport_name }}
 from .transports.rest import {{ service.name }}RestTransport
 {# TODO(https://github.com/googleapis/gapic-generator-python/issues/2121): Remove this condition when async rest is GA. #}
 {% if rest_async_io_enabled %}
-from .transports.rest_asyncio import Async{{ service.name }}RestTransport
+try:
+    from .transports.rest_asyncio import Async{{ service.name }}RestTransport
+    HAS_GOOGLE_AUTH_AIO = True
+except ImportError as e: # pragma: NO COVER
+    HAS_GOOGLE_AUTH_AIO = False
+    GOOGLE_AUTH_AIO_EXCEPTION = e
+
 {% endif %}{# if rest_async_io_enabled #}
 {% endif %}
 
@@ -87,7 +93,8 @@ class {{ service.client_name }}Meta(type):
     _transport_registry["rest"] = {{ service.name }}RestTransport
     {# TODO(https://github.com/googleapis/gapic-generator-python/issues/2121): Remove this condition when async rest is GA. #}
     {% if rest_async_io_enabled %}
-    _transport_registry["rest_asyncio"] = Async{{ service.name }}RestTransport
+    if HAS_GOOGLE_AUTH_AIO:  # pragma: NO COVER
+        _transport_registry["rest_asyncio"] = Async{{ service.name }}RestTransport
     {% endif %}{# if rest_async_io_enabled #}
     {% endif %}
 
@@ -104,6 +111,11 @@ class {{ service.client_name }}Meta(type):
             The transport class to use.
         """
         # If a specific transport is requested, return that one.
+        {# TODO(https://github.com/googleapis/gapic-generator-python/issues/2121): Remove this condition when async rest is GA. #}
+        {% if rest_async_io_enabled %}
+        if label == "rest_asyncio" and not HAS_GOOGLE_AUTH_AIO:  # pragma: NO COVER
+            raise GOOGLE_AUTH_AIO_EXCEPTION
+        {% endif %}
         if label:
             return cls._transport_registry[label]
 
@@ -552,16 +564,52 @@ class {{ service.client_name }}(metaclass={{ service.client_name }}Meta):
                 self._use_mtls_endpoint))
 
         if not transport_provided:
+            {# TODO(https://github.com/googleapis/gapic-generator-python/issues/2121): Remove this condition when async rest is GA. #}
+            {% if rest_async_io_enabled %}
+            transport_init: Union[Type[{{ service.name  }}Transport], Callable[..., {{ service.name }}Transport]] = (
+                {{ service.client_name  }}.get_transport_class(transport)
+                if isinstance(transport, str) or transport is None
+                else cast(Callable[..., {{ service.name }}Transport], transport)
+            )
+
+            if "rest_asyncio" in str(transport_init):
+                # Note: The following parameters are not support in async rest:
+                # credentials_file, scopes, client_cert_source_for_mtls, quota_project_id
+                # always_use_jwt_access, api_audience.
+                unsupported_params = {
+                    "google.api_core.client_options.ClientOptions.credentials_file": self._client_options.credentials_file,
+                    "google.api_core.client_options.ClientOptions.scopes": self._client_options.scopes,
+                    "google.api_core.client_options.ClientOptions.quota_project_id": self._client_options.quota_project_id,
+                    "google.api_core.client_options.ClientOptions.client_cert_source": self._client_options.client_cert_source,
+                    "google.api_core.client_options.ClientOptions.api_audience": self._client_options.api_audience,
+
+                }
+                provided_unsupported_params = [name for name, value in unsupported_params.items() if value is not None]
+                if provided_unsupported_params:
+                    raise NotImplementedError(
+                        f"The following provided parameters are not supported for `transport=rest_asyncio`: {', '.join(provided_unsupported_params)}"
+                    )
+                self._transport = transport_init(
+                    credentials=credentials,
+                    host=self._api_endpoint,
+                    client_info=client_info,
+                )
+                return
+
+            {% endif %}
             import google.auth._default  # type: ignore
 
             if api_key_value and hasattr(google.auth._default, "get_api_key_credentials"):
                 credentials = google.auth._default.get_api_key_credentials(api_key_value)
 
+            {# TODO(https://github.com/googleapis/gapic-generator-python/issues/2121): Remove this condition when async rest is GA. #}
+            {% if not rest_async_io_enabled %}
             transport_init: Union[Type[{{ service.name  }}Transport], Callable[..., {{ service.name }}Transport]] = (
                 {{ service.client_name  }}.get_transport_class(transport)
                 if isinstance(transport, str) or transport is None
                 else cast(Callable[..., {{ service.name }}Transport], transport)
             )
+            {% endif %}
             # initialize with the provided callable or the passed in class
             self._transport = transport_init(
                 credentials=credentials,

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

@@ -4,6 +4,14 @@
 
 {% block content %}
 
+try:
+    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 gapic_v1
 
 from typing import Any, Optional
@@ -32,8 +40,7 @@ class Async{{service.name}}RestTransport(_Base{{ service.name }}RestTransport):
     """
     def __init__(self, *,
             host: str{% if service.host %} = '{{ service.host }}'{% endif %},
-            {# TODO (https://github.com/googleapis/gapic-generator-python/issues/2129): Update the default type for credentials. #}
-            credentials: Optional[Any] = None,
+            credentials: Optional[ga_credentials_async.Credentials] = None,
             client_info: gapic_v1.client_info.ClientInfo = DEFAULT_CLIENT_INFO,
             url_scheme: str = 'https',
             ) -> None:
@@ -48,8 +55,7 @@ class Async{{service.name}}RestTransport(_Base{{ service.name }}RestTransport):
         Args:
             host ({% if service.host %}Optional[str]{% else %}str{% endif %}):
                 {{ ' ' }}The hostname to connect to {% if service.host %}(default: '{{ service.host }}'){% endif %}.
-            {# TODO (https://github.com/googleapis/gapic-generator-python/issues/2129): Update the default type for credentials. #}
-            credentials (Optional[Any]): The
+            credentials (Optional[google.auth.aio.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
@@ -72,9 +78,13 @@ class Async{{service.name}}RestTransport(_Base{{ service.name }}RestTransport):
             url_scheme=url_scheme,
             api_audience=None
         )
+        self._session = AsyncAuthorizedSession(self._credentials)
 
     @property
     def kind(self) -> str:
         return "rest_asyncio"
 
+    async def close(self):
+        await self._session.close()
+
 {% endblock %}

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

@@ -31,7 +31,6 @@ from google.iam.v1 import policy_pb2  # type: ignore
 from google.cloud.location import locations_pb2 # type: ignore
 {% endif %}
 from .base import {{service.name}}Transport, DEFAULT_CLIENT_INFO
-from google.auth import credentials as ga_credentials  # type: ignore
 
 import re
 from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
@@ -56,7 +55,7 @@ class _Base{{ service.name }}RestTransport({{service.name}}Transport):
     {# TODO: handle mtls stuff if that is relevant for rest transport #}
     def __init__(self, *,
             host: str{% if service.host %} = '{{ service.host }}'{% endif %},
-            credentials: Optional[ga_credentials.Credentials] = None,
+            credentials: Optional[Any] = None,
             client_info: gapic_v1.client_info.ClientInfo = DEFAULT_CLIENT_INFO,
             always_use_jwt_access: Optional[bool] = False,
             url_scheme: str = 'https',
@@ -66,7 +65,7 @@ class _Base{{ service.name }}RestTransport({{service.name}}Transport):
         Args:
             host ({% if service.host %}Optional[str]{% else %}str{% endif %}):
                 {{ ' ' }}The hostname to connect to {% if service.host %}(default: '{{ service.host }}'){% endif %}.
-            credentials (Optional[google.auth.credentials.Credentials]): The
+            credentials (Optional[Any]): 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

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

@@ -38,6 +38,10 @@ from google.protobuf import json_format
 
 try:
     from google.auth.aio import credentials as ga_credentials_async
+    {# TODO(https://github.com/googleapis/gapic-generator-python/issues/2121): Remove this condition when async rest is GA. #}
+    {% if rest_async_io_enabled %}
+    from google.auth.aio.transport.sessions import AsyncAuthorizedSession
+    {% endif %}
     HAS_GOOGLE_AUTH_AIO = True
 except ImportError: # pragma: NO COVER
     HAS_GOOGLE_AUTH_AIO = False
@@ -1054,7 +1058,10 @@ def test_transport_adc(transport_class):
 {% for conf in configs %}
 {{ test_macros.transport_kind_test(**conf) }}
 {% endfor %}
-
+{# TODO(https://github.com/googleapis/gapic-generator-python/issues/2121): Remove the below macro when async rest is GA. #}
+{% if rest_async_io_enabled %}
+{{ test_macros.async_rest_unsupported_params_test(service) }}
+{% endif %}
 {% if 'grpc' in opts.transport %}
 def test_transport_grpc_default():
     # A client should use the gRPC transport by default.
@@ -1668,40 +1675,11 @@ def test_client_with_default_client_info():
         )
         prep.assert_called_once_with(client_info)
 
-{% if 'grpc' in opts.transport %}
-@pytest.mark.asyncio
-async def test_transport_close_async():
-    client = {{ service.async_client_name }}(
-        credentials=async_anonymous_credentials(),
-        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()
-{% endif %}
-
 {% include 'tests/unit/gapic/%name_%version/%sub/_test_mixins.py.j2' %}
 
-def test_transport_close():
-    transports = {
-        {% if 'rest' in opts.transport %}
-        "rest": "_session",
-        {% endif %}
-        {% if 'grpc' in opts.transport %}
-        "grpc": "_grpc_channel",
-        {% endif %}
-    }
-
-    for transport, close_name in transports.items():
-        client = {{ service.client_name }}(
-            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()
+{% for conf in configs %}
+{{ test_macros.transport_close_test(**conf) }}
+{% endfor %}
 
 def test_client_ctx():
     transports = [

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

@@ -1892,9 +1892,61 @@ def test_{{ method_name }}_empty_call():
 {% macro transport_kind_test(service, transport, is_async) %}
 {% set transport_name = get_transport_name(transport, is_async) %}
 def test_transport_kind_{{ transport_name }}():
+    {% if transport_name == 'rest_asyncio' %}
+    if not HAS_GOOGLE_AUTH_AIO:
+        {# TODO(https://github.com/googleapis/google-auth-library-python/pull/1577): Update the version of google-auth once the linked PR is merged. #}
+        pytest.skip("google-auth > 2.x.x is required for async rest transport.")
+    {% endif %}
     transport = {{ get_client(service, is_async) }}.get_transport_class("{{ transport_name }}")(
         credentials={{get_credentials(is_async)}}
     )
     assert transport.kind == "{{ transport_name }}"
 
 {% endmacro %}
+
+
+{% macro transport_close_test(service, transport, is_async) %}
+{% set async_prefix = "async " if is_async else "" %}
+{% set async_decorator = "@pytest.mark.asyncio " if is_async else "" %}
+{% set transport_name = get_transport_name(transport, is_async) %}
+{% set close_session = {
+    'rest': "_session",
+    'grpc': "_grpc_channel"}
+-%}
+{{async_decorator}}
+{{async_prefix}}def test_transport_close_{{transport_name}}():
+
+    {% if transport_name == 'rest_asyncio' %}
+    if not HAS_GOOGLE_AUTH_AIO:
+        {# TODO(https://github.com/googleapis/google-auth-library-python/pull/1577): Update the version of google-auth once the linked PR is merged. #}
+        pytest.skip("google-auth > 2.x.x is required for async rest transport.")
+    {% endif %}
+    client = {{ get_client(service, is_async) }}(
+        credentials={{get_credentials(is_async)}},
+        transport="{{transport_name}}"
+    )
+    with mock.patch.object(type(getattr(client.transport, "{{close_session[transport]}}")), "close") as close:
+        {{async_prefix}}with client:
+            close.assert_not_called()
+        close.assert_called_once()
+
+{% endmacro %}
+
+{# TODO(https://github.com/googleapis/gapic-generator-python/issues/2121): Remove / Update this test macro when async rest is GA. #}
+{% macro async_rest_unsupported_params_test(service) %}
+def test_unsupported_parameter_rest_asyncio():
+    if not HAS_GOOGLE_AUTH_AIO:
+        {# TODO(https://github.com/googleapis/google-auth-library-python/pull/1577): Update the version of google-auth once the linked PR is merged. #}
+        pytest.skip("google-auth > 2.x.x is required for async rest transport.")
+    options = client_options.ClientOptions(quota_project_id="octopus")
+    with pytest.raises(NotImplementedError) as exc:
+        client = {{ get_client(service, True) }}(
+            credentials={{get_credentials(True)}},
+            transport="rest_asyncio",
+            client_options=options
+        )
+    exc.match(
+        "The following provided parameters are not supported for `transport=rest_asyncio`: google.api_core.client_options.ClientOptions.quota_project_id"
+    )
+
+{% endmacro %}

noxfile.py

@@ -440,6 +440,12 @@ def showcase_unit_w_rest_async(
     """Run the generated unit tests with async rest transport against the Showcase library."""
     with showcase_library(session, templates=templates, other_opts=other_opts, rest_async_io_enabled=True) as lib:
         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/686): Update the version of google-api-core once the linked PR is merged.
+        session.install("google-api-core[grpc]@git+https://github.com/googleapis/python-api-core.git@bc811e5c88d2f7fc9e086b61f1896b249d3ca869")
+        # 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")
         run_showcase_unit_tests(session)
 
 

tests/integration/goldens/asset/google/cloud/asset_v1/services/asset_service/transports/rest_base.py

@@ -19,7 +19,6 @@
 
 from google.protobuf import json_format
 from .base import AssetServiceTransport, DEFAULT_CLIENT_INFO
-from google.auth import credentials as ga_credentials  # type: ignore
 
 import re
 from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
@@ -45,7 +44,7 @@ class _BaseAssetServiceRestTransport(AssetServiceTransport):
 
     def __init__(self, *,
             host: str = 'cloudasset.googleapis.com',
-            credentials: Optional[ga_credentials.Credentials] = None,
+            credentials: Optional[Any] = None,
             client_info: gapic_v1.client_info.ClientInfo = DEFAULT_CLIENT_INFO,
             always_use_jwt_access: Optional[bool] = False,
             url_scheme: str = 'https',
@@ -55,7 +54,7 @@ def __init__(self, *,
         Args:
             host (Optional[str]):
                  The hostname to connect to (default: 'cloudasset.googleapis.com').
-            credentials (Optional[google.auth.credentials.Credentials]): The
+            credentials (Optional[Any]): 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