xinmi/scripts/apply_operations_etl.py

#!/usr/bin/env python3
"""Back up, rehearse and import operational history into local acceptance."""
from __future__ import annotations

import argparse
import contextlib
import datetime
import hashlib
import json
import os
import shutil
import sqlite3
import subprocess
import tempfile
from pathlib import Path

from apply_weknora_identities import sql_run as mapping_sql
from weknora_asset_etl import LEDGER as ASSET_LEDGER
from weknora_business_etl import LEDGER as BUSINESS_LEDGER
from weknora_category_etl import save_json
from weknora_etl import consistent_snapshot
from weknora_identity_etl import IDENTITY_LEDGER, TARGET_DATABASE, TARGET_TENANT
from weknora_operations_etl import LEDGER, TABLES, operations_plan, operations_sql, safe_operations_summary

PRESERVED = ("workbench_resources", "workbench_jobs", "workbench_revisions", "workbench_balance_holds")
DOCKER = ["docker", "--context", "colima-yawan", "exec", "-i", "WeKnora-postgres"]


def sql_run(query):
    result = subprocess.run(DOCKER + ["psql", "-X", "-qAt", "-v", "ON_ERROR_STOP=1", "-U", "postgres", "-d", TARGET_DATABASE],
                            input=query, text=True, capture_output=True, timeout=600)
    if result.returncode:
        # Do not expose source text included in SQL error contexts.
        lines = [line for line in result.stderr.splitlines() if line.startswith("ERROR:")]
        error = lines[0][:180] if lines and "Operations " in lines[0] else "database rejected operations transaction"
        raise RuntimeError(error)
    return result.stdout.strip()


def digest(path):
    h = hashlib.sha256()
    with path.open("rb") as f:
        for chunk in iter(lambda: f.read(1024 * 1024), b""):
            h.update(chunk)
    return h.hexdigest()


def quiescent(path):
    if any(Path(str(path) + suffix).exists() for suffix in ("-wal", "-journal")):
        raise RuntimeError("Operations fallback requires no source WAL or journal")
    check = subprocess.run(["lsof", "-t", str(path)], capture_output=True, text=True, timeout=15)
    if check.returncode != 1 or check.stdout.strip():
        raise RuntimeError("Operations fallback requires no process holding the source database")


@contextlib.contextmanager
def source_snapshot(path):
    """Use SQLite backup; compensate for missing WAL sidecars on a quiet source.

Some SQLite builds cannot read a WAL-mode file read-only when sidecars are
absent. Only a byte-identical, verified quiet copy is opened read-write to
initialize its temporary sidecars. The original is never opened for writing.
    """
    path = path.resolve()
    manager = consistent_snapshot(path)
    try:
        db, info = manager.__enter__()
    except sqlite3.OperationalError as error:
        if "unable to open database file" not in str(error):
            raise
    else:
        try:
            yield db, info
        finally:
            manager.__exit__(None, None, None)
        return
    quiescent(path)
    before = path.stat()
    with tempfile.TemporaryDirectory(prefix="operations-frozen-source-") as root:
        copied = Path(root) / "source.sqlite"
        shutil.copyfile(path, copied)
        fingerprint = digest(copied)
        if digest(path) != fingerprint or before.st_mtime_ns != path.stat().st_mtime_ns:
            raise RuntimeError("Operations source changed during quiet copy")
        quiescent(path)
        with sqlite3.connect(copied) as initializer:
            initializer.execute("SELECT count(*) FROM sqlite_master").fetchone()
            with consistent_snapshot(copied) as (db, info):
                info.update(snapshot_method="verified_quiescent_copy_then_sqlite_backup", source_bytes=before.st_size,
                            source_mtime_ns=before.st_mtime_ns, original_file_sha256=fingerprint)
                yield db, info
        quiescent(path)
        if (before.st_size, before.st_mtime_ns) != (path.stat().st_size, path.stat().st_mtime_ns) or digest(path) != fingerprint:
            raise RuntimeError("Operations source changed during snapshot read")


