xinmi/scripts/apply_weknora_business.py

#!/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())