app/db/database.py

"""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_<table>_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()