xinmi/scripts/apply_weknora_historical_identity.py

#!/usr/bin/env python3
"""Add the deleted historical owner's disabled identity with rollback rehearsal."""
import argparse
import json
from pathlib import Path

from apply_weknora_identities import sql_run, target_snapshot
from weknora_category_etl import save_json
from weknora_etl import consistent_snapshot
from weknora_historical_identity_etl import historical_identity_plan, historical_identity_sql
from weknora_identity_etl import TARGET_TENANT


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--sqlite', 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('.before.json'), args.out.with_suffix('.attempt.json')]):
        parser.error('Checkpoint exists')
    with consistent_snapshot(args.sqlite) as (db, info):
        plan = historical_identity_plan(db, TARGET_TENANT)
        if info['source_mutated']:
            raise RuntimeError('Source changed during snapshot')
        plan['source_snapshot_sha256'] = info['sha256']
        query = historical_identity_sql(plan)
    before = target_snapshot()
    save_json(args.out.with_suffix('.before.json'), before)
    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';
CREATE TEMP TABLE _old_users ON COMMIT DROP AS SELECT * FROM users;
CREATE TEMP TABLE _old_members ON COMMIT DROP AS SELECT * FROM tenant_members;
""" + query + """
DO $verify$ BEGIN
 IF EXISTS(SELECT 1 FROM _old_users b LEFT JOIN users a ON a.id=b.id WHERE a.id IS NULL OR to_jsonb(a) IS DISTINCT FROM to_jsonb(b))
 OR EXISTS(SELECT 1 FROM _old_members b LEFT JOIN tenant_members a ON a.id=b.id WHERE a.id IS NULL OR to_jsonb(a) IS DISTINCT FROM to_jsonb(b))
 THEN RAISE EXCEPTION 'Historical identity changed existing users'; END IF;
END $verify$;
"""
    sql_run(sql + 'ROLLBACK;')
    if before != target_snapshot():
        raise RuntimeError('Rollback rehearsal changed identities')
    save_json(args.out.with_suffix('.attempt.json'), {'source_snapshot_sha256': info['sha256'], 'commit_requested': True})
    sql_run(sql + 'COMMIT;')
    after = target_snapshot()
    uid = plan['users'][0]['target_id']
    user = next(u for u in after['users'] if u['id'] == uid)
    members = [m for m in after['memberships'] if m['user_id'] == uid]
    tokens = [t for t in after['token_counts'] if t['user_id'] == uid]
    checks = {'source_not_mutated': not info['source_mutated'], 'login_disabled': not user['is_active'],
              'not_admin': not user['is_system_admin'] and not user['can_access_all_tenants'],
              'membership_suspended': len(members) == 1 and members[0]['tenant_id'] == TARGET_TENANT and members[0]['status'] == 'suspended',
              'no_tokens': not tokens}
    sql_run(sql + 'COMMIT;')
    checks['repeat_no_change'] = after == target_snapshot()
    result = {'committed': True, 'source_user_id': 2, 'target_id': uid,
              'passed': all(checks.values()), 'checks': checks}
    save_json(args.out, result)
    print(json.dumps(result))
    return 0 if result['passed'] else 1


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