WeKnora deployment source with the reviewed wiki identity, summary and recovery fixes

scripts/apply_weknora_business.py

@@ -0,0 +1,133 @@
+#!/usr/bin/env python3
+"""Checkpoint historical workbench imports into the pinned acceptance database."""
+from __future__ import annotations
+
+import argparse
+import datetime
+import json
+from pathlib import Path
+
+from apply_weknora_identities import sql_run
+from weknora_asset_etl import LEDGER as ASSET_LEDGER
+from weknora_business_etl import LEDGER, business_plan, business_sql, safe_business_summary
+from weknora_category_etl import save_json
+from weknora_etl import consistent_snapshot
+from weknora_identity_etl import IDENTITY_LEDGER, TARGET_DATABASE, TARGET_TENANT
+
+PRESERVED_TABLES = ('workbench_jobs', 'workbench_events', 'workbench_revisions',
+                    'custom_agents', 'tenant_skill_catalog', 'models')
+
+
+def build_plan(db, mapping_file, artifact_root):
+    identities = json.loads(sql_run(f"SELECT json_object_agg(source_user_id,target_user_id) FROM {IDENTITY_LEDGER} "
+        f"WHERE source_system='yawan-workbench' AND source_organization_id=1 AND target_tenant_id={TARGET_TENANT};"))
+    agents = json.loads(sql_run(f"SELECT json_object_agg(source_id,target_id) FROM {ASSET_LEDGER} "
+        f"WHERE source_system='yawan-workbench' AND source_organization_id=1 AND source_kind='agent' AND target_tenant_id={TARGET_TENANT};"))
+    state = json.loads(mapping_file.read_text())
+    if state.get('tenant_id') != TARGET_TENANT:
+        raise ValueError('Knowledge mapping tenant mismatch')
+    kbs = {key: item['id'] for key, item in state['knowledge_bases'].items()}
+    return business_plan(db, identities, agents, kbs, artifact_root)
+
+
+def target_snapshot():
+    parts = ["'database',current_database()", "'knowledge_count',(SELECT count(*) FROM knowledges)",
+             "'chunk_count',(SELECT count(*) FROM chunks)"]
+    for table in ('workbench_resources', *PRESERVED_TABLES):
+        parts.append(f"'{table}',(SELECT coalesce(json_agg(json_build_object('id',id,'hash',md5(to_jsonb(r)::text)) ORDER BY id),'[]') FROM {table} r)")
+    return json.loads(sql_run('SELECT json_build_object(' + ','.join(parts) + ');'))
+
+
+def verification_sql():
+    return f"""SELECT json_build_object(
+ 'ledger_counts',(SELECT json_object_agg(source_kind,n) FROM (SELECT source_kind,count(*) n
+   FROM {LEDGER} WHERE target_tenant_id={TARGET_TENANT} GROUP BY source_kind) q),
+ 'resources',(SELECT count(*) FROM {LEDGER} l JOIN workbench_resources r ON r.id=l.target_id
+   WHERE l.target_tenant_id={TARGET_TENANT} AND l.source_kind IN ('canvas','trainingmap','campaign')),
+ 'protected',(SELECT count(*) FROM {LEDGER} l JOIN workbench_resources r ON r.id=l.target_id
+   WHERE l.target_tenant_id={TARGET_TENANT} AND r.data->'legacy_import'->>'read_only'='true'),
+ 'segments',(SELECT coalesce(sum(jsonb_array_length(r.data->'segments')),0)
+   FROM {LEDGER} l JOIN workbench_resources r ON r.id=l.target_id
+   WHERE l.target_tenant_id={TARGET_TENANT} AND l.source_kind='campaign'),
+ 'initial_data_match',(SELECT bool_and(r.data=l.native_snapshot->'data' AND r.owner_id=l.native_snapshot->>'owner_id'
+   AND r.status=l.native_snapshot->>'status' AND r.version=(l.native_snapshot->>'version')::integer)
+   FROM {LEDGER} l JOIN workbench_resources r ON r.id=l.target_id WHERE l.target_tenant_id={TARGET_TENANT}));"""
+
+
+def transaction(query, commit=False):
+    sql = "BEGIN ISOLATION LEVEL SERIALIZABLE; SET LOCAL statement_timeout='45s'; SET LOCAL lock_timeout='5s';\n"
+    tables = ('workbench_resources', *PRESERVED_TABLES)
+    for i, table in enumerate(tables):
+        sql += f'CREATE TEMP TABLE _business_before_{i} ON COMMIT DROP AS SELECT * FROM {table};\n'
+    sql += query
+    assertions = ["EXISTS (SELECT 1 FROM _business_before_0 b LEFT JOIN workbench_resources r ON r.id=b.id WHERE r.id IS NULL OR to_jsonb(r) IS DISTINCT FROM to_jsonb(b))"]
+    for i, table in enumerate(tables[1:], 1):
+        assertions.extend((f'EXISTS (SELECT * FROM {table} EXCEPT SELECT * FROM _business_before_{i})',
+                           f'EXISTS (SELECT * FROM _business_before_{i} EXCEPT SELECT * FROM {table})'))
+    sql += "DO $preserve$ BEGIN IF " + ' OR '.join(assertions) + " THEN RAISE EXCEPTION 'Business import changed existing data or execution state'; END IF; END $preserve$;\n"
+    sql += "SELECT json_build_object('inserted_resources',(SELECT count(*) FROM workbench_resources r WHERE NOT EXISTS (SELECT 1 FROM _business_before_0 b WHERE b.id=r.id)), 'existing_data_unchanged',true);\n"
+    return sql + verification_sql() + ('COMMIT;' if commit else 'ROLLBACK;')
+
+
+def verify(plan, before, after, results):
+    old = {r['id']: r for r in before['workbench_resources']}
+    new = {r['id']: r for r in after['workbench_resources']}
+    planned = {r['target']['id'] for r in plan['resources']}
+    summary = next(r for r in results if 'ledger_counts' in r)
+    counts = safe_business_summary(plan)['counts']
+    checks = {
+        'acceptance_scope': before['database'] == after['database'] == TARGET_DATABASE,
+        'existing_resources_unchanged': all(new.get(k) == v for k, v in old.items()),
+        'only_planned_resources_added': set(new) - set(old) == planned - set(old),
+        'all_resources_present': planned <= set(new),
+        'no_jobs_events_revisions_or_other_assets_changed': all(before[k] == after[k] for k in PRESERVED_TABLES),
+        'knowledge_unchanged': all(before[k] == after[k] for k in ('knowledge_count', 'chunk_count')),
+        'all_history_counts_match': all(summary['ledger_counts'].get(k, 0) == v for k, v in counts.items()),
+        'all_segments_present': summary['segments'] == counts['segment'],
+        'historical_execution_protected': summary['protected'] == len(plan['resources']),
+    }
+    if not planned.intersection(old):
+        checks['initial_content_status_and_owner_match'] = summary['initial_data_match']
+    return {'passed': all(checks.values()), '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('--kb-mapping', type=Path, required=True)
+    parser.add_argument('--artifact-root', type=Path, required=True)
+    parser.add_argument('--out', type=Path, required=True)
+    args = parser.parse_args()
+    if any(p.exists() for p in [args.out, *(args.out.with_suffix(s) for s in ('.before.json', '.manifest.json', '.attempt.json', '.committed.json', '.after.json'))]):
+        parser.error('Checkpoint exists; choose a fresh output name')
+    with consistent_snapshot(args.sqlite) as (db, info):
+        plan = build_plan(db, args.kb_mapping, args.artifact_root)
+        if info['source_mutated']:
+            raise RuntimeError('Source changed during snapshot')
+        plan['source_snapshot_sha256'] = info['sha256']
+        query = business_sql(plan)
+    before = target_snapshot()
+    save_json(args.out.with_suffix('.before.json'), before)
+    save_json(args.out.with_suffix('.manifest.json'), {**safe_business_summary(plan), 'source_snapshot_sha256': info['sha256']})
+    rehearsal = [json.loads(line) for line in sql_run(transaction(query)).splitlines() if line.startswith('{')]
+    if before != target_snapshot():
+        raise RuntimeError('Target changed during rollback rehearsal')
+    result = {'checked_at': datetime.datetime.now(datetime.timezone.utc).isoformat(), 'database': TARGET_DATABASE,
+              'tenant_id': TARGET_TENANT, 'source_mutated': info['source_mutated'],
+              'source_snapshot_sha256': info['sha256'], 'rehearsal': rehearsal, 'committed': False}
+    if args.command == 'apply':
+        save_json(args.out.with_suffix('.attempt.json'), {**result, 'commit_requested': True})
+        results = [json.loads(line) for line in sql_run(transaction(query, True)).splitlines() if line.startswith('{')]
+        result.update(committed=True, changes=results)
+        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, results)
+    save_json(args.out, result)
+    print(json.dumps(result))
+    return 0 if result.get('verification', {}).get('passed', True) else 1
+
+
+if __name__ == '__main__':
+    raise SystemExit(main())