gh-christianlouis-docuelevate/app/database.py

# app/database.py

import logging
import os
import warnings
from collections.abc import Generator
from pathlib import Path
from typing import Any

from sqlalchemy import create_engine, exc
from sqlalchemy.engine.url import make_url
from sqlalchemy.orm import Session, declarative_base, sessionmaker
from sqlalchemy.pool import NullPool, QueuePool

from app.config import settings

logger = logging.getLogger(__name__)

Base = declarative_base()

# ---------------------------------------------------------------------------
# Engine construction
# ---------------------------------------------------------------------------
DB_URL = settings.database_url
_parsed_url = make_url(DB_URL)

_connect_args: dict[str, Any] = {}
_engine_kwargs: dict[str, Any] = {
    "pool_pre_ping": True,  # detect stale / dropped connections before use
}

if _parsed_url.get_backend_name() == "sqlite":
    # SQLite does not benefit from connection pooling and is prone to
    # QueuePool exhaustion under concurrent access.  NullPool opens a fresh
    # connection for each request and closes it immediately afterwards,
    # completely avoiding the "QueuePool limit reached" TimeoutError.
    _connect_args["check_same_thread"] = False
    _engine_kwargs["poolclass"] = NullPool
else:
    # PostgreSQL / MySQL — use a bounded QueuePool with configurable limits.
    _engine_kwargs["poolclass"] = QueuePool
    _engine_kwargs.update(
        {
            "pool_size": settings.db_pool_size,
            "max_overflow": settings.db_max_overflow,
            "pool_timeout": settings.db_pool_timeout,
            "pool_recycle": settings.db_pool_recycle,
        }
    )

engine = create_engine(DB_URL, connect_args=_connect_args, **_engine_kwargs)
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)


def init_db() -> None:
    """
    Ensures the SQLite database file and its parent directory exist (if using sqlite).
    Then initializes tables and applies any pending Alembic migrations:

    - Fresh/legacy databases (no ``alembic_version`` table): creates all tables via
      ``Base.metadata.create_all()``, then stamps the Alembic version to ``head``.
    - Alembic-tracked databases (``alembic_version`` present): skips ``create_all()``
      and applies pending migrations via ``alembic upgrade head``.  Skipping
      ``create_all()`` prevents an ``OperationalError`` when the ORM model defines a
      table (e.g. ``webhook_configs``) that a pending migration also tries to create.
    """
    # 1. Parse the DB URL to see if it's sqlite
    url = make_url(DB_URL)
    if url.get_backend_name() == "sqlite":
        # 2. Extract the database path from the URL
        database_path = url.database  # e.g. "/workdir/db/database.db" or ":memory:"

        if database_path != ":memory:":
            # 3. Ensure directory exists
            db_dir = os.path.dirname(database_path)
            if db_dir and not os.path.exists(db_dir):
                logger.info(f"Creating directory for SQLite DB: {db_dir}")
                os.makedirs(db_dir, exist_ok=True)

            # 4. If the file does not exist, create an empty one
            if not os.path.exists(database_path):
                logger.info(f"Creating new SQLite database file at {database_path}")
                open(database_path, "a").close()

    # 5. Create tables only for fresh/legacy databases not yet tracked by Alembic.
    #    For Alembic-tracked databases, skip create_all to avoid conflicts where
    #    the ORM model would create a table (e.g. webhook_configs) that a pending
    #    Alembic migration also tries to create, causing an OperationalError.
    try:
        from sqlalchemy import inspect

        table_names = inspect(engine).get_table_names()
        if "alembic_version" not in table_names:
            Base.metadata.create_all(bind=engine)
            logger.info("Database initialization complete (tables created if not exist).")

        # 6. Run Alembic migrations (stamps fresh/legacy DBs to head, upgrades tracked DBs)
        _run_alembic_upgrade(engine)
    except exc.SQLAlchemyError as e:
        logger.error(f"Error initializing database: {e}")
        raise


