From 102bfe5d48961f2a7af4942c3f58d86c8c84bdd9 Mon Sep 17 00:00:00 2001 From: shmookoff Date: Fri, 15 Apr 2022 02:43:30 +0500 Subject: [PATCH 1/5] Context variables for managing sessions. --- aiogoogle/client.py | 35 +++++++++++++++++++++++++---------- 1 file changed, 25 insertions(+), 10 deletions(-) diff --git a/aiogoogle/client.py b/aiogoogle/client.py index 27755bd..84bee01 100644 --- a/aiogoogle/client.py +++ b/aiogoogle/client.py @@ -1,5 +1,7 @@ __all__ = ["Aiogoogle"] +from contextvars import ContextVar +from uuid import uuid4 from .resource import GoogleAPI from .auth.managers import Oauth2Manager, ApiKeyManager, OpenIdConnectManager, ServiceAccountManager @@ -53,7 +55,9 @@ def __init__( ): self.session_factory = session_factory - self.active_session = None + # Guarantees that each context manager gets its own active_session. + self.sessions = {} + self.active_session_id = ContextVar("active_session", default=None) # Keys self.api_key = api_key @@ -378,21 +382,32 @@ 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 _ensure_session_set(self): + active_session_id = self.active_session_id.get() + if active_session_id is None: + active_session_id = uuid4() + self.active_session_id.set(active_session_id) + + active_session = self.session_factory() + self.sessions[active_session_id] = active_session + else: + active_session = self.sessions[active_session_id] + return active_session async def send(self, *args, **kwargs): - await self._ensure_session_set() - return await self.active_session.send(*args, **kwargs) + active_session = self._ensure_session_set() + return await active_session.send(*args, **kwargs) async def __aenter__(self): - await self._ensure_session_set() - await self.active_session.__aenter__() + active_session = self._ensure_session_set() + await active_session.__aenter__() return self async def __aexit__(self, *args): - await self.active_session.__aexit__(*args) + active_session_id = self.active_session_id.get() + active_session = self.sessions[active_session_id] + await active_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.active_session_id.set(None) + From cdc4b608dc545095594e9bbe5a445242a128aa6c Mon Sep 17 00:00:00 2001 From: shmookoff Date: Fri, 15 Apr 2022 03:02:39 +0500 Subject: [PATCH 2/5] delete closing session from sessions --- aiogoogle/client.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/aiogoogle/client.py b/aiogoogle/client.py index 84bee01..fb75e1b 100644 --- a/aiogoogle/client.py +++ b/aiogoogle/client.py @@ -405,7 +405,7 @@ async def __aenter__(self): async def __aexit__(self, *args): active_session_id = self.active_session_id.get() - active_session = self.sessions[active_session_id] + active_session = self.sessions.pop(active_session_id) await active_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 From 7326d42f55d8cf964e7673c5a0bf8bd8d1a1031e Mon Sep 17 00:00:00 2001 From: shmookoff Date: Fri, 15 Apr 2022 15:46:59 +0500 Subject: [PATCH 3/5] removed active_session storage as it isn't needed. --- aiogoogle/client.py | 27 +++++++++------------------ 1 file changed, 9 insertions(+), 18 deletions(-) diff --git a/aiogoogle/client.py b/aiogoogle/client.py index fb75e1b..6365577 100644 --- a/aiogoogle/client.py +++ b/aiogoogle/client.py @@ -1,7 +1,6 @@ __all__ = ["Aiogoogle"] from contextvars import ContextVar -from uuid import uuid4 from .resource import GoogleAPI from .auth.managers import Oauth2Manager, ApiKeyManager, OpenIdConnectManager, ServiceAccountManager @@ -56,8 +55,7 @@ def __init__( self.session_factory = session_factory # Guarantees that each context manager gets its own active_session. - self.sessions = {} - self.active_session_id = ContextVar("active_session", default=None) + self.session_context: ContextVar[session_factory] = ContextVar("active_session", default=None) # Keys self.api_key = api_key @@ -383,16 +381,11 @@ async def as_anon(self, *requests, timeout=None, full_res=False, raise_for_statu ) def _ensure_session_set(self): - active_session_id = self.active_session_id.get() - if active_session_id is None: - active_session_id = uuid4() - self.active_session_id.set(active_session_id) - - active_session = self.session_factory() - self.sessions[active_session_id] = active_session - else: - active_session = self.sessions[active_session_id] - return active_session + session = self.session_context.get() + if session is None: + session = self.session_factory() + self.session_context.set(session) + return session async def send(self, *args, **kwargs): active_session = self._ensure_session_set() @@ -404,10 +397,8 @@ async def __aenter__(self): return self async def __aexit__(self, *args): - active_session_id = self.active_session_id.get() - active_session = self.sessions.pop(active_session_id) - await active_session.__aexit__(*args) + session = self.session_context.get() + 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_id.set(None) - + self.session_context.set(None) From 4913f9c810e8663af59803047c6b4f51385ec71f Mon Sep 17 00:00:00 2001 From: shmookoff Date: Fri, 15 Apr 2022 16:11:01 +0500 Subject: [PATCH 4/5] naming in __aenter__ --- aiogoogle/client.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/aiogoogle/client.py b/aiogoogle/client.py index 6365577..036f947 100644 --- a/aiogoogle/client.py +++ b/aiogoogle/client.py @@ -392,8 +392,8 @@ async def send(self, *args, **kwargs): return await active_session.send(*args, **kwargs) async def __aenter__(self): - active_session = self._ensure_session_set() - await active_session.__aenter__() + session = self._ensure_session_set() + await session.__aenter__() return self async def __aexit__(self, *args): From 0183639ce55b276a97e3ffae433fd84d17cf1aa7 Mon Sep 17 00:00:00 2001 From: shmookoff Date: Fri, 15 Apr 2022 19:16:14 +0500 Subject: [PATCH 5/5] Nesting context raise exception. --- aiogoogle/client.py | 28 +++++++++++++++++----------- 1 file changed, 17 insertions(+), 11 deletions(-) diff --git a/aiogoogle/client.py b/aiogoogle/client.py index 036f947..f5dba8f 100644 --- a/aiogoogle/client.py +++ b/aiogoogle/client.py @@ -380,24 +380,30 @@ async def as_anon(self, *requests, timeout=None, full_res=False, raise_for_statu session_factory=self.session_factory ) - def _ensure_session_set(self): - session = self.session_context.get() - if session is None: - session = self.session_factory() - self.session_context.set(session) + 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): - active_session = self._ensure_session_set() - return await 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): - session = self._ensure_session_set() - await 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): - session = self.session_context.get() + 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