xinmi/scripts/apply_weknora_identities.py

#!/usr/bin/env python3
"""Checkpointed identity import restricted to the local acceptance database."""
from __future__ import annotations

import argparse
import datetime
import json
import subprocess
from pathlib import Path

from weknora_category_etl import save_json
from weknora_etl import consistent_snapshot
from weknora_identity_etl import (ADMIN_TARGET_ID, IDENTITY_LEDGER, ORGANIZATION_LEDGER,
                                 TARGET_DATABASE, TARGET_TENANT, identity_plan, identity_sql)


def sql_run(sql):
    result = subprocess.run([
        "docker", "--context", "colima-yawan", "exec", "-i", "WeKnora-postgres",
        "psql", "-X", "-qAt", "-v", "ON_ERROR_STOP=1", "-U", "postgres", "-d", TARGET_DATABASE,
    ], input=sql, text=True, capture_output=True, timeout=60)
    if result.returncode:
        # PostgreSQL can include SQL payloads in errors. Do not print them.
        raise RuntimeError("Identity transaction not confirmed; reconcile ledger and checkpoints before retrying")
    return result.stdout.strip()


def target_snapshot():
    snapshot = json.loads(sql_run(f"""
BEGIN ISOLATION LEVEL REPEATABLE READ READ ONLY;
SELECT json_build_object(
 'database',current_database(),
 'tenant',(SELECT json_build_object('id',id,'status',status) FROM tenants WHERE id={TARGET_TENANT}),
 'users',(SELECT coalesce(json_agg(row_to_json(x) ORDER BY id),'[]') FROM
   (SELECT id,username,email,avatar,tenant_id,is_active,is_system_admin,can_access_all_tenants,
           created_at,updated_at,deleted_at FROM users) x),
 'memberships',(SELECT coalesce(json_agg(row_to_json(x) ORDER BY id),'[]') FROM tenant_members x),
 'token_counts',(SELECT coalesce(json_agg(row_to_json(x) ORDER BY user_id),'[]') FROM
   (SELECT user_id,count(*) AS count FROM auth_tokens GROUP BY user_id) x)
);
COMMIT;
"""))
    snapshot["ledgers"] = {}
    for table in (ORGANIZATION_LEDGER, IDENTITY_LEDGER):
        exists = sql_run(f"SELECT to_regclass('public.{table}') IS NOT NULL;") == "t"
        order = "source_system,source_organization_id" + (",source_user_id" if table == IDENTITY_LEDGER else "")
        snapshot["ledgers"][table] = json.loads(sql_run(f"SELECT coalesce(json_agg(x ORDER BY {order}),'[]') FROM {table} x;")) if exists else []
    return snapshot


def transaction(manifest, commit=False):
    # Save complete rows only inside PostgreSQL; credentials never enter reports.
    return f"""
BEGIN ISOLATION LEVEL SERIALIZABLE;
SET LOCAL statement_timeout = '30s';
SET LOCAL lock_timeout = '5s';
CREATE TEMP TABLE _identity_before_users ON COMMIT DROP AS SELECT * FROM users;
CREATE TEMP TABLE _identity_before_members ON COMMIT DROP AS SELECT * FROM tenant_members;
{identity_sql(manifest)}
DO $identity_preservation$
BEGIN
 IF EXISTS (SELECT 1 FROM _identity_before_users b LEFT JOIN users u ON u.id=b.id
            WHERE u.id IS NULL OR to_jsonb(u) IS DISTINCT FROM to_jsonb(b))
 THEN RAISE EXCEPTION 'Identity import changed an existing user'; END IF;
 IF EXISTS (SELECT 1 FROM _identity_before_members b LEFT JOIN tenant_members m ON m.id=b.id
            WHERE m.id IS NULL OR to_jsonb(m) IS DISTINCT FROM to_jsonb(b))
 THEN RAISE EXCEPTION 'Identity import changed an existing membership'; END IF;
END $identity_preservation$;
SELECT json_build_object(
 'inserted_users',(SELECT count(*) FROM users u WHERE NOT EXISTS (SELECT 1 FROM _identity_before_users b WHERE b.id=u.id)),
 'inserted_memberships',(SELECT count(*) FROM tenant_members m WHERE NOT EXISTS (SELECT 1 FROM _identity_before_members b WHERE b.id=m.id)),
 'identity_mappings',(SELECT count(*) FROM {IDENTITY_LEDGER} WHERE target_tenant_id={TARGET_TENANT}),
 'organization_mappings',(SELECT count(*) FROM {ORGANIZATION_LEDGER} WHERE target_tenant_id={TARGET_TENANT}),
 'existing_users_unchanged',true,'existing_memberships_unchanged',true
);
{'COMMIT' if commit else 'ROLLBACK'};
"""


