"""Async SQLAlchemy engine and session management.""" from collections.abc import AsyncGenerator from pathlib import Path from sqlalchemy import event from sqlalchemy.ext.asyncio import ( AsyncSession, async_sessionmaker, create_async_engine, ) from sqlalchemy.orm import DeclarativeBase from app.config import get_settings settings = get_settings() _is_sqlite = settings.database_url.startswith("sqlite") # Ensure the data directory exists for SQLite if _is_sqlite: db_path = settings.database_url.split("///")[-1] Path(db_path).parent.mkdir(parents=True, exist_ok=True) engine = create_async_engine( settings.database_url, echo=False, future=True, # Wait for a competing writer instead of failing instantly. Background # tasks (such as conversation summarization) can briefly hold a write # lock, and Windows enforces SQLite locking more strictly than Linux. connect_args={"timeout": 30} if _is_sqlite else {}, ) if _is_sqlite: @event.listens_for(engine.sync_engine, "connect") def _sqlite_pragmas(dbapi_connection, connection_record) -> None: cursor = dbapi_connection.cursor() cursor.execute("PRAGMA busy_timeout=30000") cursor.execute("PRAGMA foreign_keys=ON") cursor.close() AsyncSessionLocal = async_sessionmaker( bind=engine, class_=AsyncSession, expire_on_commit=False, autoflush=False, ) class Base(DeclarativeBase): """Declarative base for all ORM models.""" async def get_db() -> AsyncGenerator[AsyncSession, None]: """FastAPI dependency that yields a database session.""" async with AsyncSessionLocal() as session: yield session async def init_db() -> None: """Create tables on startup (a migration tool can replace this later).""" from app import models # noqa: F401 (register models with the metadata) async with engine.begin() as conn: await conn.run_sync(Base.metadata.create_all)