database/init_db.py

"""Idempotent SQLite schema initialization.

SQLite has no docker-entrypoint-initdb.d equivalent, so this is called on
every startup of the long-running processes (API, scheduler). Every
table/index in schema.sql uses IF NOT EXISTS, so repeated calls are safe.
"""

from __future__ import annotations

from pathlib import Path

from database.connection import connect

SCHEMA_PATH = Path(__file__).parent / "schema.sql"

JOB_RUN_PROGRESS_COLUMNS = {
    "total_rows": "INTEGER NOT NULL DEFAULT 0",
    "resolved_rows": "INTEGER NOT NULL DEFAULT 0",
    "failed_rows": "INTEGER NOT NULL DEFAULT 0",
    "last_updated_at": "TIMESTAMP",
}


def _migrate_job_runs(conn) -> None:
    """Add progress columns for databases created before resolution tracking."""
    existing = {row["name"] for row in conn.execute("PRAGMA table_info(job_runs)")}
    for name, definition in JOB_RUN_PROGRESS_COLUMNS.items():
        if name not in existing:
            conn.execute(f"ALTER TABLE job_runs ADD COLUMN {name} {definition}")


def init_db(database_url: str | None = None) -> None:
    from app.config import settings

    conn = connect(database_url or settings.DATABASE_URL)
    try:
        conn.executescript(SCHEMA_PATH.read_text())
        _migrate_job_runs(conn)
        conn.commit()
    finally:
        conn.close()


if __name__ == "__main__":
    init_db()
    print("Database schema initialized.")