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