diff --git a/aiogoogle/client.py b/aiogoogle/client.py index 27755bd..f5dba8f 100644 --- a/aiogoogle/client.py +++ b/aiogoogle/client.py @@ -1,5 +1,6 @@ __all__ = ["Aiogoogle"] +from contextvars import ContextVar from .resource import GoogleAPI from .auth.managers import Oauth2Manager, ApiKeyManager, OpenIdConnectManager, ServiceAccountManager @@ -53,7 +54,8 @@ def __init__( ): self.session_factory = session_factory - self.active_session = None + # Guarantees that each context manager gets its own active_session. + self.session_context: ContextVar[session_factory] = ContextVar("active_session", default=None) # Keys self.api_key = api_key @@ -378,21 +380,31 @@ async def as_anon(self, *requests, timeout=None, full_res=False, raise_for_statu session_factory=self.session_factory ) - async def _ensure_session_set(self): - if self.active_session is None: - self.active_session = self.session_factory() + def _get_session(self): + return self.session_context.get() + + def _set_session(self): + session = self.session_factory() + self.session_context.set(session) + return session async def send(self, *args, **kwargs): - await self._ensure_session_set() - return await self.active_session.send(*args, **kwargs) + session = self._get_session() + if session is None: + session = self._set_session() + return await session.send(*args, **kwargs) async def __aenter__(self): - await self._ensure_session_set() - await self.active_session.__aenter__() - return self + session = self._get_session() + if session is None: + session = self._set_session() + await session.__aenter__() + return self + raise RuntimeError("Nesting context managers using the same Aiogoogle object is not allowed.") async def __aexit__(self, *args): - await self.active_session.__aexit__(*args) + session = self._get_session() + await session.__aexit__(*args) # Had to add this because there's no use of keeping a closed session # Closed sessions cannot be reopened, so it's better to just get rid of the object - self.active_session = None + self.session_context.set(None)