xinmi/scripts/verify_asset_etl.py

#!/usr/bin/env python3
"""Run asset migration regressions in acceptance PostgreSQL and roll them back.

Source rules, skill bodies, and generated SQL stay in memory. Reports contain
only check names, booleans, counts, and the fixed acceptance scope.
"""
from __future__ import annotations

import argparse
import datetime
import json
from pathlib import Path

from apply_weknora_assets import build_plan, target_snapshot
from apply_weknora_identities import sql_run
from verify_identity_etl import rejected_case
from weknora_asset_etl import (BINDING_LEDGER, LEDGER, SOURCE, TARGET_DATABASE,
                               TARGET_TENANT, asset_sql, safe_asset_summary)
from weknora_category_etl import literal, save_json
from weknora_etl import consistent_snapshot
from weknora_identity_etl import IDENTITY_LEDGER


def supporting_snapshot():
    tables = (LEDGER, BINDING_LEDGER, IDENTITY_LEDGER, "users", "tenant_members",
              "tenants", "knowledge_bases")
    fields = []
    for table in tables:
        fields.extend((literal(table),
            f"(SELECT coalesce(json_agg(fingerprint ORDER BY fingerprint),'[]') "
            f"FROM (SELECT md5(to_jsonb(r)::text) AS fingerprint FROM {table} r) s)"))
    return json.loads(sql_run("SELECT json_build_object(" + ",".join(fields) + ");"))


def unchanged_check(name, pairs):
    conditions = []
    for table, baseline in pairs:
        conditions.extend((f"EXISTS (SELECT * FROM {table} EXCEPT SELECT * FROM {baseline})",
                           f"EXISTS (SELECT * FROM {baseline} EXCEPT SELECT * FROM {table})"))
    return f"""
DO $asset_unchanged$
BEGIN
 IF {' OR '.join(conditions)} THEN
  RAISE EXCEPTION 'Asset regression changed rows unexpectedly';
 END IF;
END $asset_unchanged$;
SELECT json_build_object('check',{literal(name)},'passed',true);
"""


