Kairo-Jinkendo/backend/run_seeds.py

#!/usr/bin/env python3
"""Apply idempotent data seeds with checksum-based tracking and re-run on change."""

from __future__ import annotations

import hashlib
import importlib.util
import os
import re
import sys
import time
from typing import Callable, List, Optional, Tuple

import psycopg2
import sqlparse

from db import db_params, get_connection

_SEED_PREFIX = re.compile(r"^seed_(\d+)_(.+)$")
_DEV_MARKER = ".dev."


def is_production() -> bool:
    return os.getenv("ENVIRONMENT", "development").strip().lower() == "production"


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


def init_data_seeds_table(conn) -> None:
    with conn.cursor() as cur:
        cur.execute(
            """
            CREATE TABLE IF NOT EXISTS data_seeds (
                seed_name VARCHAR(255) PRIMARY KEY,
                checksum VARCHAR(64) NOT NULL,
                executed_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
                last_run_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
            )
            """
        )
    conn.commit()


def _seed_sort_key(filename: str) -> Tuple[int, str]:
    stem = filename
    for suffix in (".sql", ".py"):
        if stem.endswith(suffix):
            stem = stem[: -len(suffix)]
            break
    match = _SEED_PREFIX.match(stem)
    if match:
        return (int(match.group(1)), stem)
    return (0, stem)


def _seed_stem(filename: str) -> str:
    if filename.endswith(".dev.sql"):
        return filename[: -len(".dev.sql")]
    if filename.endswith(".sql"):
        return filename[: -len(".sql")]
    if filename.endswith(".py"):
        return filename[: -len(".py")]
    return filename


def seed_files(seeds_dir: str) -> List[Tuple[str, str, str]]:
    """Return (seed_name, filepath, kind) sorted by numeric prefix."""
    rows: List[Tuple[str, str, str]] = []
    if not os.path.isdir(seeds_dir):
        return rows

    for filename in os.listdir(seeds_dir):
        if filename.startswith("_") or not filename.startswith("seed_"):
            continue
        if filename.endswith(".sql"):
            kind = "sql"
        elif filename.endswith(".py"):
            kind = "py"
        else:
            continue
        stem = _seed_stem(filename)
        rows.append((stem, os.path.join(seeds_dir, filename), kind))

    rows.sort(key=lambda item: _seed_sort_key(item[0]))
    return rows


def seed_applies_in_environment(filename: str) -> bool:
    if _DEV_MARKER in filename and is_production():
        return False
    return True


def file_checksum(filepath: str) -> str:
    digest = hashlib.sha256()
    with open(filepath, "rb") as handle:
        for chunk in iter(lambda: handle.read(65536), b""):
            digest.update(chunk)
    return digest.hexdigest()


def applied_seed_checksum(conn, seed_name: str) -> Optional[str]:
    with conn.cursor() as cur:
        cur.execute("SELECT checksum FROM data_seeds WHERE seed_name = %s", (seed_name,))
        row = cur.fetchone()
        return row[0] if row else None


def record_seed(conn, seed_name: str, checksum: str) -> None:
    with conn.cursor() as cur:
        cur.execute(
            """
            INSERT INTO data_seeds (seed_name, checksum, executed_at, last_run_at)
            VALUES (%s, %s, NOW(), NOW())
            ON CONFLICT (seed_name) DO UPDATE SET
                checksum = EXCLUDED.checksum,
                last_run_at = NOW()
            """,
            (seed_name, checksum),
        )


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_sql_seed(conn, filepath: str) -> None:
    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)


def run_python_seed(filepath: str) -> None:
    module_name = f"kairo_seed_{hashlib.md5(filepath.encode()).hexdigest()[:12]}"
    spec = importlib.util.spec_from_file_location(module_name, filepath)
    if spec is None or spec.loader is None:
        raise RuntimeError(f"Seed-Modul nicht ladbar: {filepath}")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    run_fn: Callable[[], None] | None = getattr(module, "run", None)
    if run_fn is None:
        raise RuntimeError(f"Seed {filepath} hat keine run()-Funktion")
    run_fn()


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 run_seed(conn, seed_name: str, filepath: str, kind: str) -> tuple[bool, object]:
    print(f"Running seed: {seed_name}")
    try:
        if kind == "sql":
            run_sql_seed(conn, filepath)
        else:
            conn.commit()
            conn.close()
            run_python_seed(filepath)
            conn = connect_with_retry(max_retries=5)

        checksum = file_checksum(filepath)
        record_seed(conn, seed_name, checksum)
        conn.commit()
        print(f"  [OK] {seed_name}")
        return True, conn
    except Exception as exc:
        try:
            conn.rollback()
        except Exception:
            pass
        print(f"  [FAIL] {seed_name}: {exc}")
        return False, conn


def pending_seeds(
    conn,
    seeds_dir: str,
    *,
    only: Optional[List[str]] = None,
    force: bool = False,
) -> List[Tuple[str, str, str]]:
    selected: List[Tuple[str, str, str]] = []
    for seed_name, filepath, kind in seed_files(seeds_dir):
        filename = os.path.basename(filepath)
        if not seed_applies_in_environment(filename):
            print(f"  [SKIP] {seed_name} (nur Nicht-Prod)")
            continue
        if only and seed_name not in only:
            continue
        checksum = file_checksum(filepath)
        is_dev_seed = _DEV_MARKER in filename
        always_run = is_dev_seed and not is_production()
        if force or always_run or applied_seed_checksum(conn, seed_name) != checksum:
            selected.append((seed_name, filepath, kind))
    return selected


def main(argv: Optional[List[str]] = None) -> int:
    argv = argv if argv is not None else sys.argv[1:]
    only: Optional[List[str]] = None
    force = False

    idx = 0
    while idx < len(argv):
        arg = argv[idx]
        if arg == "--only" and idx + 1 < len(argv):
            only = [part.strip() for part in argv[idx + 1].split(",") if part.strip()]
            idx += 2
            continue
        if arg == "--force":
            force = True
            idx += 1
            continue
        print(f"[FAIL] Unbekanntes Argument: {arg}")
        return 1

    print("=" * 60)
    print("Jinkendo Kairo — Data Seeds")
    print("=" * 60)

    seeds_dir = seeds_directory()
    if not os.path.isdir(seeds_dir):
        print(f"[OK] Kein seeds-Verzeichnis ({seeds_dir}) — nichts zu tun.")
        return 0

    try:
        conn = connect_with_retry()
        init_data_seeds_table(conn)
        to_run = pending_seeds(conn, seeds_dir, only=only, force=force)

        if not to_run:
            print("[OK] Alle Seeds aktuell — nichts auszuführen.")
            conn.close()
            return 0

        print(f"{len(to_run)} Seed(s) ausstehend:")
        for seed_name, _, _ in to_run:
            print(f"  - {seed_name}")

        for seed_name, filepath, kind in to_run:
            ok, conn = run_seed(conn, seed_name, filepath, kind)
            if not ok:
                conn.close()
                return 1

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


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