| """Database session and init helpers.""" |
|
|
| from __future__ import annotations |
|
|
| from collections.abc import AsyncGenerator |
|
|
| from sqlalchemy.ext.asyncio import ( |
| AsyncEngine, |
| AsyncSession, |
| async_sessionmaker, |
| create_async_engine, |
| ) |
|
|
| from app.config import get_settings |
| from app.models import Base |
|
|
| engine: AsyncEngine | None = None |
| SessionLocal: async_sessionmaker[AsyncSession] | None = None |
|
|
|
|
| def configure_engine(url: str | None = None) -> None: |
| """ |
| Create the async engine. |
| |
| Call dispose_engine() first when replacing an existing engine. |
| """ |
| global engine, SessionLocal |
| if engine is not None: |
| raise RuntimeError( |
| "Engine already configured; await dispose_engine() first" |
| ) |
| settings = get_settings() |
| db_url = url or settings.database_url |
| engine = create_async_engine(db_url, echo=False) |
| SessionLocal = async_sessionmaker( |
| engine, |
| class_=AsyncSession, |
| expire_on_commit=False, |
| ) |
|
|
|
|
| def get_session_factory() -> async_sessionmaker[AsyncSession]: |
| """Return the active session factory.""" |
| if SessionLocal is None: |
| configure_engine() |
| assert SessionLocal is not None |
| return SessionLocal |
|
|
|
|
| configure_engine() |
|
|
|
|
| async def dispose_engine() -> None: |
| """Dispose current engine (release file locks).""" |
| global engine, SessionLocal |
| if engine is not None: |
| await engine.dispose() |
| engine = None |
| SessionLocal = None |
|
|
|
|
| async def init_db() -> None: |
| """Create tables if they do not exist.""" |
| if engine is None: |
| configure_engine() |
| assert engine is not None |
| async with engine.begin() as conn: |
| await conn.run_sync(Base.metadata.create_all) |
|
|
|
|
| async def get_session() -> AsyncGenerator[AsyncSession, None]: |
| """Yield an async DB session.""" |
| factory = get_session_factory() |
| async with factory() as session: |
| yield session |
|
|
