"""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). Enabled by either the async pipeline OR the dedicated ``enable_parallel_section_generation`` flag, and only when concurrent section writes are safe for the active DB backend (SQLite requires ``allow_sqlite_parallel_sections``). This only governs the Temporal workflow request's ``parallel_sections`` field; the default (non-Temporal) job path is unaffected. """ if not parallel_section_writes_safe(): return False return bool( settings.enable_async_pipeline or settings.enable_parallel_section_generation ) def effective_section_concurrency() -> int: """Bounded section concurrency for a multi-section job. Decoupled from ``enable_async_pipeline``: section generation is dominated by I/O-bound LLM calls, so concurrency is safe and worthwhile regardless of the async-pipeline flag. Falls back to 1 (sequential) when concurrent section writes are not safe for the active DB backend. """ requested = max(1, int(getattr(settings, "section_generation_concurrency", 1))) if requested <= 1: return 1 if not parallel_section_writes_safe(): return 1 return requested 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) if is_sqlite: _enable_sqlite_concurrency(_engine) return _engine def _enable_sqlite_concurrency(engine: AsyncEngine) -> None: """Set WAL + busy_timeout on every SQLite connection. WAL lets readers and a single writer proceed concurrently; busy_timeout makes a blocked writer wait (instead of raising 'database is locked') when parallel section jobs commit at the same time. ``synchronous=NORMAL`` is the standard durable-enough pairing with WAL. """ from sqlalchemy import event @event.listens_for(engine.sync_engine, "connect") def _set_sqlite_pragmas(dbapi_conn: Any, _record: Any) -> None: # noqa: ANN401 cursor = dbapi_conn.cursor() try: cursor.execute("PRAGMA journal_mode=WAL") cursor.execute("PRAGMA busy_timeout=10000") cursor.execute("PRAGMA synchronous=NORMAL") finally: cursor.close() 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}")) def _sqlite_drop_not_null(connection: object, table: str, column: str) -> None: """Relax a ``NOT NULL`` constraint on one SQLite column via a table rebuild. SQLite cannot ``ALTER COLUMN``; the standard recipe is to recreate the table with the relaxed definition and copy rows. Idempotent: no-op when the column is already nullable or absent. FK declarations are intentionally dropped on rebuild (we want lenient deletes), and the ``ix__tenant_id`` helper index is recreated. """ from sqlalchemy import text info = connection.execute(text(f"PRAGMA table_info({table})")).fetchall() # type: ignore[attr-defined] if not info: return target = next((row for row in info if row[1] == column), None) if target is None or int(target[3]) == 0: # column missing, or already nullable return col_defs: list[str] = [] for _cid, name, ctype, notnull, dflt, pk in info: piece = f'"{name}" {ctype or ""}'.rstrip() if int(pk): piece += " PRIMARY KEY" if int(notnull) and name != column: piece += " NOT NULL" if dflt is not None: piece += f" DEFAULT {dflt}" col_defs.append(piece) col_names = ", ".join(f'"{row[1]}"' for row in info) tmp = f"{table}__migrate_tmp" connection.execute(text("PRAGMA foreign_keys=OFF")) # type: ignore[attr-defined] connection.execute(text(f'DROP TABLE IF EXISTS "{tmp}"')) # type: ignore[attr-defined] connection.execute(text(f'CREATE TABLE "{tmp}" ({", ".join(col_defs)})')) # type: ignore[attr-defined] connection.execute( # type: ignore[attr-defined] text(f'INSERT INTO "{tmp}" ({col_names}) SELECT {col_names} FROM "{table}"') ) connection.execute(text(f'DROP TABLE "{table}"')) # type: ignore[attr-defined] connection.execute(text(f'ALTER TABLE "{tmp}" RENAME TO "{table}"')) # type: ignore[attr-defined] connection.execute( # type: ignore[attr-defined] text(f'CREATE INDEX IF NOT EXISTS "ix_{table}_tenant_id" ON "{table}" (tenant_id)') ) def _postgres_drop_not_null(connection: object, table: str, column: str) -> None: """Relax a ``NOT NULL`` constraint on PostgreSQL (idempotent).""" from sqlalchemy import inspect, text insp = inspect(connection) if not insp.has_table(table): return col = next((c for c in insp.get_columns(table) if c["name"] == column), None) if col is None or col.get("nullable", True): return connection.execute(text(f'ALTER TABLE {table} ALTER COLUMN {column} DROP NOT NULL')) 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_drop_not_null(sync_conn, "reports", "document_id") _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", ) _sqlite_add_column_if_missing( sync_conn, "reports", "generation_section_total", "generation_section_total INTEGER", ) # SQLite stores StrEnum as VARCHAR; defaulting at ALTER time is # how legacy rows get the new column populated without a backfill # script. _sqlite_add_column_if_missing( sync_conn, "documents", "document_purpose", "document_purpose VARCHAR(32) NOT NULL DEFAULT 'report_source'", ) _sqlite_add_column_if_missing( sync_conn, "report_section_photos", "selected_for_ai", "selected_for_ai BOOLEAN NOT NULL DEFAULT 0", ) _sqlite_add_column_if_missing( sync_conn, "documents", "redaction_strategy", "redaction_strategy VARCHAR(32) NOT NULL DEFAULT 'ai_hybrid'", ) _sqlite_add_column_if_missing( sync_conn, "documents", "redaction_context_json", "redaction_context_json TEXT", ) else: _postgres_drop_not_null(sync_conn, "reports", "document_id") _postgres_add_column_if_missing( sync_conn, "reports", "generation_started_at", "generation_started_at TIMESTAMP WITH TIME ZONE", ) _postgres_add_column_if_missing( sync_conn, "reports", "generation_section_total", "generation_section_total INTEGER", ) _postgres_add_column_if_missing( sync_conn, "documents", "document_purpose", "document_purpose VARCHAR(32) NOT NULL DEFAULT 'report_source'", ) _postgres_add_column_if_missing( sync_conn, "report_section_photos", "selected_for_ai", "selected_for_ai BOOLEAN NOT NULL DEFAULT false", ) _postgres_add_column_if_missing( sync_conn, "documents", "redaction_strategy", "redaction_strategy VARCHAR(32) NOT NULL DEFAULT 'ai_hybrid'", ) _postgres_add_column_if_missing( sync_conn, "documents", "redaction_context_json", "redaction_context_json TEXT", ) 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()