Skip to content

Commit be67ab6

Browse files
nbayatinbayati
committed
feat(auth): support async blocking RAB lookups and add support to async jwt
1 parent a952527 commit be67ab6

5 files changed

Lines changed: 182 additions & 53 deletions

File tree

‎packages/google-auth/google/auth/_credentials_async.py‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import abc
1919
import inspect
2020

21+
from google.auth import _regional_access_boundary_utils
2122
from google.auth import credentials
2223

2324

@@ -189,3 +190,74 @@ def with_scopes_if_required(credentials, scopes):
189190

190191
class Signing(credentials.Signing, metaclass=abc.ABCMeta):
191192
"""Interface for credentials that can cryptographically sign messages."""
193+
194+
195+
class CredentialsWithRegionalAccessBoundary(
196+
Credentials, credentials.CredentialsWithRegionalAccessBoundary
197+
):
198+
"""Async base for credentials supporting regional access boundary configuration."""
199+
200+
def __init__(self):
201+
super().__init__()
202+
self._rab_manager.refresh_manager = (
203+
_regional_access_boundary_utils._AsyncRegionalAccessBoundaryRefreshManager()
204+
)
205+
206+
def __setstate__(self, state):
207+
super().__setstate__(state)
208+
self._rab_manager.refresh_manager = (
209+
_regional_access_boundary_utils._AsyncRegionalAccessBoundaryRefreshManager()
210+
)
211+
212+
async def _after_refresh(self, request, method, url, headers):
213+
"""Triggers the Regional Access Boundary lookup asynchronously if necessary."""
214+
await self._maybe_start_regional_access_boundary_refresh_async(request, url)
215+
216+
async def _maybe_start_regional_access_boundary_refresh_async(self, request, url):
217+
"""Starts a background refresh or performs a blocking refresh asynchronously.
218+
219+
Args:
220+
request (google.auth.aio.transport.Request): The object used to make
221+
HTTP requests.
222+
url (str): The URL of the request.
223+
"""
224+
# Do not perform a lookup if the request is for a regional endpoint.
225+
if self._is_regional_endpoint(url):
226+
return
227+
228+
# A refresh is only needed if the feature is enabled.
229+
if not self._is_regional_access_boundary_lookup_required():
230+
return
231+
232+
# Trigger background or blocking refresh if needed.
233+
await self._rab_manager.maybe_start_refresh_async(self, request)
234+
235+
async def _lookup_regional_access_boundary(self, request, fail_fast=False):
236+
"""Calls the Regional Access Boundary lookup API asynchronously.
237+
238+
Args:
239+
request (google.auth.aio.transport.Request): The object used to make
240+
HTTP requests.
241+
fail_fast (bool): Whether the lookup should fail fast (short timeout, no retries).
242+
243+
Returns:
244+
Optional[Dict[str, str]]: The Regional Access Boundary information
245+
returned by the lookup API, or None if the lookup failed.
246+
"""
247+
url_builder = self._build_regional_access_boundary_lookup_url
248+
if inspect.iscoroutinefunction(url_builder):
249+
url = await url_builder(request=request)
250+
else:
251+
url = url_builder(request=request)
252+
253+
if not url:
254+
return None
255+
256+
headers = {}
257+
self._apply(headers)
258+
259+
from google.oauth2 import _client_async
260+
261+
return await _client_async._lookup_regional_access_boundary(
262+
request, url, headers=headers, fail_fast=fail_fast
263+
)

‎packages/google-auth/google/auth/_jwt_async.py‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -91,7 +91,9 @@ def decode(token, certs=None, verify=True, audience=None):
9191

9292

9393
class Credentials(
94-
jwt.Credentials, _credentials_async.Signing, _credentials_async.Credentials
94+
jwt.Credentials,
95+
_credentials_async.Signing,
96+
_credentials_async.CredentialsWithRegionalAccessBoundary,
9597
):
9698
"""Credentials that use a JWT as the bearer token.
9799
@@ -142,6 +144,15 @@ class Credentials(
142144
new_credentials = credentials.with_claims(audience=new_audience)
143145
"""
144146

147+
def __setstate__(self, state):
148+
"""Restores the credential state and ensures the async refresh manager is attached."""
149+
super().__setstate__(state)
150+
from google.auth import _regional_access_boundary_utils
151+
152+
self._rab_manager.refresh_manager = (
153+
_regional_access_boundary_utils._AsyncRegionalAccessBoundaryRefreshManager()
154+
)
155+
145156

