xinmi/scripts/apply_weknora_models.py

#!/usr/bin/env python3
"""Checkpointed model import into the isolated acceptance database."""
from __future__ import annotations

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

from apply_weknora_identities import sql_run
from weknora_category_etl import save_json
from weknora_etl import consistent_snapshot
from weknora_identity_etl import TARGET_DATABASE, TARGET_TENANT
from weknora_model_etl import LEDGER, model_plan, model_sql, model_verification_sql, read_credentials


def resolve_source_credential(value, provider_id):
    if not value.startswith('keychain:'):
        return value
    expected = f'keychain:yawen-workbench.providers:provider-{provider_id}'
    if value != expected:
        raise ValueError('Source credential reference does not match the reviewed provider')
    result = subprocess.run(['/usr/bin/security', 'find-generic-password', '-s', value[9:], '-w'],
                            capture_output=True, text=True, timeout=10)
    if result.returncode or not result.stdout.strip():
        raise RuntimeError('Source provider credential could not be resolved')
    return result.stdout.strip()


def native_encryptor(environment: Path):
    env = dict(line.split("=", 1) for line in environment.read_text().splitlines()
               if line and not line.startswith("#") and "=" in line)
    key = env.get("SYSTEM_AES_KEY", "")
    if env.get("DB_NAME") != TARGET_DATABASE or len(key.encode()) != 32:
        raise ValueError("Acceptance encryption configuration is invalid")

    def encrypt(value):
        # Node's standard crypto API matches internal/utils/crypto.go. Secrets
        # enter through stdin, never argv, an environment override, or a file.
        script = """const fs = require('node:fs'), crypto = require('node:crypto');
const input = JSON.parse(fs.readFileSync(0, 'utf8'));
const nonce = crypto.randomBytes(12);
const cipher = crypto.createCipheriv('aes-256-gcm', Buffer.from(input.key), nonce);
const data = Buffer.concat([nonce, cipher.update(input.value, 'utf8'), cipher.final(), cipher.getAuthTag()]);
process.stdout.write('enc:v1:' + data.toString('base64url'));
"""
        result = subprocess.run(["node", "-e", script], input=json.dumps({"key": key, "value": value}),
                                capture_output=True, text=True, timeout=10)
        if result.returncode:
            raise RuntimeError("Native credential encryption failed")
        return result.stdout
    return encrypt


def target_snapshot():
    return json.loads(sql_run("""SELECT json_build_object(
 'database',current_database(),
 'models',(SELECT coalesce(json_agg(json_build_object('id',id,'tenant_id',tenant_id,
   'name',name,'type',type,'fingerprint',md5(to_jsonb(m)::text)) ORDER BY id),'[]') FROM models m),
 'knowledge_count',(SELECT count(*) FROM knowledges),
 'chunk_count',(SELECT count(*) FROM chunks));"""))


def transaction(query, commit=False):
    return f"""BEGIN ISOLATION LEVEL SERIALIZABLE;
SET LOCAL statement_timeout='30s';
SET LOCAL lock_timeout='5s';
CREATE TEMP TABLE _model_before ON COMMIT DROP AS SELECT * FROM models;
{query}
DO $preservation$
BEGIN
 IF EXISTS (SELECT 1 FROM _model_before b LEFT JOIN models m ON m.id=b.id
            WHERE m.id IS NULL OR to_jsonb(m) IS DISTINCT FROM to_jsonb(b))
 THEN RAISE EXCEPTION 'Model import changed an existing row'; END IF;
END $preservation$;
SELECT json_build_object('inserted_models',
 (SELECT count(*) FROM models m WHERE NOT EXISTS (SELECT 1 FROM _model_before b WHERE b.id=m.id)),
 'existing_models_unchanged',true);
{model_verification_sql()}
{'COMMIT' if commit else 'ROLLBACK'};
"""


def verify(plan, before, after, summaries):
    old = {m['id']: m for m in before['models']}
    new = {m['id']: m for m in after['models']}
    expected = {m['target_id'] for m in plan['models']}
    checks = {
        'acceptance_scope': before['database'] == after['database'] == TARGET_DATABASE,
        'existing_models_unchanged': all(new.get(k) == v for k, v in old.items()),
        'only_planned_models_added': set(new) - set(old) == expected - set(old),
        'all_models_present': expected <= set(new),
        'knowledge_and_chunks_unchanged': all(before[k] == after[k] for k in ('knowledge_count', 'chunk_count')),
        'ledger_and_encrypted_credentials': any(s.get('mapped') == len(expected)
            and len(s.get('models', [])) == len(expected)
            and all(m['encrypted_credential'] for m in s['models']) for s in summaries),
    }
    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('--environment', 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('Checkpoint exists; use a new output name')
    with consistent_snapshot(args.sqlite) as (db, info):
        plan = model_plan(db, TARGET_TENANT)
        if info['source_mutated']:
            raise RuntimeError('Source changed while taking the snapshot')
        plan['source_snapshot_sha256'] = info['sha256']
        query = model_sql(plan, read_credentials(db, plan, resolve_source_credential), native_encryptor(args.environment))
    before = target_snapshot()
    save_json(args.out.with_suffix('.before.json'), before)
    save_json(args.out.with_suffix('.manifest.json'), plan)
    rehearsal = [json.loads(line) for line in sql_run(transaction(query)).splitlines() if line.startswith('{')]
    result = {'checked_at': datetime.datetime.now(datetime.timezone.utc).isoformat(),
              'database': TARGET_DATABASE, 'tenant_id': TARGET_TENANT,
              'source_snapshot_sha256': info['sha256'], 'source_mutated': info['source_mutated'],
              'planned_models': len(plan['models']), 'skipped_models': len(plan['skipped']),
              'rehearsal': rehearsal, 'committed': False}
    if args.command == 'apply':
        save_json(args.out.with_suffix('.attempt.json'), {**result, 'commit_requested': True,
                  'recovery': 'Reconcile model ledger and committed receipt; never delete models automatically'})
        summaries = [json.loads(line) for line in sql_run(transaction(query, True)).splitlines() if line.startswith('{')]
        result.update(committed=True, changes=summaries)
        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, summaries)
    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())