Skip to content

Commit 3799568

Browse files
feat(auth): Support sync credentials in AsyncAuthorizedSession (#18542)
This PR allows AsyncAuthorizedSession to support standard credential types by wrapping them in an adapter --------- Co-authored-by: Anthonios Partheniou <partheniou@google.com>
1 parent 97bc619 commit 3799568

3 files changed

Lines changed: 386 additions & 28 deletions

File tree

‎packages/google-api-core/tests/unit/operations_v1/test_operations_rest_client.py‎

Lines changed: 12 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1241,24 +1241,20 @@ def test_operations_base_transport_with_adc():
12411241
)
12421242
def test_operations_auth_adc(client_class):
12431243
# If no credentials are provided, we should use ADC credentials.
1244+
is_async = "async" in str(client_class).lower()
1245+
if is_async and parse_version_to_tuple(auth_version) < (2, 60, 0):
1246+
# Older versions of google-auth do not accept the synchronous credentials
1247+
# returned by ADC in `AsyncAuthorizedSession`.
1248+
pytest.skip("ADC with the async REST transport requires google-auth >= 2.60.0")
1249+
12441250
with mock.patch.object(google.auth, "default", autospec=True) as adc:
12451251
adc.return_value = (ga_credentials.AnonymousCredentials(), None)
1246-
1247-
if "async" in str(client_class).lower():
1248-
# TODO(): Add support for adc to async REST transport.
1249-
# NOTE: Ideally, the logic for adc shouldn't be called if transport
1250-
# is set to async REST. If the user does not configure credentials
1251-
# of type `google.auth.aio.credentials.Credentials`,
1252-
# we should raise an exception to avoid the adc workflow.
1253-
with pytest.raises(google.auth.exceptions.InvalidType):
1254-
client_class()
1255-
else:
1256-
client_class()
1257-
adc.assert_called_once_with(
1258-
scopes=None,
1259-
default_scopes=(),
1260-
quota_project_id=None,
1261-
)
1252+
client_class()
1253+
adc.assert_called_once_with(
1254+
scopes=None,
1255+
default_scopes=(),
1256+
quota_project_id=None,
1257+
)
12621258

12631259

12641260
# TODO(https://github.com/googleapis/python-api-core/issues/705): Add

‎packages/google-auth/google/auth/aio/transport/sessions.py‎

Lines changed: 117 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -24,8 +24,9 @@
2424
from contextlib import asynccontextmanager
2525
from typing import TYPE_CHECKING, Mapping, Optional, Union
2626

27+
import google.auth.credentials
2728
import google.auth.transport._mtls_helper
28-
from google.auth import _exponential_backoff, exceptions
29+
from google.auth import _exponential_backoff, _helpers, exceptions
2930
from google.auth.aio import transport
3031
from google.auth.aio.credentials import Credentials
3132
from google.auth.aio.transport import mtls
@@ -104,6 +105,106 @@ async def with_timeout(coro):
104105
_remaining_time()
105106

106107

