spb/uqo-chat

1"""Async engine / session factory."""23from __future__ import annotations45from collections.abc import AsyncIterator67from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine89from app.core.config import get_settings10from app.models import Base1112_settings = get_settings()13engine = create_async_engine(14    _settings.database_url,15    echo=False,16    pool_pre_ping=True,17    connect_args={"timeout": 30} if _settings.database_url.startswith("sqlite") else {},18)19SessionLocal = async_sessionmaker(engine, expire_on_commit=False, class_=AsyncSession)202122async def init_db() -> None:23    async with engine.begin() as conn:24        if _settings.database_url.startswith("sqlite"):25            await conn.exec_driver_sql("PRAGMA journal_mode=WAL")26            await conn.exec_driver_sql("PRAGMA foreign_keys=ON")27        await conn.run_sync(Base.metadata.create_all)28        await conn.run_sync(_add_missing_columns)293031def _add_missing_columns(sync_conn) -> None:  # noqa: ANN00132    """Poor man's migration: ALTER TABLE ADD COLUMN for new nullable columns (SQLite/Postgres)."""33    from sqlalchemy import inspect, text3435    insp = inspect(sync_conn)36    for table in Base.metadata.sorted_tables:37        if not insp.has_table(table.name):38            continue39        existing = {c["name"] for c in insp.get_columns(table.name)}40        for col in table.columns:41            if col.name in existing:42                continue43            ctype = col.type.compile(sync_conn.dialect)44            sync_conn.execute(text(f'ALTER TABLE {table.name} ADD COLUMN "{col.name}" {ctype}'))454647async def get_session() -> AsyncIterator[AsyncSession]:48    async with SessionLocal() as session:49        yield session50