def _run_alembic_upgrade(engine: Any) -> None:
    """Run Alembic migrations programmatically to apply pending schema changes.

    For fresh databases (created via ``Base.metadata.create_all()``), the
    Alembic version is stamped to ``head`` because all tables already exist.
    For existing databases with Alembic tracking, any pending migrations are
    applied via ``alembic upgrade head``.

    Args:
        engine: The SQLAlchemy engine connected to the target database.
    """
    from alembic import command
    from alembic.config import Config
    from sqlalchemy import inspect

    inspector = inspect(engine)
    table_names = inspector.get_table_names()

    # Locate the migrations directory relative to this file
    migrations_dir = str(Path(__file__).resolve().parent.parent / "migrations")

    # Build an Alembic Config that points at our migration scripts.
    # The sqlalchemy.url is intentionally left empty because we pass the
    # live connection via config.attributes["connection"] below.
    alembic_cfg = Config()
    alembic_cfg.set_main_option("script_location", migrations_dir)
    alembic_cfg.set_main_option("sqlalchemy.url", "")

    with engine.begin() as connection:
        alembic_cfg.attributes["connection"] = connection

        if "alembic_version" not in table_names:
            # Fresh database or one that predates Alembic tracking.
            # Base.metadata.create_all() already created everything, so
            # stamp the current version to head (no migrations need to run).
            logger.info("No Alembic version table found — stamping database to latest revision.")
            command.stamp(alembic_cfg, "head")
        else:
            # Existing database with Alembic version tracking — apply pending migrations.
            logger.info("Running pending Alembic migrations…")
            command.upgrade(alembic_cfg, "head")
            logger.info("Alembic migration check complete.")


def _run_schema_migrations(engine: Any) -> None:
    """Apply lightweight schema migrations for columns added after the initial release.

    .. deprecated::
        This function is deprecated and will be removed in a future release.
        All schema migrations are now managed exclusively through Alembic.
        Run ``alembic upgrade head`` (or let ``init_db()`` handle it
        automatically) instead of calling this function directly.

    Each migration is idempotent and safe to run multiple times.
    """
    warnings.warn(
        "_run_schema_migrations() is deprecated. "
        "All schema changes are now managed by Alembic migrations. "
        "Use 'alembic upgrade head' or init_db() instead.",
        DeprecationWarning,
        stacklevel=2,
    )

    from sqlalchemy import inspect, text

    inspector = inspect(engine)

    # Migration: Add 'detail' column to processing_logs (added for verbose worker log output)
    if "processing_logs" in inspector.get_table_names():
        columns = [col["name"] for col in inspector.get_columns("processing_logs")]
        if "detail" not in columns:
            logger.info("Migrating processing_logs: adding 'detail' column")
            with engine.begin() as conn:
                conn.execute(text("ALTER TABLE processing_logs ADD COLUMN detail TEXT"))
            logger.info("Migration complete: 'detail' column added to processing_logs")

    # Migration: Add file path columns to files table
    if "files" in inspector.get_table_names():
        columns = [col["name"] for col in inspector.get_columns("files")]
        if "original_file_path" not in columns:
            logger.info("Migrating files: adding 'original_file_path' column")
            with engine.begin() as conn:
                conn.execute(text("ALTER TABLE files ADD COLUMN original_file_path VARCHAR"))
            logger.info("Migration complete: 'original_file_path' column added to files")

        if "processed_file_path" not in columns:
            logger.info("Migrating files: adding 'processed_file_path' column")
            with engine.begin() as conn:
                conn.execute(text("ALTER TABLE files ADD COLUMN processed_file_path VARCHAR"))
            logger.info("Migration complete: 'processed_file_path' column added to files")

        # Migration: Add deduplication columns to files table
        if "is_duplicate" not in columns:
            logger.info("Migrating files: adding 'is_duplicate' column")
            with engine.begin() as conn:
                conn.execute(text("ALTER TABLE files ADD COLUMN is_duplicate BOOLEAN DEFAULT FALSE NOT NULL"))
            logger.info("Migration complete: 'is_duplicate' column added to files")

        if "duplicate_of_id" not in columns:
            logger.info("Migrating files: adding 'duplicate_of_id' column")
            with engine.begin() as conn:
                conn.execute(text("ALTER TABLE files ADD COLUMN duplicate_of_id INTEGER"))
            logger.info("Migration complete: 'duplicate_of_id' column added to files")

        # Migration: Add search/OCR fields to files table
        if "ocr_text" not in columns:
            logger.info("Migrating files: adding 'ocr_text' column")
            with engine.begin() as conn:
                conn.execute(text("ALTER TABLE files ADD COLUMN ocr_text TEXT"))
            logger.info("Migration complete: 'ocr_text' column added to files")

        if "ai_metadata" not in columns:
            logger.info("Migrating files: adding 'ai_metadata' column")
            with engine.begin() as conn:
                conn.execute(text("ALTER TABLE files ADD COLUMN ai_metadata TEXT"))
            logger.info("Migration complete: 'ai_metadata' column added to files")

        if "document_title" not in columns:
            logger.info("Migrating files: adding 'document_title' column")
            with engine.begin() as conn:
                conn.execute(text("ALTER TABLE files ADD COLUMN document_title VARCHAR"))
            logger.info("Migration complete: 'document_title' column added to files")

        if "ocr_quality_score" not in columns:
            logger.info("Migrating files: adding 'ocr_quality_score' column")
            with engine.begin() as conn:
                conn.execute(text("ALTER TABLE files ADD COLUMN ocr_quality_score INTEGER"))
            logger.info("Migration complete: 'ocr_quality_score' column added to files")

        # Migration: Drop unique index on filehash to allow duplicate records
        try:
            indexes = inspector.get_indexes("files")
            unique_filehash_indexes = [
                index for index in indexes if index.get("unique") and "filehash" in index.get("column_names", [])
            ]
            if unique_filehash_indexes:
                logger.info("Migrating files: dropping unique index on 'filehash'")
                with engine.begin() as conn:
                    preparer = conn.dialect.identifier_preparer
                    for index in unique_filehash_indexes:
                        quoted_idx = preparer.quote(index["name"])
                        conn.execute(text(f"DROP INDEX IF EXISTS {quoted_idx}"))
                logger.info("Migration complete: unique index on 'filehash' removed")
        except Exception as exc:
            logger.warning(f"Skipping filehash unique index drop: {exc}")

    # Migration: Create saved_searches table for user-defined filter combinations
    if "saved_searches" not in inspector.get_table_names():
        logger.info("Migrating: creating 'saved_searches' table")
        with engine.begin() as conn:
            conn.execute(
                text(
                    """
                    CREATE TABLE saved_searches (
                        id INTEGER PRIMARY KEY AUTOINCREMENT,
                        user_id VARCHAR NOT NULL,
                        name VARCHAR NOT NULL,
                        filters TEXT NOT NULL,
                        created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                        updated_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                        UNIQUE (user_id, name)
                    )
                    """
                )
            )
            conn.execute(text("CREATE INDEX IF NOT EXISTS ix_saved_searches_user_id ON saved_searches (user_id)"))
        logger.info("Migration complete: 'saved_searches' table created")

    # Migration: Add performance indexes for common query patterns
    _ensure_indexes(engine, inspector)


