Implement session_factory context manager

Author: teamhide

README.md

@@ -61,6 +61,27 @@ async def create_user(self):
 
 Do not use explicit `commit()`. `Transactional` class automatically do.
 
+### Query with asyncio.gather()
+When executing queries concurrently through `asyncio.gather()`, you must use the `session_factory` context manager rather than the globally used session.
+
+```python
+from core.db import session_factory
+
+
+async def get_by_id(self, *, user_id) -> User:
+    stmt = select(User)
+    async with session_factory() as read_session:
+        return await read_session.execute(query).scalars().first()
+
+
+async def main() -> None:
+    user_1, user_2 = await asyncio.gather(
+        get_by_id(user_id=1),
+        get_by_id(user_id=2),
+    )
+```
+If you do not use a database connection like `session.add()`, it is recommended to use a globally provided session.
+
 ### Multiple databases
 
 Go to `core/config.py` and edit `WRITER_DB_URL` and `READER_DB_URL` in the config class.

app/user/adapter/output/persistence/sqlalchemy/user.py

@@ -1,8 +1,8 @@
-from sqlalchemy import select, or_, and_
+from sqlalchemy import and_, or_, select
 
 from app.user.domain.entity.user import User
 from app.user.domain.repository.user import UserRepo
-from core.db.session import session
+from core.db.session import session, session_factory
 
 
 class UserSQLAlchemyRepo(UserRepo):
@@ -21,7 +21,9 @@ async def get_users(
             limit = 12
 
         query = query.limit(limit)
-        result = await session.execute(query)
+        async with session_factory() as read_session:
+            result = await read_session.execute(query)
+
         return result.scalars().all()
 
     async def get_user_by_email_or_nickname(
@@ -30,25 +32,28 @@ async def get_user_by_email_or_nickname(
         email: str,
         nickname: str,
     ) -> User | None:
-        stmt = await session.execute(
-            select(User).where(or_(User.email == email, User.nickname == nickname)),
-        )
-        return stmt.scalars().first()
+        async with session_factory() as read_session:
+            stmt = await read_session.execute(
+                select(User).where(or_(User.email == email, User.nickname == nickname)),
+            )
+            return stmt.scalars().first()
 
     async def get_user_by_id(self, *, user_id: int) -> User | None:
-        query = await session.execute(select(User).where(User.id == user_id))
-        return query.scalars().first()
+        async with session_factory() as read_session:
+            stmt = await read_session.execute(select(User).where(User.id == user_id))
+            return stmt.scalars().first()
 
     async def get_user_by_email_and_password(
         self,
         *,
         email: str,
         password: str,
     ) -> User | None:
-        stmt = await session.execute(
-            select(User).where(and_(User.email == email, password == password))
-        )
-        return stmt.scalars().first()
+        async with session_factory() as read_session:
+            stmt = await read_session.execute(
+                select(User).where(and_(User.email == email, password == password))
+            )
+            return stmt.scalars().first()
 
     async def save(self, *, user: User) -> None:
         session.add(user)

core/db/__init__.py

@@ -1,8 +1,9 @@
-from .session import Base, session
+from .session import Base, session, session_factory
 from .transactional import Transactional
 
 __all__ = [
     "Base",
     "session",
     "Transactional",
+    "session_factory",
 ]

core/db/session.py

@@ -1,13 +1,15 @@
+from contextlib import asynccontextmanager
 from contextvars import ContextVar, Token
 from enum import Enum
+from typing import AsyncGenerator
 
 from sqlalchemy.ext.asyncio import (
     AsyncSession,
     async_scoped_session,
     async_sessionmaker,
     create_async_engine,
 )
-from sqlalchemy.orm import Session, DeclarativeBase
+from sqlalchemy.orm import DeclarativeBase, Session
 from sqlalchemy.sql.expression import Delete, Insert, Update
 
 from core.config import config
@@ -46,16 +48,29 @@ def get_bind(self, mapper=None, clause=None, **kw):
             return engines[EngineType.READER].sync_engine
 
 
-async_session_factory = async_sessionmaker(
+_async_session_factory = async_sessionmaker(
     class_=AsyncSession,
     sync_session_class=RoutingSession,
     expire_on_commit=False,
 )
 session = async_scoped_session(
-    session_factory=async_session_factory,
+    session_factory=_async_session_factory,
     scopefunc=get_session_context,
 )
 
 
 class Base(DeclarativeBase):
     ...
+
+
+@asynccontextmanager
+async def session_factory() -> AsyncGenerator[AsyncSession, None]:
+    _session = async_sessionmaker(
+        class_=AsyncSession,
+        sync_session_class=RoutingSession,
+        expire_on_commit=False,
+    )()
+    try:
+        yield _session
+    finally:
+        await _session.close()