# 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() # Parse the DATABASE_URL DB_URL = settings.database_url _db_url = make_url(DB_URL) if _db_url.get_backend_name() == "sqlite": # SQLite does not benefit from connection pooling; NullPool avoids contention. engine = create_engine(DB_URL, connect_args={"check_same_thread": False}, poolclass=NullPool) else: # PostgreSQL / MySQL / other: use a configurable QueuePool. engine = create_engine( DB_URL, poolclass=QueuePool, pool_size=settings.db_pool_size, max_overflow=settings.db_max_overflow, pool_timeout=settings.db_pool_timeout, pool_recycle=settings.db_pool_recycle, ) 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()