"""Async SQLAlchemy engine, session factory, and database initialisation.""" import logging from collections.abc import AsyncGenerator from typing import Any from sqlalchemy.ext.asyncio import ( AsyncEngine, AsyncSession, async_sessionmaker, create_async_engine, ) from sqlalchemy.orm import DeclarativeBase # starlette.exceptions.HTTPException is the base class of fastapi.HTTPException; # catching the Starlette one filters out both routine 4xx flows that propagate # through yield-based dependencies in FastAPI >= 0.106. from starlette.exceptions import HTTPException as StarletteHTTPException from app.config import settings logger = logging.getLogger(__name__) class Base(DeclarativeBase): """Declarative base shared by all ORM models.""" _engine: AsyncEngine | None = None _session_factory: async_sessionmaker[AsyncSession] | None = None def is_sqlite_database() -> bool: """True when the configured async URL targets SQLite.""" return "sqlite" in (settings.database_url or "").lower() def parallel_section_writes_safe() -> bool: """Whether concurrent section DB writers are safe for the active backend.""" if is_sqlite_database() and not settings.allow_sqlite_parallel_sections: return False return True def multi_section_parallel_enabled() -> bool: """Parallel multi-section jobs (in-process or Temporal) when async pipeline is on.""" return bool(settings.enable_async_pipeline and parallel_section_writes_safe()) def get_engine() -> AsyncEngine: """Return (or lazily create) the async SQLAlchemy engine. SQLite — uses ``check_same_thread=False`` (required for async). PostgreSQL / other — enables ``pool_pre_ping`` so stale connections are transparently recycled rather than causing 500 errors under load. Returns: The singleton ``AsyncEngine`` instance. """ global _engine if _engine is None: is_sqlite = is_sqlite_database() kwargs: dict[str, Any] = { "echo": settings.dev_mode, } if is_sqlite: kwargs["connect_args"] = {"check_same_thread": False} else: # For PostgreSQL / MySQL: recycle stale connections and limit pool # size so we don't exhaust DB connection slots under high concurrency. kwargs["pool_pre_ping"] = True kwargs["pool_size"] = 10 kwargs["max_overflow"] = 20 kwargs["pool_recycle"] = 1800 # recycle connections every 30 min _engine = create_async_engine(settings.database_url, **kwargs) return _engine def get_session_factory() -> async_sessionmaker[AsyncSession]: """Return (or lazily create) the async session factory. Returns: An ``async_sessionmaker`` bound to the singleton engine. """ global _session_factory if _session_factory is None: _session_factory = async_sessionmaker( get_engine(), expire_on_commit=False, ) return _session_factory async def get_db() -> AsyncGenerator[AsyncSession, None]: """FastAPI dependency that yields a transactional database session. Commits on success, rolls back on any exception. Note on logging: from FastAPI 0.106 onwards, exceptions raised in path operations — including ``HTTPException`` for routine 404 / 409 / 413 / 422 responses — propagate to yield-based dependencies. We therefore handle ``HTTPException`` separately and roll back quietly; logging it with ``logger.exception`` would emit an ERROR-level stack trace for every routine 4xx response and bury genuine database errors in the noise. Yields: ``AsyncSession`` scoped to a single request. """ factory = get_session_factory() async with factory() as session: try: yield session await session.commit() except StarletteHTTPException: # Routine HTTP responses (4xx) — not a database failure. Roll back # any open transaction so the session is clean, but stay quiet so # the logs don't drown in stack traces for every NotFound / Conflict. try: await session.rollback() except Exception: logger.exception("Rollback failed while propagating HTTPException") raise except Exception as _exc: logger.exception("DB session error — rolling back: %s", _exc) try: await session.rollback() except Exception as _rb_exc: logger.error("Rollback itself failed: %s", _rb_exc) raise def _sqlite_add_column_if_missing(connection: object, table: str, column: str, ddl_suffix: str) -> None: """Best-effort SQLite ALTER for deployments that created tables before new columns existed.""" from sqlalchemy import inspect, text insp = inspect(connection) if not insp.has_table(table): return cols = {c["name"] for c in insp.get_columns(table)} if column in cols: return connection.execute(text(f"ALTER TABLE {table} ADD COLUMN {ddl_suffix}")) def _postgres_add_column_if_missing( connection: object, table: str, column: str, ddl_suffix: str, ) -> None: """Best-effort PostgreSQL ALTER for nullable columns added after first deploy.""" from sqlalchemy import inspect, text insp = inspect(connection) if not insp.has_table(table): return cols = {c["name"] for c in insp.get_columns(table)} if column in cols: return connection.execute(text(f"ALTER TABLE {table} ADD COLUMN IF NOT EXISTS {ddl_suffix}")) async def migrate_schema() -> None: """Apply lightweight schema upgrades when models gain nullable columns.""" engine = get_engine() def _upgrade(sync_conn: object) -> None: if is_sqlite_database(): _sqlite_add_column_if_missing( sync_conn, "documents", "survey_level", "survey_level INTEGER" ) _sqlite_add_column_if_missing( sync_conn, "reports", "survey_level", "survey_level INTEGER" ) _sqlite_add_column_if_missing( sync_conn, "reports", "generation_started_at", "generation_started_at DATETIME", ) else: _postgres_add_column_if_missing( sync_conn, "reports", "generation_started_at", "generation_started_at TIMESTAMP WITH TIME ZONE", ) async with engine.begin() as conn: await conn.run_sync(_upgrade) async def init_db() -> None: """Create all tables if they do not already exist (idempotent). Example:: await init_db() """ from app.db import models as _models # noqa: F401 — registers ORM metadata engine = get_engine() async with engine.begin() as conn: await conn.run_sync(Base.metadata.create_all) await migrate_schema()