def verify_states(manifest, before, after):
    checks = []

    def check(name, passed):
        checks.append({"check": name, "passed": bool(passed)})

    check("acceptance_scope", before["database"] == after["database"] == TARGET_DATABASE
          and before["tenant"] == after["tenant"] == {"id": TARGET_TENANT, "status": "active"})
    old_users = {u["id"]: u for u in before["users"]}
    users = {u["id"]: u for u in after["users"]}
    old_members = {m["id"]: m for m in before["memberships"]}
    members = {m["id"]: m for m in after["memberships"]}
    check("existing_users_unchanged", all(users.get(key) == value for key, value in old_users.items()))
    check("existing_memberships_unchanged", all(members.get(key) == value for key, value in old_members.items()))
    check("existing_admin_unchanged", ADMIN_TARGET_ID in old_users and users.get(ADMIN_TARGET_ID) == old_users[ADMIN_TARGET_ID])
    expected_ids = {u["target_id"] for u in manifest["users"] if u["action"] == "create_suspended"}
    check("only_planned_users_added", set(users) - set(old_users) == expected_ids - set(old_users))
    added_members = [m for key, m in members.items() if key not in old_members]
    check("only_planned_memberships_added", len(added_members) == len(expected_ids - set(old_users))
          and {m["user_id"] for m in added_members} == expected_ids - set(old_users))
    token_counts = {t["user_id"]: t["count"] for t in after["token_counts"]}
    for item in manifest["users"]:
        if item["action"] != "create_suspended":
            continue
        uid = item["target_id"]
        user = users.get(uid, {})
        scoped = [m for m in members.values() if m["user_id"] == uid and m["deleted_at"] is None]
        expected_contact = old_users.get(uid, {"username": item["target_username"], "email": item["target_email"]})
        check(f"user_{item['source_user']['id']}_identity", user.get("username") == expected_contact["username"]
              and user.get("email") == expected_contact["email"] and user.get("tenant_id") == TARGET_TENANT
              and user.get("deleted_at") is None and user.get("is_system_admin") is False
              and user.get("can_access_all_tenants") is False)
        check(f"user_{item['source_user']['id']}_isolated", len(scoped) == 1
              and scoped[0]["tenant_id"] == TARGET_TENANT and scoped[0]["role"] == "contributor")
        # Exact reruns preserve later activation and password-reset decisions.
        if uid not in old_users:
            check(f"user_{item['source_user']['id']}_login_disabled", user.get("is_active") is False
                  and len(scoped) == 1 and scoped[0]["status"] == "suspended" and token_counts.get(uid, 0) == 0)
    skipped_names = {u["source_user"]["username"].casefold() for u in manifest["skipped"]}
    check("disabled_source_accounts_not_added", all(u["username"].casefold() not in skipped_names
          for uid, u in users.items() if uid not in old_users))
    identities = {r["source_user_id"]: r for r in after["ledgers"][IDENTITY_LEDGER] if r["target_tenant_id"] == TARGET_TENANT}
    organizations = [r for r in after["ledgers"][ORGANIZATION_LEDGER] if r["target_tenant_id"] == TARGET_TENANT]
    check("durable_identity_mapping", set(identities) == {u["source_user"]["id"] for u in manifest["users"]}
          and all(identities[u["source_user"]["id"]]["target_user_id"] == u["target_id"]
                  and identities[u["source_user"]["id"]]["source_snapshot"] == u["source_user"] for u in manifest["users"]))
    check("durable_organization_and_roles", len(organizations) == 1
          and organizations[0]["source_snapshot"] == manifest["source_organization"]
          and organizations[0]["source_role_definitions"] == manifest["source_role_definitions"])
    return {"passed": all(c["passed"] for c in checks), "checks": checks}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("command", choices=("plan", "apply"))
    parser.add_argument("--sqlite", type=Path, required=True)
    parser.add_argument("--out", type=Path, required=True)
    args = parser.parse_args()
    suffixes = (".before.json", ".manifest.json", ".attempt.json", ".committed.json", ".after.json")
    if any(p.exists() for p in (args.out, *(args.out.with_suffix(s) for s in suffixes))):
        parser.error("Output checkpoint exists; use a new output name")
    with consistent_snapshot(args.sqlite) as (snapshot, info):
        manifest = identity_plan(snapshot, TARGET_TENANT)
        if info["source_mutated"]:
            raise RuntimeError("Source changed while taking snapshot; retry during a stable migration window")
        manifest["source_snapshot"] = info
        manifest["source_snapshot_sha256"] = info["sha256"]
    before = target_snapshot()
    save_json(args.out.with_suffix(".before.json"), before)
    save_json(args.out.with_suffix(".manifest.json"), manifest)
    rehearsal = [json.loads(line) for line in sql_run(transaction(manifest)).splitlines() if line.startswith("{")]
    result = {"checked_at": datetime.datetime.now(datetime.timezone.utc).isoformat(),
              "tenant_id": TARGET_TENANT, "database": TARGET_DATABASE,
              "source_snapshot_sha256": info["sha256"], "source_mutated": info["source_mutated"],
              "mapped_admins": 1, "planned_suspended_users": 2, "skipped_disabled_users": len(manifest["skipped"]),
              "rehearsal": rehearsal, "committed": False}
    if args.command == "apply":
        save_json(args.out.with_suffix(".attempt.json"), {**result, "commit_requested": True,
                  "recovery": "Inspect identity ledger and use a new output name to reconcile; do not delete users automatically"})
        result["changes"] = [json.loads(line) for line in sql_run(transaction(manifest, True)).splitlines() if line.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_states(manifest, before, after)
    save_json(args.out, result)
    print(json.dumps(result, ensure_ascii=False))
    return 0 if result.get("verification", {}).get("passed", True) else 1


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