108+
class _SyncCredentialsAdapter(Credentials):
109+
"""Adapts synchronous credentials to the asynchronous credentials interface.
110+
111+
:class:`AsyncAuthorizedSession` wraps :class:`google.auth.credentials.Credentials`
112+
(e.g. application default credentials) with this adapter so that they can be
113+
used with an asynchronous transport. Calls are delegated to the wrapped
114+
credentials using a synchronous transport, and blocking calls such as
115+
refreshing the access token run in a worker thread so that the event loop
116+
is not blocked.
117+
118+
Args:
119+
credentials (google.auth.credentials.Credentials): The synchronous
120+
credentials to adapt.
121+
"""
122+
123+
def __init__(self, credentials: google.auth.credentials.Credentials):
124+
self._credentials = credentials
125+
# Synchronous credentials cannot use the asynchronous transport of the
126+
# session, so they are called with a synchronous transport instead.
127+
self._sync_request_instance = None
128+
# Synchronous credentials are not safe to refresh concurrently, which
129+
# concurrent requests would otherwise do from multiple worker threads.
130+
# Instead, at most one refresh is in flight and concurrent callers share it.
131+
self._pending_refresh: Optional["asyncio.Task[None]"] = None
132+
133+
@property
134+
def _sync_request(self):
135+
if self._sync_request_instance is None:
136+
# Imported here because `requests` is an optional dependency of
137+
# google-auth. It is installed alongside `aiohttp` by the `aiohttp` extra.
138+
from google.auth.transport import requests as sync_requests
139+
140+
self._sync_request_instance = sync_requests.Request()
141+
return self._sync_request_instance
142+
143+
def close(self):
144+
if (
145+
self._sync_request_instance is not None
146+
and hasattr(self._sync_request_instance, "session")
147+
and self._sync_request_instance.session is not None
148+
):
149+
self._sync_request_instance.session.close()
150+
151+
@property
152+
def token(self):
153+
"""Optional[str]: The bearer token that can be used in HTTP headers to make
154+
authenticated requests."""
155+
return self._credentials.token
156+
157+
@property
158+
def expiry(self):
159+
"""Optional[datetime]: When the token expires and is no longer valid.
160+
If this is None, the token is assumed to never expire."""
161+
return self._credentials.expiry
162+
163+
@property
164+
@_helpers.copy_docstring(google.auth.credentials.Credentials)
165+
def valid(self):
166+
return self._credentials.valid
167+
168+
@property
169+
@_helpers.copy_docstring(google.auth.credentials.Credentials)
170+
def expired(self):
171+
return self._credentials.expired
172+
173+
async def _refresh_shared(self):
174+
"""Refreshes the wrapped credentials, joining a refresh already in flight.
175+
176+
The refresh is shielded from cancellation: a caller that is cancelled
177+
while waiting (e.g. because of a timeout) stops waiting, but the refresh
178+
completes so that the next caller joins it rather than starting a
179+
second, concurrent refresh.
180+
"""
181+
if self._pending_refresh is None or self._pending_refresh.done():
182+
self._pending_refresh = asyncio.create_task(
183+
asyncio.to_thread(self._credentials.refresh, self._sync_request)
184+
)
185+
await asyncio.shield(self._pending_refresh)
186+
187+
@_helpers.copy_docstring(Credentials)
188+
async def apply(self, headers, token=None):
189+
self._credentials.apply(headers, token=token)
190+
191+
@_helpers.copy_docstring(Credentials)
192+
async def refresh(self, request):
193+
await self._refresh_shared()
194+
195+
@_helpers.copy_docstring(Credentials)
196+
async def before_request(self, request, method, url, headers):
197+
if not self._credentials.valid:
198+
await self._refresh_shared()
199+
await asyncio.to_thread(
200+
self._credentials.before_request,
201+
self._sync_request,
202+
method,
203+
url,
204+
headers,
205+
)
206+
207+
107208
class AsyncAuthorizedSession:
108209
"""This is an asynchronous implementation of :class:`google.auth.requests.AuthorizedSession` class.
109210
We utilize an instance of a class that implements :class:`google.auth.aio.transport.Request` configured
@@ -126,8 +227,9 @@ class AsyncAuthorizedSession:
126227
credentials' headers to the request and refreshing credentials as needed.
127228
128229
Args:
129-
credentials (google.auth.aio.credentials.Credentials):
130-
The credentials to add to the request.
230+
credentials (Union[google.auth.aio.credentials.Credentials, google.auth.credentials.Credentials]):
231+
The credentials to add to the request. Synchronous credentials
232+
(e.g. application default credentials) are also supported.
131233
auth_request (Optional[google.auth.aio.transport.Request]):
132234
An instance of a class that implements
133235
:class:`~google.auth.aio.transport.Request` used to make requests
@@ -139,17 +241,22 @@ class AsyncAuthorizedSession:
139241
- google.auth.exceptions.TransportError: If `auth_request` is `None`
140242
and the external package `aiohttp` is not installed.
141243
- google.auth.exceptions.InvalidType: If the provided credentials are
142-
not of type `google.auth.aio.credentials.Credentials`.
244+
not of type `google.auth.aio.credentials.Credentials` or
245+
`google.auth.credentials.Credentials`.
143246
"""
144247

145248
def __init__(
146-
self, credentials: Credentials, auth_request: Optional[transport.Request] = None
249+
self,
250+
credentials: Union[Credentials, google.auth.credentials.Credentials],
251+
auth_request: Optional[transport.Request] = None,
147252
):
148-
if not isinstance(credentials, Credentials):
253+
if isinstance(credentials, google.auth.credentials.Credentials):
254+
credentials = _SyncCredentialsAdapter(credentials)
255+
elif not isinstance(credentials, Credentials):
149256
raise exceptions.InvalidType(
150-
f"The configured credentials of type {type(credentials)} are invalid and must be of type `google.auth.aio.credentials.Credentials`"
257+
f"The configured credentials of type {type(credentials)} are invalid and must be of type `google.auth.aio.credentials.Credentials` or `google.auth.credentials.Credentials`"
151258
)
152-
self._credentials = credentials
259+
self._credentials: Credentials = credentials
153260
_auth_request = auth_request
154261
if not _auth_request and AIOHTTP_INSTALLED:
155262
_auth_request = AiohttpRequest()
@@ -846,6 +953,8 @@ async def close(self) -> None:
846953
if inspect.isawaitable(res):
847954
await res
848955
finally:
956+
if hasattr(self._credentials, "close"):
957+
self._credentials.close()
849958
for old_request in self._old_auth_requests:
850959
try:
851960
if hasattr(old_request, "close"):

0 commit comments

Comments
 (0)