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 app.config import settings

logger = logging.getLogger(__name__)

Base = declarative_base()

# Parse the DATABASE_URL
DB_URL = settings.database_url
engine = create_engine(DB_URL, connect_args={"check_same_thread": False})
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 runs Base.metadata.create_all(bind=engine) to initialize tables and
    applies any pending Alembic migrations.
    """
    # 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. Now create tables if they don't exist yet
    try:
        Base.metadata.create_all(bind=engine)
        logger.info("Database initialization complete (tables created if not exist).")

        # 6. Run Alembic migrations for existing databases
        _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:
                    for index in unique_filehash_indexes:
                        conn.execute(text(f"DROP INDEX IF EXISTS {index['name']}"))
                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")


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