146157
class OnDemandCredentials(
147158
jwt.OnDemandCredentials, _credentials_async.Signing, _credentials_async.Credentials

‎packages/google-auth/google/auth/_regional_access_boundary_utils.py‎

Lines changed: 64 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -171,12 +171,11 @@ def apply_headers(self, headers):
171171
else:
172172
headers.pop(_REGIONAL_ACCESS_BOUNDARY_HEADER, None)
173173

174-
def maybe_start_refresh(self, credentials, request):
175-
"""Starts a background thread to refresh the Regional Access Boundary if needed.
174+
def _should_refresh(self):
175+
"""Checks if the Regional Access Boundary data needs a refresh and is not in cooldown.
176176
177-
Args:
178-
credentials (google.auth.credentials.Credentials): The credentials to refresh.
179-
request (google.auth.transport.Request): The object used to make HTTP requests.
177+
Returns:
178+
bool: True if a refresh is required, False otherwise.
180179
"""
181180
rab_data = self._data
182181

@@ -187,10 +186,22 @@ def maybe_start_refresh(self, credentials, request):
187186
and _helpers.utcnow()
188187
< (rab_data.expiry - REGIONAL_ACCESS_BOUNDARY_REFRESH_THRESHOLD)
189188
):
190-
return
189+
return False
191190

192191
# Don't start a new refresh if the cooldown is still in effect.
193192
if rab_data.cooldown_expiry and _helpers.utcnow() < rab_data.cooldown_expiry:
193+
return False
194+
195+
return True
196+
197+
def maybe_start_refresh(self, credentials, request):
198+
"""Starts a background thread to refresh the Regional Access Boundary if needed.
199+
200+
Args:
201+
credentials (google.auth.credentials.Credentials): The credentials to refresh.
202+
request (google.auth.transport.Request): The object used to make HTTP requests.
203+
"""
204+
if not self._should_refresh():
194205
return
195206

196207
# If all checks pass, start the background refresh.
@@ -199,6 +210,22 @@ def maybe_start_refresh(self, credentials, request):
199210
else:
200211
self.refresh_manager.start_refresh(credentials, request, self)
201212

213+
async def maybe_start_refresh_async(self, credentials, request):
214+
"""Starts a background refresh or performs a blocking refresh asynchronously.
215+
216+
Args:
217+
credentials (google.auth.credentials.Credentials): The credentials to refresh.
218+
request (google.auth.aio.transport.Request): The object used to make HTTP requests.
219+
"""
220+
if not self._should_refresh():
221+
return
222+
223+
# If all checks pass, start the refresh.
224+
if self._use_blocking_regional_access_boundary_lookup:
225+
await self.start_blocking_refresh_async(credentials, request)
226+
else:
227+
self.refresh_manager.start_refresh(credentials, request, self)
228+
202229
def start_blocking_refresh(self, credentials, request):
203230
"""Initiates a blocking lookup of the Regional Access Boundary.
204231
@@ -228,6 +255,37 @@ def start_blocking_refresh(self, credentials, request):
228255

229256
self.process_regional_access_boundary_info(regional_access_boundary_info)
230257

258+
async def start_blocking_refresh_async(self, credentials, request):
259+
"""Initiates a blocking lookup of the Regional Access Boundary asynchronously.
260+
261+
If the lookup raises an exception, it is caught and logged as a warning,
262+
and the lookup is treated as a failure (entering cooldown). Exceptions
263+
are not propagated to the caller.
264+
265+
Args:
266+
credentials (google.auth.credentials.Credentials): The credentials to refresh.
267+
request (google.auth.aio.transport.Request): The object used to make HTTP requests.
268+
"""
269+
try:
270+
# The fail_fast parameter is set to True to ensure we don't block the calling
271+
# thread for too long. This will do two things: 1) set a timeout to 3s
272+
# instead of the default 120s and 2) ensure we do not retry at all
273+
regional_access_boundary_info = (
274+
await credentials._lookup_regional_access_boundary(
275+
request, fail_fast=True
276+
)
277+
)
278+
except Exception as e:
279+
if _helpers.is_logging_enabled(_LOGGER):
280+
_LOGGER.warning(
281+
"Regional Access Boundary lookup raised an exception: %s",
282+
e,
283+
exc_info=True,
284+
)
285+
regional_access_boundary_info = None
286+
287+
self.process_regional_access_boundary_info(regional_access_boundary_info)
288+
231289
def process_regional_access_boundary_info(self, regional_access_boundary_info):
232290
"""Processes the regional access boundary info and updates the state.
233291

‎packages/google-auth/google/auth/credentials.py‎

Lines changed: 29 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -434,18 +434,14 @@ def _set_blocking_regional_access_boundary_lookup(self):
434434
self._rab_manager.enable_blocking_lookup()
435435
return self
436436

437-
def _maybe_start_regional_access_boundary_refresh(self, request, url):
438-
"""
439-
Starts a background thread to refresh the Regional Access Boundary if needed.
440-
441-
This method checks if a refresh is necessary and if one is not already
442-
in progress or in a cooldown period. If so, it starts a background
443-
thread to perform the lookup.
437+
def _is_regional_endpoint(self, url):
438+
"""Checks if the request URL is for a regional endpoint.
444439
445440
Args:
446-
request (google.auth.transport.Request): The object used to make
447-
HTTP requests.
448441
url (str): The URL of the request.
442+
443+
Returns:
444+
bool: True if the URL is a regional endpoint, False otherwise.
449445
"""
450446
try:
451447
# Do not perform a lookup if the request is for a regional endpoint.
@@ -454,16 +450,35 @@ def _maybe_start_regional_access_boundary_refresh(self, request, url):
454450
hostname.endswith(".rep.googleapis.com")
455451
or hostname.endswith(".rep.sandbox.googleapis.com")
456452
):
457-
return
458-
except (ValueError, TypeError):
453+
return True
454+
except (ValueError, TypeError, AttributeError):
459455
# If the URL is malformed, proceed with the default lookup behavior.
460456
pass
461457