def _ensure_indexes(engine: Any, inspector: Any) -> None:
    """
    Create performance indexes for common query patterns.

    Each ``CREATE INDEX IF NOT EXISTS`` is idempotent and safe to run
    on every startup.  The indexes target the columns most frequently
    used in file listing/filtering, status computation and log retrieval.
    """
    from sqlalchemy import text

    _PERF_INDEXES = [
        ("ix_files_created_at", "files", "created_at"),
        ("ix_files_mime_type", "files", "mime_type"),
        ("ix_processing_logs_file_id", "processing_logs", "file_id"),
        ("ix_processing_logs_timestamp", "processing_logs", "timestamp"),
        ("ix_file_processing_steps_status", "file_processing_steps", "status"),
    ]

    table_names = inspector.get_table_names()
    columns_by_table: dict[str, set[str]] = {}
    with engine.begin() as conn:
        preparer = conn.dialect.identifier_preparer
        for idx_name, table, column in _PERF_INDEXES:
            if table in table_names:
                if table not in columns_by_table:
                    columns_by_table[table] = {col["name"] for col in inspector.get_columns(table)}
                if column in columns_by_table[table]:
                    # SECURITY: Quoted identifiers to prevent SQL injection during index creation
                    quoted_idx = preparer.quote(idx_name)
                    quoted_table = preparer.quote(table)
                    quoted_col = preparer.quote(column)
                    conn.execute(text(f"CREATE INDEX IF NOT EXISTS {quoted_idx} ON {quoted_table} ({quoted_col})"))

    logger.info("Performance indexes ensured")


def get_db() -> Generator[Session, None, None]:
    """
    Dependency for FastAPI routes or general DB usage.
    Yields a SQLAlchemy session, and closes it upon exit.
    """
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()