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 sqlalchemy.pool import NullPool, QueuePool

from app.config import settings

logger = logging.getLogger(__name__)

Base = declarative_base()

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}

if _parsed_url.get_backend_name() == "sqlite":
    _connect_args["check_same_thread"] = False
    _engine_kwargs["poolclass"] = NullPool
else:
    _engine_kwargs.update(
        {
            "poolclass": QueuePool,
            "pool_size": int(os.getenv("DB_POOL_SIZE", "10")),
            "max_overflow": int(os.getenv("DB_MAX_OVERFLOW", "20")),
            "pool_timeout": int(os.getenv("DB_POOL_TIMEOUT", "30")),
            "pool_recycle": int(os.getenv("DB_POOL_RECYCLE", "1800")),
        }
    )

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 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:
        import app.models  # noqa: F401

        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: 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}")


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