def build_plan(db, mapping_file):
    ids = json.loads(mapping_sql(f"SELECT json_object_agg(source_user_id,target_user_id) FROM {IDENTITY_LEDGER} WHERE source_system='yawan-workbench' AND source_organization_id=1 AND target_tenant_id={TARGET_TENANT};"))
    agents = json.loads(mapping_sql(f"SELECT json_object_agg(source_id,target_id) FROM {ASSET_LEDGER} WHERE source_system='yawan-workbench' AND source_organization_id=1 AND source_kind='agent' AND target_tenant_id={TARGET_TENANT};"))
    business = json.loads(mapping_sql(f"SELECT json_object_agg(source_kind||':'||source_id,target_id) FROM {BUSINESS_LEDGER} WHERE source_system='yawan-workbench' AND source_organization_id=1 AND target_tenant_id={TARGET_TENANT};"))
    model_rows = json.loads(mapping_sql(f"SELECT coalesce(json_agg(json_build_object('provider_id',source_provider_id,'snapshot',source_snapshot,'target_id',target_model_id)),'[]') FROM yawan_legacy_model_mappings WHERE source_system='yawan-workbench' AND source_organization_id=1 AND target_tenant_id={TARGET_TENANT};"))
    models = {}
    for row in model_rows:
        name = row["snapshot"].get("model_name", row["snapshot"].get("model", ""))
        if name:
            models[f"{row['provider_id']}:{name}"] = row["target_id"]
    state = json.loads(mapping_file.read_text())
    if state.get("tenant_id") != TARGET_TENANT:
        raise ValueError("Operations KB mapping tenant differs")
    return operations_plan(db, ids, agents, {k: v["id"] for k, v in state["knowledge_bases"].items()}, business, models)


def target_snapshot():
    tables = (*TABLES, "workbench_history", *PRESERVED)
    entries = ["'database',current_database()"]
    for table in tables:
        key = "tenant_id::text||':'||user_id" if table == "workbench_balances" else "id"
        entries.append(f"'{table}',(SELECT coalesce(json_agg(json_build_object('id',{key},'hash',md5(to_jsonb(r)::text)) ORDER BY {key}),'[]') FROM {table} r)")
    return json.loads(sql_run("SELECT json_build_object(" + ",".join(entries) + ");"))


def backup(path):
    if path.exists():
        raise RuntimeError("Operations backup exists")
    args = DOCKER + ["pg_dump", "-U", "postgres", "-d", TARGET_DATABASE, "--format=custom"]
    for table in (*TABLES, "workbench_history", *PRESERVED):
        args.extend(["--table", table])
    if mapping_sql(f"SELECT to_regclass('{LEDGER}') IS NOT NULL;") == "t":
        args.extend(["--table", LEDGER])
    with path.open("xb") as output:
        os.chmod(path, 0o600)
        result = subprocess.run(args, stdout=output, stderr=subprocess.PIPE, timeout=120)
    if result.returncode or not path.stat().st_size:
        raise RuntimeError("Operations target backup failed")
    os.chmod(path, 0o400)
    return {"path": str(path), "bytes": path.stat().st_size, "sha256": digest(path)}


def export_plan(path, plan):
    with path.open("x") as out:
        os.chmod(path, 0o600)
        json.dump(plan, out, ensure_ascii=True, sort_keys=True, separators=(",", ":"))
    os.chmod(path, 0o400)
    return {"path": str(path), "bytes": path.stat().st_size, "sha256": digest(path), "mode": "0400"}


def transaction(query, commit=False):
    sql = "BEGIN ISOLATION LEVEL SERIALIZABLE; SET LOCAL statement_timeout='180s'; SET LOCAL lock_timeout='5s';\n"
    for i, table in enumerate((*TABLES, "workbench_history", *PRESERVED)):
        sql += f"CREATE TEMP TABLE _ops_before_{i} ON COMMIT DROP AS SELECT * FROM {table};\n"
    sql += query
    for i, table in enumerate((*TABLES, "workbench_history", *PRESERVED)):
        if table == "workbench_balances":
            join = "n.tenant_id=b.tenant_id AND n.user_id=b.user_id"
        else:
            join = "n.id=b.id"
        assertion = f"EXISTS (SELECT 1 FROM _ops_before_{i} b LEFT JOIN {table} n ON {join} WHERE to_jsonb(n) IS DISTINCT FROM to_jsonb(b))"
        if table in PRESERVED:
            assertion += f" OR (SELECT count(*) FROM {table})<>(SELECT count(*) FROM _ops_before_{i})"
        sql += f"DO $preserve$ BEGIN IF {assertion} THEN RAISE EXCEPTION 'Operations existing data changed in {table}'; END IF; END $preserve$;\n"
    return sql + ("COMMIT;" if commit else "ROLLBACK;")


