ute2e/tools/review-db.py

"""Development review support: credentials stay in the local YAML, never in output."""
import argparse
import hashlib
import json
from pathlib import Path
from urllib.parse import urlparse
import datetime
import psycopg
from psycopg.rows import dict_row
from psycopg import sql
import yaml
import sqlparse
import secrets
import bcrypt

ROOT = Path(__file__).resolve().parents[2]
API = ROOT / 'unreal_tran_api'
WORK = ROOT.parent / '.codex-tmp' / 'review-20260905'
REPORT = ROOT / 'ute2e' / 'reports' / 'review-20260905'
MIGRATIONS = ['20260904_033_teaching_command_receipt.sql', '20260905_035_teaching_step_run_kingbase_constraints.sql']
PARENTS = ['ut_teaching_assignment', 'ut_teaching_run', 'ut_teaching_run_member', 'ut_teaching_event']
TABLES = ['ut_teaching_command_receipt', 'ut_teaching_step_run']

def connect():
    conf = yaml.safe_load((API / 'config/application-auth-local.yml').read_text(encoding='utf-8-sig'))
    source = conf['spring']['datasource']
    # Existing configuration can wrap the JDBC properties in Druid.
    if 'url' not in source:
        source = source['druid']
    url = urlparse(source['url'].replace('jdbc:kingbase8:', 'postgresql:', 1))
    connection = psycopg.connect(host=url.hostname, port=url.port, dbname=url.path.lstrip('/'), user=source['username'], password=source['password'], row_factory=dict_row, autocommit=True)
    with connection.cursor() as cursor:
        cursor.execute('select current_database() as db, current_schema() as schema')
        identity = cursor.fetchone()
        if identity != {'db': 'unreal_tran', 'schema': 'public'}:
            raise RuntimeError('Unexpected development database/schema; stopped')
    return connection

def apply(cursor, name):
    source = (API / 'database' / name).read_text(encoding='utf-8-sig')
    results = []
    for statement in sqlparse.split(source):
        cursor.execute(statement)
        if cursor.description:
            results.extend(cursor.fetchall())
    if any(value != 0 for row in results for value in row.values()):
        raise RuntimeError(f'Migration post-check did not return zero: {name}')
    return {'file': name, 'sha256': hashlib.sha256((API / 'database' / name).read_bytes()).hexdigest(), 'checks': results}

