fix: Fix Session ContextVar Manager

amisadmin · amisadmin · commit 101a1939c0a0 · 2022-09-28T18:44:10.000+08:00
diff --git a/sqlalchemy_database/_abc_async_database.py b/sqlalchemy_database/_abc_async_database.py
@@ -37,9 +37,9 @@ async def asgi_dispatch(self, request, call_next):
             ```Python
             app = FastAPI()
             db = Database.create("sqlite:///test.db")
-            app.add_middleware(BaseHTTPMiddleware, db.asgi_dispatch)
+            app.add_middleware(BaseHTTPMiddleware, dispatch=db.asgi_dispatch)
             ```
         """
-        async with self:
+        async with self.__call__():
             response = await call_next(request)
         return response
diff --git a/sqlalchemy_database/_abc_async_database.pyi b/sqlalchemy_database/_abc_async_database.pyi
@@ -16,6 +16,8 @@ from sqlalchemy.ext.asyncio import AsyncConnection
 from sqlalchemy.sql import Executable
 from typing_extensions import Concatenate, ParamSpec
 
+from sqlalchemy_database.database import AsyncSessionContextVarManager
+
 try:
     from sqlmodel import Session
     from sqlmodel.ext.asyncio.session import AsyncSession
@@ -92,3 +94,5 @@ class AbcAsyncDatabase(metaclass=abc.ABCMeta):
         **kwargs: _P.kwargs,
     ) -> Union[_T, _R]: ...
     async def asgi_dispatch(self, request, call_next): ...
+    def __call__(self) -> AsyncSessionContextVarManager:
+        pass
diff --git a/sqlalchemy_database/database.py b/sqlalchemy_database/database.py
@@ -1,5 +1,4 @@
 from contextvars import ContextVar
-from threading import Lock
 from typing import (
     Any,
     AsyncGenerator,
@@ -69,8 +68,6 @@ def __init__(self, engine: AsyncEngine, **session_options):
                 await session.commit()
             ```
         """
-        self._session_lock = Lock()
-        self._session_enter_count = 0
         self._session_context_var: ContextVar[Optional[AsyncSession]] = ContextVar("_session_context_var", default=None)
 
     @property
@@ -82,37 +79,22 @@ def session(self) -> Optional[AsyncSession]:
         Example:
             ```Python
             app = FastAPI()
-            app.add_middleware(BaseHTTPMiddleware, db.asgi_dispatch)
+            app.add_middleware(BaseHTTPMiddleware, dispatch=db.asgi_dispatch)
 
             @app.get('/get_user')
             async def get_user(id:int):
                 return await db.session.get(User,id)
             ```
         In ordinary methods, session will return None. You can get it through:
             ```Python
-            async with db:
+            async with db():
                 db.session.get(User,id)
             ```
         """
-        return self._session_context_var.get() if self._session_enter_count > 0 else None
+        return self._session_context_var.get()
 
-    async def __aenter__(self):
-        with self._session_lock:
-            session = self.session
-            if session is None:
-                session = self.session_maker()
-                self._session_context_var_token = self._session_context_var.set(session)
-            self._session_enter_count += 1
-        return session
-
-    async def __aexit__(self, exc_type, exc_value, traceback):
-        with self._session_lock:
-            self._session_enter_count -= 1
-            if self._session_enter_count <= 0:
-                session = self._session_context_var.get()
-                if session is not None:
-                    await session.close()
-                self._session_context_var.reset(self._session_context_var_token)
+    def __call__(self):
+        return AsyncSessionContextVarManager(self)
 
     @classmethod
     def create(cls, url: str, *, session_options: Mapping[str, Any] = None, **kwargs) -> "AsyncDatabase":
@@ -428,8 +410,6 @@ def get_user(session:Session,id:int):
             APIs which will be properly adapted to the greenlet context.
         """
         need_close = False
-        if executor is None and is_session:
-            executor = self.session
         if executor is None or not isinstance(executor, (AsyncSession, AsyncConnection)):
             if is_session:
                 executor = self.session
@@ -456,38 +436,15 @@ def __init__(self, engine: Engine, **session_options):
         self.engine: Engine = engine
         session_options.setdefault("class_", Session)
         self.session_maker: Callable[..., Session] = sessionmaker(self.engine, **session_options)
-        self._session_lock = Lock()
-        self._session_enter_count = 0
         self._session_context_var: ContextVar[Optional[Session]] = ContextVar("_session_context_var", default=None)
 
     @property
     def session(self) -> Optional[Session]:
         """Return an instance of Session local to the current context."""
