2020import functools
2121import inspect
2222import logging
23- import os
2423import threading
2524from typing import NamedTuple , Optional , TYPE_CHECKING
2625
2726from google .auth import _helpers
28- from google .auth import environment_vars
2927
30- if TYPE_CHECKING :
28+ if TYPE_CHECKING : # pragma: NO COVER
3129 import google .auth .credentials
3230 import google .auth .transport
3331
3432_LOGGER = logging .getLogger (__name__ )
3533
3634
37- @functools .lru_cache ()
38- def is_regional_access_boundary_enabled ():
39- """Checks if Regional Access Boundary is enabled via environment variable.
40-
41- The environment variable is interpreted as a boolean with the following
42- (case-insensitive) rules:
43- - "true", "1" are considered true.
44- - Any other value (or unset) is considered false.
45-
46- Returns:
47- bool: True if Regional Access Boundary is enabled, False otherwise.
48- """
49- value = os .environ .get (environment_vars .GOOGLE_AUTH_TRUST_BOUNDARY_ENABLED )
50- if value is None :
51- return False
52-
53- return value .lower () in ("true" , "1" )
54-
55-
5635# The default lifetime for a cached Regional Access Boundary.
5736DEFAULT_REGIONAL_ACCESS_BOUNDARY_TTL = datetime .timedelta (hours = 6 )
5837
@@ -455,6 +434,61 @@ def start_refresh(self, credentials, request, rab_manager):
455434 self ._worker .start ()
456435
457436
437+ def _prepare_async_lookup_callable (request ):
438+ """Unwraps a request callable, clones the transport, and returns the new callable.
439+
440+ Args:
441+ request: The original request callable (e.g. functools.partial or raw Request).
442+
443+ Returns:
444+ Tuple[Callable, Any, bool]: A tuple containing the new lookup callable, the
445+ underlying request object, and a boolean indicating if it was cloned.
446+ """
447+ is_partial = isinstance (request , functools .partial )
448+ base_callable = request .func if is_partial else request
449+
450+ if not hasattr (base_callable , "_clone" ):
451+ return request , base_callable , False
452+
453+ cloned_callable = base_callable ._clone ()
454+ is_cloned = cloned_callable is not base_callable
455+
456+ if is_partial :
457+ new_request = functools .partial (
458+ cloned_callable , * request .args , ** request .keywords
459+ )
460+ else :
461+ new_request = cloned_callable
462+
463+ return new_request , cloned_callable , is_cloned
464+
465+
466+ async def _close_cloned_request (lookup_request , is_cloned ):
467+ """Safely closes the underlying cloned request transport, if applicable.
468+
469+ Args:
470+ lookup_request (Any): The request object/transport to close.
471+ is_cloned (bool): Whether the request was actually cloned.
472+ """
473+ if not is_cloned or not hasattr (lookup_request , "close" ):
474+ return
475+
476+ is_async = False
477+ try :
478+ maybe_coro = lookup_request .close ()
479+ if is_async := inspect .isawaitable (maybe_coro ):
480+ await maybe_coro
481+ except Exception as e :
482+ if _helpers .is_logging_enabled (_LOGGER ):
483+ adapter_type = " asynchronous " if is_async else " "
484+ _LOGGER .warning (
485+ "Failed to cleanly close cloned%srequest transport: %s" ,
486+ adapter_type ,
487+ e ,
488+ exc_info = True ,
489+ )
490+
491+
458492class _AsyncRegionalAccessBoundaryRefreshManager (object ):
459493 """Manages a task for background refreshing of the Regional Access Boundary in async flows."""
460494
@@ -491,11 +525,28 @@ def start_refresh(self, credentials, request, rab_manager):
491525 # A refresh is already in progress.
492526 return
493527
528+ try :
529+ (
530+ lookup_callable ,
531+ lookup_request ,
532+ is_cloned ,
533+ ) = _prepare_async_lookup_callable (request )
534+ except Exception as e :
535+ if _helpers .is_logging_enabled (_LOGGER ):
536+ _LOGGER .warning (
537+ "Synchronous cloning of request for Regional Access Boundary lookup failed: %s" ,
538+ e ,
539+ exc_info = True ,
540+ )
541+ rab_manager .process_regional_access_boundary_info (None )
542+ return
543+
494544 async def _worker ():
495545 try :
496- # credentials._lookup_regional_access_boundary should be async in the async creds class
497546 regional_access_boundary_info = (
498- await credentials ._lookup_regional_access_boundary (request )
547+ await credentials ._lookup_regional_access_boundary (
548+ lookup_callable
549+ )
499550 )
500551 except Exception as e :
501552 if _helpers .is_logging_enabled (_LOGGER ):
@@ -505,6 +556,8 @@ async def _worker():
505556 exc_info = True ,
506557 )
507558 regional_access_boundary_info = None
559+ finally :
560+ await _close_cloned_request (lookup_request , is_cloned )
508561
509562 rab_manager .process_regional_access_boundary_info (
510563 regional_access_boundary_info
@@ -514,7 +567,15 @@ async def _worker():
514567 try :
515568 self ._worker_task = asyncio .create_task (coro )
516569 except Exception :
570+ # Clean up cloned request if task creation fails
517571 coro .close ()
572+ try :
573+ asyncio .get_running_loop ().create_task (
574+ _close_cloned_request (lookup_request , is_cloned )
575+ )
576+ except RuntimeError :
577+ pass
578+ rab_manager .process_regional_access_boundary_info (None )
518579 raise
519580
520581
0 commit comments