Merge pull request #35 from febus982/multithread_session

Make session handler thread safe and async task safe

Author: febus982

docs/lifecycle.md

@@ -17,70 +17,37 @@ From SQLAlchemy documentation:
 recommends we create `Session` object at the beginning of a logical operation where
 database access is potentially anticipated.
 
-The repository is not built for parallel execution, it keeps a `Session` object scoped to
-its lifecycle to avoid unnecessary queries, and executes a transaction for each _operation_
-to maintain isolation. This means you can create a repository object almost whenever you want,
-as long as you don't run parallel operations.
-
-The session in the repository is not thread safe and [is not safe on concurrent asyncio tasks](https://docs.sqlalchemy.org/en/20/orm/session_basics.html#is-the-session-thread-safe-is-asyncsession-safe-to-share-in-concurrent-tasks) 
-therefore the repository has the same limitations.
-
-This means:
-
-* Do not assign a repository object to a global variable
-  (Check the [Notes on multithreaded applications](/manager/session/#note-on-multithreaded-applications))
-* Do not share a repository instance in multiple threads or processes (e.g. using `asyncio.to_thread`).
-* Do not use the same repository in concurrent asyncio task (e.g. using `asyncio.gather`)
-
-Even using multiple repository instances will work fine, however as they will have completely
-different sessions, it's likely that the second repository will fire additional SELECT queries
-to get the state of the object prior to saving it.
-
-/// details | Example
-```python
-from sqlalchemy import String
-from sqlalchemy.orm import Mapped, mapped_column
-from sqlalchemy_bind_manager import SQLAlchemyBindManager ,SQLAlchemyConfig
-from sqlalchemy_bind_manager.repository import SQLAlchemyRepository
-
-config = SQLAlchemyConfig(
-    engine_url="sqlite:///./sqlite.db",
-    engine_options=dict(connect_args={"check_same_thread": False}, echo=True),
-    session_options=dict(expire_on_commit=False),
-)
-
-sa_manager = SQLAlchemyBindManager(config={})
-
-class MyModel(sa_manager.get_bind().model_declarative_base):
-    id: Mapped[int] = mapped_column(primary_key=True)
-    name: Mapped[str] = mapped_column(String(30))
-
-def update_my_model():
-    # Create 2 instances of the same repository
-    repo = SQLAlchemyRepository(sa_manager.get_bind(), model_class=MyModel)
-    repo2 = SQLAlchemyRepository(sa_manager.get_bind(), model_class=MyModel)
-    
-    o = repo.get(1)
-    o.name = "John"
-    
-    repo2.save(o)
-
-update_my_model()
-```
-///
+The repository keeps a `Session` object scoped to its lifecycle to avoid unnecessary queries,
+and executes a transaction for each _operation_ to maintain isolation. This means you can create
+a repository object almost whenever you want, as long as you don't run parallel operations.
+
+The repository is safe in multithreaded applications and in concurrent asyncio tasks, this means
+that potentially you can save it in a global variable, and it will have a different `Session`
+in each thread or asyncio task.
+
+Even if the repository can be used with concurrency or parallelism, remember SQLAlchemy models
+belong to a single `Session`, so sharing the same models in multiple threads or asyncio tasks
+will cause problems.
+
+What you can do is:
 
-The recommendation is of course to try to write your application in a way you cna use
-a single repository instance, where possible.
+* Save the repositories in global variables and start a thread / asyncio task to handle
+  a scoped request (e.g. one thread per HTTP request)
+
+What you cannot do is:
+
+* Get a list of models
+* Save the models using `save()` in parallel threads / tasks (each task will have a different session)
+
+/// tip | The recommendation is of course to try to use a single repository instance, where possible.
 
 For example a strategy similar to this would be optimal, if possible:
 
 * Create repositories
 * Retrieve all the models you need
-* Do the changes you need, as per business logic
+* Do the changes you need, as per business logic, eventually using multiple threads / tasks
 * Save all the changed models as needed
 
-/// tip | Using multiple repository instances is the only way to safely use concurrent asyncio tasks
-
 ///
 
 ### Unit of work

docs/repository/repository.md docs/repository/usage.mddocs/repository/repository.md renamed to docs/repository/usage.md

@@ -45,7 +45,7 @@ nav:
       - Session usage: manager/session.md
       - Alembic integration: manager/alembic.md
   - Repository:
-      - Repository usage: repository/repository.md
+      - Repository usage: repository/usage.md
       - Unit of work: repository/uow.md
   - Components life cycle: lifecycle.md
 

mkdocs.yml

@@ -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/_session_handler.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