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) 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()