458+
return False
459+
460+
def _maybe_start_regional_access_boundary_refresh(self, request, url):
461+
"""
462+
Starts a background thread to refresh the Regional Access Boundary if needed.
463+
464+
This method checks if a refresh is necessary and if one is not already
465+
in progress or in a cooldown period. If so, it starts a background
466+
thread to perform the lookup.
467+
468+
Args:
469+
request (google.auth.transport.Request): The object used to make
470+
HTTP requests.
471+
url (str): The URL of the request.
472+
"""
473+
# Do not perform a lookup if the request is for a regional endpoint.
474+
if self._is_regional_endpoint(url):
475+
return
476+
462477
# A refresh is only needed if the feature is enabled.
463478
if not self._is_regional_access_boundary_lookup_required():
464479
return
465480

466-
# Start the background refresh if needed.
481+
# Trigger background or blocking refresh if needed
467482
self._rab_manager.maybe_start_refresh(self, request)
468483

469484
def _is_regional_access_boundary_lookup_required(self):
@@ -475,11 +490,11 @@ def _is_regional_access_boundary_lookup_required(self):
475490
Returns:
476491
bool: True if a Regional Access Boundary lookup is required, False otherwise.
477492
"""
478-
# 1. Check if the feature is enabled.
493+
# Check if the feature is enabled.
479494
if not _regional_access_boundary_utils.is_regional_access_boundary_enabled():
480495
return False
481496

482-
# 2. Skip for non-default universe domains.
497+
# Skip for non-default universe domains.
483498
if self.universe_domain != DEFAULT_UNIVERSE_DOMAIN:
484499
return False
485500

@@ -526,7 +541,6 @@ def _lookup_regional_access_boundary(
526541

527542
headers: Dict[str, str] = {}
528543
self._apply(headers)
529-
self._rab_manager.apply_headers(headers)
530544
return _client._lookup_regional_access_boundary(
531545
request, url, headers=headers, fail_fast=fail_fast
532546
)

‎packages/google-auth/google/oauth2/_service_account_async.py‎

Lines changed: 5 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -24,13 +24,14 @@
2424

2525
from google.auth import _credentials_async as credentials_async
2626
from google.auth import _helpers
27-
from google.auth import _regional_access_boundary_utils
2827
from google.oauth2 import _client_async
2928
from google.oauth2 import service_account
3029

3130

3231
class Credentials(
33-
service_account.Credentials, credentials_async.Scoped, credentials_async.Credentials
32+
service_account.Credentials,
33+
credentials_async.Scoped,
34+
credentials_async.CredentialsWithRegionalAccessBoundary,
3435
):
3536
"""Service account credentials
3637
@@ -67,15 +68,11 @@ class Credentials(
6768
credentials = credentials.with_quota_project('myproject-123')
6869
"""
6970

70-
def __init__(self, *args, **kwargs):
71-
super().__init__(*args, **kwargs)
72-
self._rab_manager.refresh_manager = (
73-
_regional_access_boundary_utils._AsyncRegionalAccessBoundaryRefreshManager()
74-
)
75-
7671
def __setstate__(self, state):
7772
"""Restores the credential state and ensures the async refresh manager is attached."""
7873
super().__setstate__(state)
74+
from google.auth import _regional_access_boundary_utils
75+
7976
self._rab_manager.refresh_manager = (
8077
_regional_access_boundary_utils._AsyncRegionalAccessBoundaryRefreshManager()
8178
)
@@ -89,29 +86,6 @@ async def refresh(self, request):
8986
self.token = access_token
9087
self.expiry = expiry
9188

92-
async def _lookup_regional_access_boundary(self, request, fail_fast=False):
93-
"""Calls the Regional Access Boundary lookup API to retrieve the Regional Access Boundary information.
94-
95-
Args:
96-
request (google.auth.aio.transport.Request): The object used to make
97-
HTTP requests.
98-
fail_fast (bool): Whether the lookup should fail fast.
99-
100-
Returns:
101-
Optional[Dict[str, str]]: The Regional Access Boundary information.
102-
"""
103-
url = self._build_regional_access_boundary_lookup_url(request=request)
104-
if not url:
105-
return None
106-
107-
headers = {}
108-
self._apply(headers)
109-
self._rab_manager.apply_headers(headers)
110-
111-
return await _client_async._lookup_regional_access_boundary(
112-
request, url, headers=headers, fail_fast=fail_fast
113-
)
114-
11589

11690
class IDTokenCredentials(
11791
service_account.IDTokenCredentials,

0 commit comments

Comments
 (0)