def regression_sql(plan, query):
    agent, skill = plan["agents"][0]["target"], plan["skills"][0]["target"]
    aid, sid = literal(agent["id"]), literal(skill["id"])
    other_owner = literal(next(uid for uid in plan["identity_mapping"].values()
                               if uid != agent["created_by"]))
    other_skill_owner = literal(next(uid for uid in plan["identity_mapping"].values()
                                     if uid != skill["created_by"]))
    asset_scope = f"source_system={literal(SOURCE)} AND source_organization_id=1"
    agent_ledger = f"{asset_scope} AND source_kind='agent' AND target_id={aid}"
    skill_ledger = f"{asset_scope} AND source_kind='skill' AND target_id={sid}"
    binding = plan["bindings"][0]
    binding_where = f"{asset_scope} AND source_binding_id={binding['source']['id']}"
    kb = literal(next(iter(plan["kb_mapping"].values())))
    model = literal(agent["config"]["model_id"])
    owner = literal(agent["created_by"])
    foreign_tenant = TARGET_TENANT + 1
    cases = [
        ("mapped_agent_wrong_tenant_rejected", f"UPDATE custom_agents SET tenant_id={foreign_tenant} WHERE id={aid};",
         "Mapped native asset removed; refusing recreation"),
        ("mapped_skill_wrong_tenant_rejected", f"UPDATE tenant_skill_catalog SET tenant_id={foreign_tenant} WHERE id={sid};",
         "Mapped native asset removed; refusing recreation"),
        ("mapped_agent_wrong_owner_rejected", f"UPDATE custom_agents SET created_by={other_owner} WHERE id={aid};",
         "Mapped native asset removed; refusing recreation"),
        ("mapped_skill_wrong_owner_rejected", f"UPDATE tenant_skill_catalog SET created_by={other_skill_owner} WHERE id={sid};",
         "Mapped native asset removed; refusing recreation"),
        ("mapped_skill_kind_change_rejected", f"UPDATE tenant_skill_catalog SET instruction_only=false WHERE id={sid};",
         "Mapped native asset removed; refusing recreation"),
        ("asset_ledger_wrong_tenant_rejected", f"UPDATE {LEDGER} SET target_tenant_id={foreign_tenant} WHERE {agent_ledger};",
         "Asset durable mapping conflict"),
        ("asset_ledger_wrong_id_rejected", f"UPDATE {LEDGER} SET target_id='rollback-only-id' WHERE {agent_ledger};",
         "Asset durable mapping conflict"),
        ("asset_ledger_source_change_rejected", f"UPDATE {LEDGER} SET source_snapshot='{{}}'::jsonb WHERE {skill_ledger};",
         "Asset durable mapping conflict"),
        ("asset_ledger_native_plan_change_rejected", f"UPDATE {LEDGER} SET native_snapshot='{{}}'::jsonb WHERE {agent_ledger};",
         "Asset durable mapping conflict"),
        ("asset_ledger_changed_duty_rejected", f"UPDATE {LEDGER} SET native_snapshot=jsonb_set(native_snapshot,'{{config,role}}','\"invalid-role\"'::jsonb) WHERE {agent_ledger};",
         "Asset durable mapping conflict"),
        ("hard_deleted_agent_not_recreated", f"DELETE FROM custom_agents WHERE id={aid};",
         "Mapped native asset removed; refusing recreation"),
        ("hard_deleted_skill_not_recreated", f"DELETE FROM tenant_skill_catalog WHERE id={sid};",
         "Mapped native asset removed; refusing recreation"),
        ("agent_id_without_ledger_rejected", f"DELETE FROM {LEDGER} WHERE {agent_ledger};",
         "Native Agent ID or owner/name collision"),
        ("skill_id_without_ledger_rejected", f"DELETE FROM {LEDGER} WHERE {skill_ledger};",
         "Native skill ID or owner/name collision"),
        ("agent_owner_name_without_ledger_rejected", f"DELETE FROM {LEDGER} WHERE {agent_ledger}; "
         f"UPDATE custom_agents SET id='rollback-only-agent-id' WHERE id={aid};",
         "Native Agent ID or owner/name collision"),
        ("skill_owner_name_without_ledger_rejected", f"DELETE FROM {LEDGER} WHERE {skill_ledger}; "
         f"UPDATE tenant_skill_catalog SET id='rollback-only-skill-id' WHERE id={sid};",
         "Native skill ID or owner/name collision"),
        ("inactive_target_tenant_rejected", f"UPDATE tenants SET status='suspended' WHERE id={TARGET_TENANT};",
         "Asset migration requires active target tenant"),
        ("owner_identity_mapping_change_rejected", f"UPDATE {IDENTITY_LEDGER} SET target_user_id='rollback-only-owner' "
         f"WHERE {asset_scope} AND target_user_id={owner};",
         "Asset owner durable identity or membership mismatch"),
        ("owner_membership_deletion_rejected", f"UPDATE tenant_members SET deleted_at=now() WHERE tenant_id={TARGET_TENANT} AND user_id={owner};",
         "Asset owner durable identity or membership mismatch"),
        ("knowledge_base_wrong_tenant_rejected", f"UPDATE knowledge_bases SET tenant_id={foreign_tenant} WHERE id={kb};",
         "Asset knowledge base mapping missing or outside tenant"),
        ("deleted_model_rejected", f"UPDATE models SET deleted_at=now() WHERE id={model};",
         "Asset model mapping missing or not usable"),
        ("binding_wrong_tenant_rejected", f"UPDATE {BINDING_LEDGER} SET target_tenant_id={foreign_tenant} WHERE {binding_where};",
         "Legacy skill binding provenance conflict"),
        ("binding_agent_change_rejected", f"UPDATE {BINDING_LEDGER} SET target_agent_id='rollback-only-agent' WHERE {binding_where};",
         "Legacy skill binding provenance conflict"),
        ("binding_source_change_rejected", f"UPDATE {BINDING_LEDGER} SET source_snapshot='{{}}'::jsonb WHERE {binding_where};",
         "Legacy skill binding provenance conflict"),
        ("binding_access_review_change_rejected", f"UPDATE {BINDING_LEDGER} SET access_review='[\"rollback-only-review\"]'::jsonb WHERE {binding_where};",
         "Legacy skill binding provenance conflict"),
        ("binding_effective_flag_change_rejected", f"UPDATE {BINDING_LEDGER} SET effective=NOT effective WHERE {binding_where};",
         "Legacy skill binding provenance conflict"),
    ]
    if plan["prompt_templates"]:
        template_id = plan["prompt_templates"][0]["id"]
        cases.append(("template_source_change_rejected",
            f"UPDATE {LEDGER} SET source_snapshot='{{}}'::jsonb WHERE {asset_scope} "
            f"AND source_kind='template' AND source_id={template_id};",
            "Legacy prompt template provenance conflict"))
    tables = ("custom_agents", "tenant_skill_catalog", LEDGER, BINDING_LEDGER, "models")
    pairs = [(table, f"_asset_before_{index}") for index, table in enumerate(tables)]
    sql = """BEGIN ISOLATION LEVEL SERIALIZABLE;
SET LOCAL statement_timeout='30s';
SET LOCAL lock_timeout='5s';
SET LOCAL log_statement='none';
SET LOCAL log_min_error_statement='panic';
"""
    for table, baseline in pairs:
        sql += f"CREATE TEMP TABLE {baseline} ON COMMIT DROP AS SELECT * FROM {table};\n"
    sql += query + unchanged_check("exact_repeat_and_existing_models_unchanged", pairs)
    sql += f"""
SELECT json_build_object('check','all_39_source_bindings_retained_only_30_effective','passed',
 (SELECT count(*)=39 AND count(*) FILTER (WHERE effective)=30 FROM {BINDING_LEDGER}
  WHERE {asset_scope} AND target_tenant_id={TARGET_TENANT})
 AND (SELECT count(*)=15 FROM {LEDGER} WHERE {asset_scope} AND source_kind='agent')
 AND (SELECT count(*)=9 FROM {LEDGER} WHERE {asset_scope} AND source_kind='skill')
 AND NOT EXISTS (
  (SELECT b.target_agent_id,b.target_skill_id FROM {BINDING_LEDGER} b
   WHERE b.source_system={literal(SOURCE)} AND b.source_organization_id=1 AND b.effective
   EXCEPT
   SELECT a.id,s.value FROM custom_agents a JOIN {LEDGER} l
    ON l.source_kind='agent' AND l.target_id=a.id AND l.target_tenant_id=a.tenant_id
   CROSS JOIN LATERAL jsonb_array_elements_text(a.config->'instruction_skill_ids') s(value)
   WHERE l.source_system={literal(SOURCE)} AND l.source_organization_id=1)
  UNION ALL
  (SELECT a.id,s.value FROM custom_agents a JOIN {LEDGER} l
    ON l.source_kind='agent' AND l.target_id=a.id AND l.target_tenant_id=a.tenant_id
   CROSS JOIN LATERAL jsonb_array_elements_text(a.config->'instruction_skill_ids') s(value)
   WHERE l.source_system={literal(SOURCE)} AND l.source_organization_id=1
   EXCEPT
   SELECT b.target_agent_id,b.target_skill_id FROM {BINDING_LEDGER} b
   WHERE b.source_system={literal(SOURCE)} AND b.source_organization_id=1 AND b.effective)));
SELECT json_build_object('check','cross_owner_private_skills_blocked','passed',
 (SELECT count(*)=9 FROM {BINDING_LEDGER} WHERE {asset_scope} AND NOT effective)
 AND NOT EXISTS (SELECT 1 FROM {BINDING_LEDGER} b
  JOIN custom_agents a ON a.id=b.target_agent_id AND a.tenant_id=b.target_tenant_id
  JOIN tenant_skill_catalog s ON s.id=b.target_skill_id AND s.tenant_id=b.target_tenant_id
  WHERE b.source_system={literal(SOURCE)} AND b.source_organization_id=1
  AND b.effective IS DISTINCT FROM (s.visibility='global' OR s.created_by=a.created_by)));
"""
    for name, mutation, expected in cases:
        sql += rejected_case(name, mutation, expected, query)
    sql += f"""
SAVEPOINT pre_duty_replay;
UPDATE {LEDGER} SET native_snapshot=native_snapshot #- '{{config,role}}' #- '{{config,content_tags}}' #- '{{config,fusion_agent_ids}}'
 WHERE {asset_scope} AND source_kind='agent';
CREATE TEMP TABLE _pre_duty_ledger ON COMMIT DROP AS SELECT * FROM {LEDGER};
"""
    sql += query + unchanged_check("pre_duty_ledger_replay_preserves_native_edits_and_audit",
        [("custom_agents", pairs[0][1]), (LEDGER, "_pre_duty_ledger")])
    sql += "ROLLBACK TO SAVEPOINT pre_duty_replay;\n"
    duty_migration = (Path(__file__).resolve().parents[1] / "migrations/versioned/000205_agent_duties.up.sql").read_text()
    sql += f"""
SAVEPOINT duty_backfill;
UPDATE custom_agents SET config=config - 'role' - 'content_tags' WHERE id={aid};
{duty_migration}
SELECT json_build_object('check','missing_duties_backfilled_without_other_changes','passed',
 (SELECT a.config->>'role'=l.source_snapshot->>'role'
   AND a.config->'content_tags'=l.source_snapshot->'content_tags'
   AND (to_jsonb(a) #- '{{config,role}}' #- '{{config,content_tags}}') =
       (to_jsonb(b) #- '{{config,role}}' #- '{{config,content_tags}}')
  FROM custom_agents a JOIN {LEDGER} l ON l.target_id=a.id AND l.source_kind='agent'
  JOIN {pairs[0][1]} b ON b.id=a.id WHERE a.id={aid}));
UPDATE custom_agents SET config=config || '{{"role":"general","content_tags":[]}}'::jsonb WHERE id={aid};
CREATE TEMP TABLE _duty_edited_agents ON COMMIT DROP AS SELECT * FROM custom_agents;
{duty_migration}
"""
    sql += unchanged_check("duty_backfill_preserves_explicit_native_edits", [("custom_agents", "_duty_edited_agents")])
    sql += "ROLLBACK TO SAVEPOINT duty_backfill;\n"
    sql += f"""
UPDATE custom_agents SET name='Rollback-only native Agent edit',description='local edit',
 enabled=false,visibility='private',
 config=jsonb_set(jsonb_set(config,'{{system_prompt}}','\"Rollback-only native prompt\"'::jsonb),
 '{{instruction_skill_ids}}','[]'::jsonb) || '{{"role":"general","content_tags":[]}}'::jsonb WHERE id={aid} AND tenant_id={TARGET_TENANT};
UPDATE tenant_skill_catalog SET name='Rollback-only native skill edit',description='local edit',
 instructions='Rollback-only native instructions',version='77',enabled=false,visibility='private'
 WHERE id={sid} AND tenant_id={TARGET_TENANT};
CREATE TEMP TABLE _edited_agents ON COMMIT DROP AS SELECT * FROM custom_agents;
CREATE TEMP TABLE _edited_skills ON COMMIT DROP AS SELECT * FROM tenant_skill_catalog;
"""
    sql += query + unchanged_check("native_edits_disabling_and_binding_removal_preserved",
        [("custom_agents", "_edited_agents"), ("tenant_skill_catalog", "_edited_skills"), *pairs[2:]])
    sql += f"""
UPDATE custom_agents SET deleted_at=now() WHERE id={aid} AND tenant_id={TARGET_TENANT};
UPDATE tenant_skill_catalog SET deleted_at=now() WHERE id={sid} AND tenant_id={TARGET_TENANT};
CREATE TEMP TABLE _deleted_agents ON COMMIT DROP AS SELECT * FROM custom_agents;
CREATE TEMP TABLE _deleted_skills ON COMMIT DROP AS SELECT * FROM tenant_skill_catalog;
"""
    sql += query + unchanged_check("soft_deleted_assets_not_revived",
        [("custom_agents", "_deleted_agents"), ("tenant_skill_catalog", "_deleted_skills"), *pairs[2:]])
    sql += "ROLLBACK;\n"
    return sql, len(cases) + 8


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--sqlite", type=Path, required=True)
    parser.add_argument("--kb-mapping", type=Path, required=True)
    parser.add_argument("--manifest", type=Path, required=True)
    parser.add_argument("--out", type=Path, required=True)
    args = parser.parse_args()
    if args.out.exists():
        parser.error("Output already exists")
    manifest = json.loads(args.manifest.read_text())
    with consistent_snapshot(args.sqlite) as (db, info):
        plan = build_plan(db, args.kb_mapping)
        expected = {k: v for k, v in manifest.items() if k != "source_snapshot_sha256"}
        if safe_asset_summary(plan) != expected or info["source_mutated"]:
            raise RuntimeError("Asset source differs from the reviewed manifest")
        if (len(plan["agents"]), len(plan["skills"]), len(plan["bindings"]),
                sum(b["effective"] for b in plan["bindings"])) != (15, 9, 39, 30):
            raise RuntimeError("Asset counts differ from the reviewed acceptance batch")
        plan["source_snapshot_sha256"] = info["sha256"]
        query, expected_checks = regression_sql(plan, asset_sql(plan))
    before, supporting_before = target_snapshot(), supporting_snapshot()
    output = sql_run(query)
    checks = [json.loads(line) for line in output.splitlines() if line.startswith("{")]
    after, supporting_after = target_snapshot(), supporting_snapshot()
    checks.extend((
        {"check": "all_regressions_executed", "passed": len(checks) == expected_checks},
        {"check": "all_test_changes_rolled_back", "passed": before == after and supporting_before == supporting_after},
        {"check": "acceptance_scope", "passed": before["database"] == after["database"] == TARGET_DATABASE},
    ))
    report = {"checked_at": datetime.datetime.now(datetime.timezone.utc).isoformat(),
              "database": TARGET_DATABASE, "tenant_id": TARGET_TENANT,
              "committed": False, "source_mutated": info["source_mutated"],
              "checks": checks, "passed": all(c["passed"] for c in checks)}
    save_json(args.out, report)
    print(json.dumps(report))
    return 0 if report["passed"] else 1


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