xinmi/scripts/verify_operations_etl.py

#!/usr/bin/env python3
"""Rollback-only regressions for operations import idempotency and isolation."""
import argparse
import copy
import json
from pathlib import Path

from apply_operations_etl import sql_run, target_snapshot
from weknora_category_etl import save_json
from weknora_identity_etl import TARGET_TENANT
from weknora_operations_etl import LEDGER, operations_sql


def run_case(query, setup="", assertion="", rejected=False):
    before = target_snapshot()
    try:
        sql_run("BEGIN; SET LOCAL statement_timeout='180s'; SET LOCAL lock_timeout='5s';\n" + setup + query + assertion + "\nROLLBACK;")
        result = not rejected
    except RuntimeError:
        result = rejected
    after = target_snapshot()
    return {"passed": result and before == after, "rollback_unchanged": before == after,
            "expected_rejection": rejected}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--export", type=Path, required=True)
    parser.add_argument("--out", type=Path, required=True)
    args = parser.parse_args()
    if args.out.exists():
        parser.error("Operations regression checkpoint already exists")
    original = json.loads(args.export.read_text())
    usage = next(i for i in original["items"] if i["target_table"] == "workbench_usage")
    log = next(i for i in original["items"] if i["source_kind"] == "systemlog")
    plan = copy.copy(original)
    plan["items"] = [i for i in original["items"] if i["source_kind"] in ("pointaccount", "pointtransaction") or i in (usage, log)]
    query = operations_sql(plan)
    uid = usage["target_key"]
    hid = usage["history_id"]
    source_id = usage["source_id"]
    balance = next(i for i in plan["items"] if i["target_table"] == "workbench_balances")
    owner = balance["target_key"]
    alternate_owner = next(i for i in plan["identity_mapping"].values() if i != usage["owner_id"])
    no_new = "DO $$ BEGIN IF EXISTS(SELECT 1 FROM _ops_new) THEN RAISE EXCEPTION 'Operations repeat added rows'; END IF; END $$;"
    cases = {"exact_repeat": run_case(query, assertion=no_new)}
    edit = f"UPDATE workbench_usage SET action='edited locally' WHERE id='{uid}'; UPDATE workbench_balances SET balance=balance+1 WHERE tenant_id={TARGET_TENANT} AND user_id='{owner}';"
    retained = f"DO $$ BEGIN IF (SELECT action FROM workbench_usage WHERE id='{uid}')<>'edited locally' OR (SELECT balance FROM workbench_balances WHERE tenant_id={TARGET_TENANT} AND user_id='{owner}')<>{balance['native_snapshot']['balance'] + 1} THEN RAISE EXCEPTION 'Operations local edit overwritten'; END IF; END $$;"
    cases["native_edits_preserved"] = run_case(query, edit, retained)
    cases["deleted_native_not_recreated"] = run_case(query, f"DELETE FROM workbench_usage WHERE id='{uid}';", rejected=True)
    cases["deleted_history_not_recreated"] = run_case(query, f"DELETE FROM workbench_history WHERE id='{hid}';", rejected=True)
    cases["changed_owner_rejected"] = run_case(query, f"UPDATE workbench_usage SET owner_id='{alternate_owner}' WHERE id='{uid}';", rejected=True)
    cases["native_without_provenance_rejected"] = run_case(query, f"DELETE FROM {LEDGER} WHERE source_system='yawan-workbench' AND source_organization_id=1 AND source_kind='usageevent' AND source_id={source_id};", rejected=True)
    cases["changed_provenance_rejected"] = run_case(query, f"UPDATE {LEDGER} SET source_snapshot=source_snapshot||'{{\"changed\":true}}'::jsonb WHERE source_kind='usageevent' AND source_id={source_id};", rejected=True)
    cases["changed_tenant_rejected"] = run_case(query, f"UPDATE workbench_usage SET tenant_id=99999 WHERE id='{uid}';", rejected=True)
    result = {"passed": all(case["passed"] for case in cases.values()), "cases": cases,
              "private_source_export": str(args.export), "tests_rollback_only": True}
    save_json(args.out, result)
    print(json.dumps(result))
    return 0 if result["passed"] else 1


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