gh-christianlouis-docuelevate/app/database.py

# app/database.py

import logging
import os
from collections.abc import Generator
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.
    Logs a message if a new SQLite DB file is created.
    """
    # 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 lightweight schema migrations for existing databases
        _run_schema_migrations(engine)
    except exc.SQLAlchemyError as e:
        logger.error(f"Error initializing database: {e}")
        raise


def _run_schema_migrations(engine: Any) -> None:
    """
    Apply lightweight schema migrations for columns added after the initial release.
    Each migration is idempotent and safe to run multiple times.
    """
    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")

    # 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:
        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]:
                    conn.execute(text(f"CREATE INDEX IF NOT EXISTS {idx_name} ON {table} ({column})"))

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