def isolate(cursor, schema, interrupt=False):
    if schema not in {'review_20260905_first', 'review_20260905_resume'}:
        raise RuntimeError('Unapproved temporary schema')
    cursor.execute(f'DROP SCHEMA IF EXISTS {schema} CASCADE')
    cursor.execute(f'CREATE SCHEMA {schema}')
    cursor.execute(f'SET search_path TO {schema}')
    cursor.execute('CREATE TABLE ut_teaching_assignment(id BIGINT PRIMARY KEY); CREATE TABLE ut_teaching_run(id BIGINT PRIMARY KEY, assignment_id BIGINT NOT NULL); CREATE TABLE ut_teaching_run_member(id BIGINT PRIMARY KEY, run_id BIGINT NOT NULL); CREATE TABLE ut_teaching_event(id BIGINT PRIMARY KEY, run_id BIGINT NOT NULL, assignment_id BIGINT NOT NULL)')
    if interrupt:
        # Simulate a committed interruption after the first table / first unique index.
        for name, marker in [(MIGRATIONS[0], "COMMENT='教学写命令幂等回执';"), (MIGRATIONS[1], '$$;')]:
            source = (API / 'database' / name).read_text(encoding='utf-8-sig')
            partial = source[:source.index(marker) + len(marker)]
            for statement in sqlparse.split(partial + '\nCOMMIT;'):
                cursor.execute(statement)
    first = [apply(cursor, name) for name in MIGRATIONS]
    repeat = [apply(cursor, name) for name in MIGRATIONS]
    cursor.execute('SET search_path TO public')
    cursor.execute(f'DROP SCHEMA {schema} CASCADE')
    return {'interrupted': interrupt, 'first': first, 'repeat': repeat}

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('action', choices=['prepare', 'check', 'test-admin', 'cleanup', 'clone-physical'])
    parser.add_argument('--code')
    args = parser.parse_args()
    WORK.mkdir(parents=True, exist_ok=True)
    REPORT.mkdir(parents=True, exist_ok=True)
    with connect() as db, db.cursor() as cursor:
        if args.action == 'cleanup':
            backup = json.loads((WORK / 'database-before.json').read_text(encoding='utf-8'))
            original_ids = [row['id'] for row in backup['tables']['ut_teaching_assignment']]
            cursor.execute("SELECT id,assignment_code FROM ut_teaching_assignment WHERE assignment_code LIKE 'UTE2E-REVIEW-%%' AND NOT (id = ANY(%s))", (original_ids,))
            temporary = cursor.fetchall()
            task_ids = [row['id'] for row in temporary]
            result = {'tasks': temporary, 'deleted': {}, 'time': datetime.datetime.now().isoformat()}
            with db.transaction():
                if task_ids:
                    cursor.execute('SELECT id FROM ut_teaching_run WHERE assignment_id = ANY(%s)', (task_ids,))
                    run_ids = [row['id'] for row in cursor.fetchall()]
                    # Every affected row is tied to a test task absent from the recovery snapshot.
                    for table, column, ids in [
                        ('ut_teaching_step_run','assignment_id',task_ids),
                        ('ut_teaching_command_receipt','assignment_id',task_ids),
                        ('ut_teaching_xr_session','assignment_id',task_ids),
                        ('ut_assistance_alert','assignment_id',task_ids),
                        ('ut_assistance_evaluation','assignment_id',task_ids),
                        ('ut_data_observation','run_id',run_ids),
                        ('ut_teaching_event','assignment_id',task_ids),
                        ('ut_teaching_intervention','assignment_id',task_ids),
                        ('ut_teaching_assignment_fault','assignment_id',task_ids),
                        ('ut_teaching_run_member','run_id',run_ids),
                        ('ut_teaching_run','assignment_id',task_ids),
                        ('ut_teaching_member','assignment_id',task_ids),
                        ('ut_teaching_assignment','id',task_ids),
                    ]:
                        cursor.execute(sql.SQL('DELETE FROM {} WHERE {} = ANY(%s)').format(sql.Identifier(table),sql.Identifier(column)), (ids,))
                        result['deleted'][table] = cursor.rowcount
                since = int(datetime.datetime.fromisoformat(backup['time']).timestamp())
                cursor.execute("SELECT id FROM ut_sys_user WHERE remark LIKE 'UTE2E-REVIEW-%%' AND username LIKE 'rv_%%' AND add_time >= %s", (since,))
                user_ids = [row['id'] for row in cursor.fetchall()]
                temp_admin = WORK / 'test-admin.json'
                if temp_admin.exists():
                    admin = json.loads(temp_admin.read_text(encoding='utf-8'))
                    cursor.execute('SELECT id FROM ut_sys_user WHERE id=%s AND username=%s AND is_protected=0', (admin['id'],admin['username']))
                    matched = cursor.fetchone()
                    if matched:
                        user_ids.append(matched['id'])
                cursor.execute("UPDATE ut_auth_session SET is_revoked=1,revoked_time=UNIX_TIMESTAMP(),revoke_reason='review cleanup' WHERE user_id = ANY(%s)", (user_ids,))
                for table in ['ut_sys_user_role','ut_auth_user_credential','ut_sys_user']:
                    column = 'id' if table == 'ut_sys_user' else 'user_id'
                    extra = ', status=0, token_version=token_version+1' if table == 'ut_sys_user' else ''
                    cursor.execute(sql.SQL('UPDATE {} SET is_delete=1,delete_token=id' + extra + ' WHERE {} = ANY(%s) AND is_delete=0').format(sql.Identifier(table),sql.Identifier(column)), (user_ids,))
                result['testUsersDisabled'] = user_ids
            result['originalRowsUnchanged'] = {}
            for table, rows in backup['tables'].items():
                ids = [row['id'] for row in rows]
                cursor.execute(sql.SQL('SELECT * FROM {} WHERE id = ANY(%s)').format(sql.Identifier(table)), (ids,))
                after = cursor.fetchall()
                normalize = lambda values: json.dumps(sorted(values,key=lambda row:row['id']),sort_keys=True,default=str)
                result['originalRowsUnchanged'][table] = normalize(rows) == normalize(after)
            cursor.execute("SELECT COUNT(*) AS count FROM ut_sys_user WHERE username LIKE 'rv_%%' AND remark LIKE 'UTE2E-REVIEW-%%' AND is_delete=0 AND add_time >= %s", (since,))
            result['remainingActiveTestUsers'] = cursor.fetchone()['count']
            (REPORT / 'cleanup.json').write_text(json.dumps(result,ensure_ascii=False,indent=2,default=str),encoding='utf-8')
            if not all(result['originalRowsUnchanged'].values()) or result['remainingActiveTestUsers'] != 0:
                raise RuntimeError('Cleanup verification failed; inspect cleanup.json before continuing')
            if temp_admin.exists():
                temp_admin.unlink()
            print(json.dumps(result,ensure_ascii=True))
            return
        if args.action == 'clone-physical':
            if not args.code or not args.code.startswith('UTE2E-REVIEW-') or not args.code.endswith('-PHYSICAL'):
                raise RuntimeError('Invalid isolated fixture prefix')
            cursor.execute("SELECT * FROM ut_teaching_assignment WHERE assignment_code='UAT-TEACH-PHYSICAL-001' AND is_delete=0")
            record = cursor.fetchone()
            record.pop('id')
            record.update(assignment_code=args.code, assignment_name=args.code, audience_type='COMMON', collaboration_mode='INDIVIDUAL', status_code='PUBLISHED', data_version=0, schedule_start_at=None, due_at=None, digital_human_allowed=0)
            definition = json.loads(record['definition_snapshot'])
            steps = definition['steps']
            for index, step in enumerate(steps):
                step['physicalEventRule'] = {'source':'MANUAL','eventType':'review.step.confirmed','progressPercent':round((index+1)*100/len(steps))}
            definition['modeConfig'] = {'type':'PHYSICAL','manualConfirmAllowed':True}
            record['definition_snapshot'] = json.dumps(definition,ensure_ascii=False,separators=(',',':'))
            record['definition_checksum'] = hashlib.sha256(record['definition_snapshot'].encode()).hexdigest()
            cursor.execute(sql.SQL('INSERT INTO ut_teaching_assignment ({}) VALUES ({}) RETURNING id').format(sql.SQL(',').join(map(sql.Identifier, record)),sql.SQL(',').join(sql.Placeholder() for _ in record)),list(record.values()))
            result = cursor.fetchone()
            print(json.dumps({'id':str(result['id']),'code':args.code,'version':0}))
            return
        if args.action == 'test-admin':
            username = 'rv_admin_' + secrets.token_hex(4)
            password = 'Rv!' + secrets.token_hex(7)
            with db.transaction():
                cursor.execute("SELECT id,department_id FROM ut_sys_user WHERE is_delete=0 AND is_protected=1 LIMIT 1")
                department = cursor.fetchone()['department_id']
                cursor.execute("SELECT id FROM ut_sys_role WHERE role_code='admin' AND status=1 AND is_delete=0")
                role = cursor.fetchone()['id']
                cursor.execute("INSERT INTO ut_sys_user(username,username_normalized,display_name,department_id,default_role_id,status,is_protected,must_change_password,remark) VALUES(%s,%s,%s,%s,%s,1,0,0,%s) RETURNING id", (username,username,'UTE2E-REVIEW-ADMIN',department,role,'Temporary review administrator'))
                user_id = cursor.fetchone()['id']
                cursor.execute('INSERT INTO ut_sys_user_role(user_id,role_id) VALUES(%s,%s)', (user_id,role))
                digest = '{bcrypt}' + bcrypt.hashpw(password.encode(), bcrypt.gensalt()).decode()
                cursor.execute('INSERT INTO ut_auth_user_credential(user_id,password_hash) VALUES(%s,%s)', (user_id,digest))
            (WORK / 'test-admin.json').write_text(json.dumps({'id':user_id,'username':username,'password':password}),encoding='utf-8')
            print('Created isolated temporary review administrator; credentials stored only in local scratch file')
            return
        if args.action == 'prepare':
            cursor.execute("SELECT table_name FROM information_schema.tables WHERE table_schema='public' AND table_name = ANY(%s)", (TABLES,))
            if cursor.fetchall():
                raise RuntimeError('Target tables already exist; inspect instead of reapplying')
            backup = {'time': datetime.datetime.now().isoformat(), 'tables': {}}
            for name in PARENTS:
                cursor.execute(f'SELECT * FROM {name}')
                backup['tables'][name] = cursor.fetchall()
            cursor.execute("SELECT tablename,indexname,indexdef FROM pg_indexes WHERE schemaname='public' AND tablename = ANY(%s)", (PARENTS,))
            backup['indexes'] = cursor.fetchall()
            cursor.execute("SELECT table_name,column_name,data_type,column_default,is_nullable FROM information_schema.columns WHERE table_schema='public' AND table_name = ANY(%s) ORDER BY table_name,ordinal_position", (PARENTS,))
            backup['columns'] = cursor.fetchall()
            (WORK / 'database-before.json').write_text(json.dumps(backup, ensure_ascii=False, indent=2, default=str), encoding='utf-8')
            result = {'environment': 'development unreal_tran/public', 'time': backup['time'], 'backup': str(WORK / 'database-before.json'), 'validation': []}
            for schema, interrupted in [('review_20260905_first', False), ('review_20260905_resume', True)]:
                result['validation'].append(isolate(cursor, schema, interrupted))
            result['executed'] = [apply(cursor, name) for name in MIGRATIONS]
            (REPORT / 'database-migration.json').write_text(json.dumps(result, ensure_ascii=False, indent=2), encoding='utf-8')
            print('PASS: recovery snapshot, isolated first/repeat/interruption recovery, development migrations 033/035')
        for name in TABLES:
            cursor.execute(f'SELECT COUNT(*) AS count FROM {name}')
            print(name, cursor.fetchone())

if __name__ == '__main__':
    main()