You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
34 lines
1014 B
Python
34 lines
1014 B
Python
"""Async session factory and context manager."""
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import AsyncGenerator
|
|
from contextlib import asynccontextmanager
|
|
|
|
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
|
|
|
|
from libs.db.engine import get_engine
|
|
|
|
_session_factory: async_sessionmaker[AsyncSession] | None = None
|
|
|
|
|
|
def get_session_factory(dsn: str | None = None) -> async_sessionmaker[AsyncSession]:
|
|
global _session_factory
|
|
if _session_factory is None:
|
|
engine = get_engine(dsn)
|
|
_session_factory = async_sessionmaker(
|
|
engine, expire_on_commit=False, class_=AsyncSession
|
|
)
|
|
return _session_factory
|
|
|
|
|
|
@asynccontextmanager
|
|
async def get_session(dsn: str | None = None) -> AsyncGenerator[AsyncSession, None]:
|
|
factory = get_session_factory(dsn)
|
|
async with factory() as session:
|
|
try:
|
|
yield session
|
|
await session.commit()
|
|
except Exception:
|
|
await session.rollback()
|
|
raise
|