-        return self._session_context_var.get() if self._session_enter_count > 0 else None
+        return self._session_context_var.get()
 
-    def __enter__(self):
-        with self._session_lock:
-            session = self.session
-            if session is None:
-                session = self.session_maker()
-                self._session_context_var_token = self._session_context_var.set(session)
-            self._session_enter_count += 1
-        return session
-
-    def __exit__(self, exc_type, exc_value, traceback):
-        with self._session_lock:
-            self._session_enter_count -= 1
-            if self._session_enter_count <= 0:
-                session = self._session_context_var.get()
-                if session is not None:
-                    session.close()
-                self._session_context_var.reset(self._session_context_var_token)
-
-    async def __aenter__(self):
-        return self.__enter__()
-
-    async def __aexit__(self, exc_type, exc_value, traceback):
-        return self.__exit__(exc_type, exc_value, traceback)
+    def __call__(self):
+        return SessionContextVarManager(self)
 
     @classmethod
     def create(cls, url: str, *, session_options: Optional[Mapping[str, Any]] = None, **kwargs) -> "Database":
@@ -514,8 +471,6 @@ def execute(
         **kw: Any,
     ) -> Union[Result, _T]:
         need_close = False
-        if executor is None and is_session:
-            executor = self.session
         if executor is None or not isinstance(executor, (Session, Connection)):
             need_close = True
             if is_session:
@@ -688,3 +643,45 @@ async def __aenter__(self):
     async def __aexit__(self, exc_type, exc_val, exc_tb):
         if self.need_close:
             await self.executor.close()
+
+
+class AsyncSessionContextVarManager:
+    def __init__(self, db: AsyncDatabase):
+        self.db = db
+        self.token = None
+
+    async def __aenter__(self):
+        session = self.db.session_maker()
+        self.token = self.db._session_context_var.set(session)
+        return session
+
+    async def __aexit__(self, exc_type, exc_value, traceback):
+        session = self.db._session_context_var.get()
+        if exc_type is not None:
+            await session.rollback()
+        await session.close()
+        self.db._session_context_var.reset(self.token)
+
+
+class SessionContextVarManager:
+    def __init__(self, db: Database):
+        self.db = db
+        self.token = None
+
+    def __enter__(self):
+        session = self.db.session_maker()
+        self.token = self.db._session_context_var.set(session)
+        return session
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        session = self.db._session_context_var.get()
+        if exc_type is not None:
+            session.rollback()
+        session.close()
+        self.db._session_context_var.reset(self.token)
+
+    async def __aenter__(self):
+        return self.__enter__()
+
+    async def __aexit__(self, exc_type, exc_value, traceback):
+        self.__exit__(exc_type, exc_value, traceback)
diff --git a/tests/test_AbcAsyncDatabase.py b/tests/test_AbcAsyncDatabase.py
@@ -127,7 +127,7 @@ def get_user(session: Session, user_id: int):
 
 
 async def test_async_session_context_var(db, fake_users):
-    async with db:
+    async with db():
         # test db function
         user = await db.async_get(User, 1)
         assert user.id == 1
diff --git a/tests/test_AsyncDatabase.py b/tests/test_AsyncDatabase.py
@@ -175,16 +175,16 @@ async def test_sqlmodel_session(fake_users):
 
 
 async def test_async_session_context_var(fake_users):
-    async with db as session:
+    async with db() as session:
         # test enter return session
         user = await session.get(User, 1)
         assert user.id == 1
 
         # test nested session
-        async with db as session2:
+        async with db() as session2:
             user = await session2.get(User, 1)
             assert user.id == 1
-            assert session is session2
+            assert session is not session2
         # test db.session
         user = await db.session.get(User, 1)
         assert user.id == 1
diff --git a/tests/test_Database.py b/tests/test_Database.py
@@ -173,16 +173,16 @@ def test_sqlmodel_session(fake_users):
 
 
 def test_session_context_var(fake_users):
-    with db as session:
+    with db() as session:
         # test enter return session
         user = session.get(User, 1)
         assert user.id == 1
 
         # test nested session
-        with db as session2:
+        with db() as session2:
             user = session2.get(User, 1)
             assert user.id == 1
-            assert session is session2
+            assert session is not session2
 
         # test db.session
         user = db.session.get(User, 1)
