Make session scoped to threads and asyncio tasks

Author: febus982

sqlalchemy_bind_manager/_session_handler.py

@@ -1,7 +1,6 @@
 import asyncio
 from contextlib import asynccontextmanager, contextmanager
 from typing import AsyncIterator, Iterator
-from uuid import uuid4
 
 from sqlalchemy.ext.asyncio import (
     AsyncSession,
@@ -17,91 +16,87 @@
 
 
 class SessionHandler:
-    _session_class: scoped_session
-    session: Session
+    scoped_session: scoped_session
 
     def __init__(self, bind: SQLAlchemyBind):
         if not isinstance(bind, SQLAlchemyBind):
             raise UnsupportedBind("Bind is not an instance of SQLAlchemyBind")
         else:
-            u = uuid4()
-            self._session_class = scoped_session(
-                bind.session_class, scopefunc=lambda: str(u)
-            )
-            self.session = self._session_class()
+            self.scoped_session = scoped_session(bind.session_class)
 
     def __del__(self):
-        if getattr(self, "_session_class", None):
-            self._session_class.remove()
+        if getattr(self, "scoped_session", None):
+            self.scoped_session.remove()
 
     @contextmanager
     def get_session(self, read_only: bool = False) -> Iterator[Session]:
+        session = self.scoped_session()
         try:
-            self.session.begin()
-            yield self.session
+            session.begin()
+            yield session
             if not read_only:
-                self.commit()
+                self.commit(session)
         finally:
-            self.session.close()
+            session.close()
 
-    def commit(self) -> None:
+    def commit(self, session: Session) -> None:
         """Commits the session and handles rollback on errors.
 
+        :param session: The session object.
+        :type session: Session
         :raises Exception: Any error is re-raised after the rollback.
         """
         try:
-            self.session.commit()
+            session.commit()
         except:
-            self.session.rollback()
+            session.rollback()
             raise
 
 
 class AsyncSessionHandler:
-    _session_class: async_scoped_session
-    session: AsyncSession
+    scoped_session: async_scoped_session
 
     def __init__(self, bind: SQLAlchemyAsyncBind):
         if not isinstance(bind, SQLAlchemyAsyncBind):
             raise UnsupportedBind("Bind is not an instance of SQLAlchemyAsyncBind")
         else:
-            u = uuid4()
-            self._session_class = async_scoped_session(
-                bind.session_class, scopefunc=lambda: str(u)
+            self.scoped_session = async_scoped_session(
+                bind.session_class, asyncio.current_task
             )
-            self.session = self._session_class()
 
     def __del__(self):
-        if not getattr(self, "_session_class", None):
+        if not getattr(self, "scoped_session", None):
             return
 
         try:
             loop = asyncio.get_event_loop()
             if loop.is_running():
-                loop.create_task(self._session_class.remove())
+                loop.create_task(self.scoped_session.remove())
             else:
-                loop.run_until_complete(self._session_class.remove())
+                loop.run_until_complete(self.scoped_session.remove())
         except RuntimeError:
-            asyncio.run(self._session_class.remove())
+            asyncio.run(self.scoped_session.remove())
 
     @asynccontextmanager
     async def get_session(self, read_only: bool = False) -> AsyncIterator[AsyncSession]:
+        session = self.scoped_session()
         try:
-            await self.session.begin()
-            yield self.session
+            await session.begin()
+            yield session
             if not read_only:
-                await self.commit()
+                await self.commit(session)
         finally:
-            await self.session.close()
+            await session.close()
 
-    async def commit(self) -> None:
+    async def commit(self, session: AsyncSession) -> None:
         """Commits the session and handles rollback on errors.
 
         :param session: The session object.
-        :type session: Session
+        :type session: AsyncSession
         :raises Exception: Any error is re-raised after the rollback.
         """
         try:
-            await self.session.commit()
+            await session.commit()
         except:
-            await self.session.rollback()
+            await session.rollback()
             raise

sqlalchemy_bind_manager/_unit_of_work/__init__.py

@@ -16,36 +16,36 @@
 
 
 class UnitOfWork:
-    _transaction_handler: SessionHandler
+    _session_handler: SessionHandler
 
     def __init__(
         self, bind: SQLAlchemyBind, repositories: Iterable[Type[SQLAlchemyRepository]]
     ) -> None:
         super().__init__()
-        self._transaction_handler = SessionHandler(bind)
+        self._session_handler = SessionHandler(bind)
         for r in repositories:
-            setattr(self, r.__name__, r(session=self._transaction_handler.session))
+            setattr(self, r.__name__, r(session=self._session_handler.scoped_session()))
 
     @contextmanager
     def transaction(self, read_only: bool = False) -> Iterator[Session]:
-        with self._transaction_handler.get_session(read_only=read_only) as _s:
+        with self._session_handler.get_session(read_only=read_only) as _s:
             yield _s
 
 
 class AsyncUnitOfWork:
-    _transaction_handler: AsyncSessionHandler
+    _session_handler: AsyncSessionHandler
 
     def __init__(
         self,
         bind: SQLAlchemyAsyncBind,
         repositories: Iterable[Type[SQLAlchemyAsyncRepository]],
     ) -> None:
         super().__init__()
-        self._transaction_handler = AsyncSessionHandler(bind)
+        self._session_handler = AsyncSessionHandler(bind)
         for r in repositories:
-            setattr(self, r.__name__, r(session=self._transaction_handler.session))
+            setattr(self, r.__name__, r(session=self._session_handler.scoped_session()))
 
     @asynccontextmanager
     async def transaction(self, read_only: bool = False) -> AsyncIterator[AsyncSession]:
-        async with self._transaction_handler.get_session(read_only=read_only) as _s:
+        async with self._session_handler.get_session(read_only=read_only) as _s:
             yield _s

tests/session_handler/test_session_lifecycle.py

@@ -1,47 +1,50 @@
+import asyncio
+from multiprocessing.pool import ThreadPool
+from time import sleep
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
-from sqlalchemy.ext.asyncio import async_scoped_session
-from sqlalchemy.orm import scoped_session
+from sqlalchemy.ext.asyncio import AsyncSession, async_scoped_session
+from sqlalchemy.orm import Session, scoped_session
 
 from sqlalchemy_bind_manager._bind_manager import SQLAlchemyAsyncBind
 
 
 async def test_session_is_removed_on_cleanup(session_handler_class, sa_bind):
-    uow = session_handler_class(sa_bind)
-    original_session_remove = uow._session_class.remove
+    sh = session_handler_class(sa_bind)
+    original_session_remove = sh.scoped_session.remove
 
     with patch.object(
-        uow._session_class,
+        sh.scoped_session,
         "remove",
         wraps=original_session_remove,
     ) as mocked_remove:
         # This should trigger the garbage collector and close the session
-        uow = None
+        sh = None
 
     mocked_remove.assert_called_once()
 
 
-def test_session_is_removed_on_cleanup_even_if_loop_is_not_running(
+async def test_session_is_removed_on_cleanup_even_if_loop_is_not_running(
     session_handler_class, sa_bind
 ):
     # This test makes sense only for async implementation
     if not isinstance(sa_bind, SQLAlchemyAsyncBind):
         return
 
     # Running the test without a loop will trigger the loop creation
-    uow = session_handler_class(sa_bind)
-    original_session_remove = uow._session_class.remove
+    sh = session_handler_class(sa_bind)
+    original_session_remove = sh.scoped_session.remove
 
     with patch.object(
-        uow._session_class,
+        sh.scoped_session,
         "remove",
         wraps=original_session_remove,
     ) as mocked_close, patch(
         "asyncio.get_event_loop", side_effect=RuntimeError()
     ) as mocked_get_event_loop:
         # This should trigger the garbage collector and close the session
-        uow = None
+        sh = None
 
     mocked_get_event_loop.assert_called_once()
     mocked_close.assert_called_once()
@@ -55,7 +58,7 @@ async def test_commit_is_called_only_if_not_read_only(
     sa_bind,
     sync_async_cm_wrapper,
 ):
-    uow = session_handler_class(sa_bind)
+    sh = session_handler_class(sa_bind)
 
     # Populate a database entry to be used for tests
     model1 = model_class(
@@ -64,13 +67,13 @@ async def test_commit_is_called_only_if_not_read_only(
 
     with patch.object(
         session_handler_class, "commit", return_value=None
-    ) as mocked_uow_commit:
+    ) as mocked_sh_commit:
         async with sync_async_cm_wrapper(
-            uow.get_session(read_only=read_only_flag)
+            sh.get_session(read_only=read_only_flag)
         ) as _session:
             _session.add(model1)
 
-    assert mocked_uow_commit.call_count == int(not read_only_flag)
+    assert mocked_sh_commit.call_count == int(not read_only_flag)
 
 
 @pytest.mark.parametrize("commit_fails", [True, False])
@@ -80,23 +83,88 @@ async def test_rollback_is_called_if_commit_fails(
     sa_bind,
     sync_async_wrapper,
 ):
-    uow = session_handler_class(sa_bind)
+    sh = session_handler_class(sa_bind)
 
     failure_exception = Exception("Some Error")
     mocked_session = (
         AsyncMock(spec=async_scoped_session)
         if isinstance(sa_bind, SQLAlchemyAsyncBind)
         else MagicMock(spec=scoped_session)
     )
-    uow.session = mocked_session
     if commit_fails:
         mocked_session.commit.side_effect = failure_exception
 
     try:
-        await sync_async_wrapper(uow.commit())
+        await sync_async_wrapper(sh.commit(mocked_session))
     except Exception as e:
         assert commit_fails is True
         assert e == failure_exception
 
     assert mocked_session.commit.call_count == 1
     assert mocked_session.rollback.call_count == int(commit_fails)
+
+
+async def test_session_is_different_on_different_asyncio_tasks(
+    session_handler_class, sa_bind
+):
+    # This test makes sense only for async implementation
+    if not isinstance(sa_bind, SQLAlchemyAsyncBind):
+        return
+
+    # Running the test without a loop will trigger the loop creation
+    sh = session_handler_class(sa_bind)
+
+    s1 = sh.scoped_session()
+    s2 = sh.scoped_session()
+
+    assert isinstance(s1, AsyncSession)
+    assert isinstance(s2, AsyncSession)
+    assert s1 is s2
+
+    async def _get_sh_session():
+        return sh.scoped_session()
+
+    s = await asyncio.gather(
+        _get_sh_session(),
+        _get_sh_session(),
+    )
+
+    assert isinstance(s[0], AsyncSession)
+    assert isinstance(s[1], AsyncSession)
+    assert s[0] is not s[1]
+
+
+async def test_session_is_different_on_different_threads(
+    session_handler_class, sa_bind
+):
+    # This test makes sense only for sync implementation
+    if isinstance(sa_bind, SQLAlchemyAsyncBind):
+        return
+
+    # Running the test without a loop will trigger the loop creation
+    sh = session_handler_class(sa_bind)
+
+    s1 = sh.scoped_session()
+    s2 = sh.scoped_session()
+
+    assert isinstance(s1, Session)
+    assert isinstance(s2, Session)
+    assert s1 is s2
+
+    def _get_session():
+        # This sleep is to make sure the task doesn't
+        # resolve immediately and multiple instances
+        # end up in different threads
+        sleep(1)
+        return sh.scoped_session()
+
+    with ThreadPool() as pool:
+        s3_task = pool.apply_async(_get_session)
+        s4_task = pool.apply_async(_get_session)
+
+        s3 = s3_task.get()
+        s4 = s4_task.get()
+
+    assert isinstance(s3, Session)
+    assert isinstance(s4, Session)
+    assert s3 is not s4

tests/unit_of_work/test_lifecycle.py

@@ -17,4 +17,6 @@ class ChildRepoClass(repository_class):
         repo = getattr(uow, repo_class.__name__)
         assert not hasattr(repo, "_session_handler")
         assert hasattr(repo, "_external_session")
-        assert getattr(repo, "_external_session") is uow._transaction_handler.session
+        assert (
+            getattr(repo, "_external_session") is uow._session_handler.scoped_session()
+        )