Kansho/backend/db_init.py

#!/usr/bin/env python3
"""Wait for PostgreSQL, load schema.sql, apply numbered SQL migrations. Fail-fast."""
from __future__ import annotations

import os
import re
import sys
import time
from pathlib import Path

from sql_compat import postgres_connect_kwargs, split_sql, sqlite_schema_to_postgres, use_postgres

SCHEMA_PATH = Path(__file__).resolve().parent / "schema.sql"
MIGRATIONS_DIR = Path(__file__).resolve().parent / "migrations"
SQLITE_HISTORY_MIGRATIONS = (
    "001_frame",
    "002_platform",
    "003_dialogue_memory",
    "004_mvp_journal",
    "005_provider_settings",
    "006_writing_profile_dialogue_style",
    "007_journal_day_scratch",
    "008_journal_source_refs",
    "009_conversation_signals",
    "010_profile_governance",
    "011_profile_review",
    "012_profile_shell",
    "013_journal_generate_narration",
    "014_identity_registry",
    "015_journal_generation_settings",
    "016_generation_instruction_fragments",
    "017_generation_guidelines",
    "018_debug_runs",
    "019_debug_run_placement",
    "020_voice_style_context",
    "021_voice_legacy_immutable",
)
_LEADING_DIGITS = re.compile(r"^(\d{3})_.*\.sql$")
_schema_ready = False


def mark_schema_dirty() -> None:
    """Call after DROP SCHEMA so the next init_db reloads schema.sql."""
    global _schema_ready
    _schema_ready = False


def _connect_raw():
    import psycopg

    kwargs = postgres_connect_kwargs()
    if "conninfo" in kwargs:
        return psycopg.connect(kwargs["conninfo"], autocommit=False)
    return psycopg.connect(autocommit=False, **kwargs)


def wait_for_postgres(max_retries: int = 30) -> None:
    print("Checking PostgreSQL connection...")
    last_error = None
    for attempt in range(1, max_retries + 1):
        try:
            conn = _connect_raw()
            conn.close()
            print("PostgreSQL ready")
            return
        except Exception as exc:  # noqa: BLE001 — fail-fast after retries
            last_error = exc
            print(f"  waiting ({attempt}/{max_retries})")
            time.sleep(2)
    print(f"PostgreSQL not ready: {last_error}")
    sys.exit(1)


def _table_exists(conn, name: str) -> bool:
    row = conn.execute(
        """
        SELECT 1 FROM information_schema.tables
        WHERE table_schema = 'public' AND table_name = %s
        """,
        (name,),
    ).fetchone()
    return row is not None


def _load_schema(conn) -> None:
    sql = sqlite_schema_to_postgres(SCHEMA_PATH.read_text(encoding="utf-8"))
    for statement in split_sql(sql):
        conn.execute(statement)
    print("Schema loaded from schema.sql (Postgres dialect)")


def _applied_ids(conn) -> set[str]:
    rows = conn.execute("SELECT id FROM schema_migrations").fetchall()
    return {row[0] for row in rows}


def _record(conn, migration_id: str) -> None:
    conn.execute(
        "INSERT INTO schema_migrations (id) VALUES (%s) ON CONFLICT (id) DO NOTHING",
        (migration_id,),
    )


def _migration_files() -> list[tuple[str, Path]]:
    if not MIGRATIONS_DIR.is_dir():
        return []
    rows: list[tuple[str, Path]] = []
    for path in sorted(MIGRATIONS_DIR.iterdir()):
        match = _LEADING_DIGITS.match(path.name)
        if not match or path.suffix != ".sql":
            continue
        rows.append((path.stem, path))
    return rows


def apply_schema_and_migrations() -> None:
    conn = _connect_raw()
    try:
        if not _table_exists(conn, "schema_migrations"):
            _load_schema(conn)
            conn.commit()
        applied = _applied_ids(conn)
        if "001_frame" not in applied:
            for migration_id in SQLITE_HISTORY_MIGRATIONS:
                _record(conn, migration_id)
            conn.commit()
            print("Recorded SQLite history 001-021 as applied (greenfield Postgres schema)")
        applied = _applied_ids(conn)
        for stem, path in _migration_files():
            if stem in applied:
                continue
            sql = path.read_text(encoding="utf-8")
            for statement in split_sql(sql):
                conn.execute(statement)
            _record(conn, stem)
            conn.commit()
            print(f"Applied migration {stem}")
    except Exception as exc:
        conn.rollback()
        print(f"Database initialization failed: {exc}")
        raise
    finally:
        conn.close()


def ensure_postgres_ready() -> None:
    global _schema_ready
    if not use_postgres():
        return
    if _schema_ready:
        return
    wait_for_postgres()
    try:
        apply_schema_and_migrations()
    except Exception:
        sys.exit(1)
    _schema_ready = True


def main() -> None:
    os.environ.setdefault("KANSHO_DB_BACKEND", "postgres")
    ensure_postgres_ready()
    print("db_init complete")


if __name__ == "__main__":
    main()