Kairo-Jinkendo/backend/run_migrations.py

#!/usr/bin/env python3
"""Apply numbered SQL migrations with schema_migrations tracking."""

from __future__ import annotations

import os
import re
import shutil
import subprocess
import sys
import time
from typing import List, Tuple

import psycopg2
import sqlparse

from db import db_params, get_connection

_LEADING_DIGITS = re.compile(r"^(\d+)")


def init_migrations_table(conn) -> None:
    with conn.cursor() as cur:
        cur.execute(
            """
            CREATE TABLE IF NOT EXISTS schema_migrations (
                id SERIAL PRIMARY KEY,
                migration VARCHAR(255) UNIQUE NOT NULL,
                executed_at TIMESTAMP DEFAULT NOW()
            )
            """
        )
    conn.commit()


def _migration_sort_key(stem: str) -> Tuple[int, str]:
    match = _LEADING_DIGITS.match(stem)
    return (int(match.group(1)) if match else 0, stem)


def migration_files(migrations_dir: str) -> List[Tuple[str, str]]:
    rows: List[Tuple[str, str]] = []
    for filename in os.listdir(migrations_dir):
        if not filename.endswith(".sql") or not filename[0].isdigit():
            continue
        stem = filename[:-4]
        rows.append((stem, os.path.join(migrations_dir, filename)))
    rows.sort(key=lambda item: _migration_sort_key(item[0]))
    return rows


def executed_migrations(conn) -> set[str]:
    with conn.cursor() as cur:
        cur.execute("SELECT migration FROM schema_migrations")
        return {row[0] for row in cur.fetchall()}


def pending_migrations(conn, migrations_dir: str) -> List[Tuple[str, str]]:
    done = executed_migrations(conn)
    return [(name, path) for name, path in migration_files(migrations_dir) if name not in done]


def _split_statements(sql_text: str) -> List[str]:
    parts = sqlparse.split(sql_text.strip())
    return [part.strip() for part in parts if part and part.strip()]


def _run_with_psql(filepath: str) -> tuple[bool, str]:
    psql = shutil.which("psql")
    if not psql:
        return False, ""

    p = db_params()
    env = os.environ.copy()
    env["PGPASSWORD"] = str(p["password"])
    cmd = [
        psql,
        "-h",
        p["host"],
        "-p",
        str(p["port"]),
        "-U",
        p["user"],
        "-d",
        p["dbname"],
        "-v",
        "ON_ERROR_STOP=1",
        "-1",
        "-f",
        filepath,
    ]
    proc = subprocess.run(cmd, env=env, capture_output=True, text=True, timeout=7200)
    if proc.returncode != 0:
        tail = ((proc.stderr or "") + "\n" + (proc.stdout or "")).strip()
        return False, tail[:8000] or f"exit {proc.returncode}"
    return True, (proc.stdout or "").strip()


def _record_migration(conn, migration_name: str) -> None:
    with conn.cursor() as cur:
        cur.execute(
            """
            INSERT INTO schema_migrations (migration)
            VALUES (%s)
            ON CONFLICT (migration) DO NOTHING
            """,
            (migration_name,),
        )


def run_migration(conn, migration_name: str, filepath: str) -> bool:
    print(f"Running migration: {migration_name}")
    try:
        if shutil.which("psql"):
            ok, diag = _run_with_psql(filepath)
            if not ok:
                print(f"  [FAIL] psql:\n{diag or '(kein Output)'}")
                conn.rollback()
                return False
        else:
            with open(filepath, "r", encoding="utf-8") as handle:
                body = handle.read()
            statements = _split_statements(body)
            with conn.cursor() as cur:
                for stmt in statements:
                    cur.execute(stmt)

        _record_migration(conn, migration_name)
        conn.commit()
        print(f"  [OK] {migration_name}")
        return True
    except Exception as exc:
        conn.rollback()
        print(f"  [FAIL] {migration_name}: {exc}")
        return False


def connect_with_retry(max_retries: int = 30):
    p = db_params()
    for attempt in range(max_retries):
        try:
            conn = get_connection()
            conn.autocommit = False
            print(f"[OK] Connected to database: {p['dbname']}")
            return conn
        except psycopg2.OperationalError:
            if attempt >= max_retries - 1:
                raise
            print(f"Waiting for database... ({attempt + 1}/{max_retries})")
            time.sleep(2)


def migrations_directory() -> str:
    docker_path = "/app/migrations"
    if os.path.isdir(docker_path):
        return docker_path
    return os.path.join(os.path.dirname(os.path.abspath(__file__)), "migrations")


def main() -> int:
    print("=" * 60)
    print("Jinkendo Kairo — Database Migrations")
    print("=" * 60)

    migrations_dir = migrations_directory()
    if not os.path.isdir(migrations_dir):
        print(f"[FAIL] migrations directory missing: {migrations_dir}")
        return 1

    try:
        conn = connect_with_retry()
        init_migrations_table(conn)
        pending = pending_migrations(conn, migrations_dir)

        if not pending:
            print("[OK] Keine ausstehenden Migrationen.")
            conn.close()
            return 0

        print(f"{len(pending)} ausstehende Migration(en):")
        for name, _ in pending:
            print(f"  - {name}")

        for migration_name, filepath in pending:
            if not run_migration(conn, migration_name, filepath):
                conn.close()
                return 1

        conn.close()
        print("[OK] Migrationen abgeschlossen.")
        return 0
    except Exception as exc:
        print(f"[FAIL] {exc}")
        return 1


if __name__ == "__main__":
    sys.exit(main())