def verify(plan, before, after):
    checks = {"acceptance_only": before["database"] == after["database"] == TARGET_DATABASE}
    for table in (*TABLES, "workbench_history", *PRESERVED):
        old = {r["id"]: r["hash"] for r in before[table]}
        new = {r["id"]: r["hash"] for r in after[table]}
        checks[table + "_existing_unchanged"] = all(new.get(k) == v for k, v in old.items())
        if table in PRESERVED:
            checks[table + "_no_new"] = old == new
        else:
            if table == "workbench_history":
                expected = {i["history_id"] for i in plan["items"]}
            else:
                expected = {f"{TARGET_TENANT}:{i['target_key']}" if table == "workbench_balances" else i["target_key"]
                            for i in plan["items"] if i["target_table"] == table}
            checks[table + "_planned_additions_only"] = set(new) - set(old) == expected - set(old)
            checks[table + "_all_planned_present"] = expected <= set(new)
    return {"passed": all(checks.values()), "checks": checks}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("command", choices=("plan", "apply", "repeat"))
    parser.add_argument("--sqlite", type=Path)
    parser.add_argument("--kb-mapping", type=Path)
    parser.add_argument("--export", type=Path)
    parser.add_argument("--out", type=Path, required=True)
    args = parser.parse_args()
    if args.out.exists():
        parser.error("Operations checkpoint already exists")
    if args.command == "repeat":
        plan = json.loads(args.export.read_text())
        info = {"sha256": plan["source_snapshot_sha256"], "source_mutated": False, "snapshot_method": "frozen_verified_export"}
    else:
        with source_snapshot(args.sqlite) as (db, info):
            plan = build_plan(db, args.kb_mapping)
            plan["source_snapshot_sha256"] = info["sha256"]
        if info["source_mutated"]:
            raise RuntimeError("Operations source snapshot changed")
    query = operations_sql(plan)
    result = {"checked_at": datetime.datetime.now(datetime.timezone.utc).isoformat(), "summary": safe_operations_summary(plan),
              "snapshot": info, "committed": False}
    if args.command != "repeat":
        result["private_export"] = export_plan(args.out.with_suffix(".private.json"), plan)
    save_json(args.out.with_suffix(".manifest.json"), result)
    result["backup"] = backup(args.out.with_suffix(".backup.dump"))
    before = target_snapshot()
    save_json(args.out.with_suffix(".before.json"), before)
    result["rehearsal"] = [json.loads(s) for s in sql_run(transaction(query)).splitlines() if s.startswith("{")]
    result["rollback_unchanged"] = before == target_snapshot()
    if not result["rollback_unchanged"]:
        raise RuntimeError("Operations rollback rehearsal changed target")
    if args.command in ("apply", "repeat"):
        save_json(args.out.with_suffix(".attempt.json"), result)
        result["applied"] = [json.loads(s) for s in sql_run(transaction(query, True)).splitlines() if s.startswith("{")]
        result["committed"] = True
        save_json(args.out.with_suffix(".committed.json"), result)
        after = target_snapshot()
        save_json(args.out.with_suffix(".after.json"), after)
        result["verification"] = verify(plan, before, after)
        if args.command == "repeat":
            result["repeat_no_change"] = before == after and all(r.get("inserted_history") == 0 and r.get("inserted_native") == 0 for r in result["applied"])
    save_json(args.out, result)
    print(json.dumps(result, ensure_ascii=False))
    return 0 if result.get("verification", {}).get("passed", True) and result.get("repeat_no_change", True) else 1


if __name__ == "__main__":
    raise SystemExit(main())