kangshifu-board/scripts/migrate-unified-database.py

"""Plan, back up, migrate and verify the current library in the existing Parse DB.

No credentials are accepted on the command line. The operator supplies a private
libpq service/passfile configuration. Nothing switches a running service here.
"""
import argparse
from contextlib import closing
import json
import os
from pathlib import Path
import re
import sqlite3
import sys
sys.dont_write_bytecode = True
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / 'server/voc-library'))
import unified_store as store


def write(path, body):
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
    with path.open('x', encoding='utf-8') as f:
        f.write(json.dumps(body, ensure_ascii=False, indent=2) + '\n')
    path.chmod(0o600)


def source_snapshot(registry, database):
    store.need(bool(registry) != bool(database), 'ONE_CURRENT_LIBRARY_BINDING_REQUIRED')
    if registry:
        root = Path(registry).resolve()
        pointer_raw = (root / 'current.json').read_bytes()
        pointer = json.loads(pointer_raw)
        path = (root / pointer['file']).resolve()
        store.need(path.is_relative_to(root) and path.is_file(), 'INVALID_SOURCE_POINTER')
        return path, pointer, store.sha(pointer_raw)
    path = Path(database).resolve()
    store.need(path.is_file(), 'SOURCE_DATABASE_REQUIRED')
    return path, None, None


def backup(source, destination):
    store.need(not destination.exists(), 'BACKUP_TARGET_EXISTS')
    with closing(sqlite3.connect(source.as_uri() + '?mode=ro', uri=True)) as raw:
        with closing(sqlite3.connect(destination)) as saved:
            raw.backup(saved)
            store.need(saved.execute('PRAGMA integrity_check').fetchone()[0] == 'ok', 'BACKUP_INVALID')
    destination.chmod(0o600)


def verify(config, sources=None):
    results = []
    for dataset in ('research', 'monitoring'):
        with store.connect_pg(config, readonly=True) as pg:
            desc = store.descriptor(pg, dataset)
            c = store.reconstruct(pg, desc)
            try:
                if sources:
                    with closing(sqlite3.connect(sources[dataset].as_uri() + '?mode=ro', uri=True)) as source:
                        store.compare_manifest(store.inventory(source), store.inventory(c))
                results.append({'dataset': dataset, 'generation': desc['generation'], 'revision': desc['revision'],
                    'fingerprint': desc['manifest']['fingerprint'], 'tables': len(desc['manifest']['tables']),
                    'rows': sum(t['count'] for t in desc['manifest']['tables']), 'allRowsVerified': True})
            finally:
                c.close()
    archived = []
    with store.connect_pg(config, readonly=True) as pg:
        active = store.descriptor(pg, 'research')
        for generation, storage, manifest in pg.execute('SELECT generation,storage_schema,manifest FROM '
                'ksf_voc_control.snapshots WHERE dataset=\'research\' AND storage_schema<>%s ORDER BY generation', (active['storage_schema'],)).fetchall():
            c = store.reconstruct(pg, {'storage_schema': storage, 'manifest': manifest});c.close()
            archived.append({'generation': generation, 'fingerprint': manifest['fingerprint'], 'allRowsVerified': True})
    return {'ok': True, 'schema': store.VERSION, 'database': store.DATABASE, 'datasets': results,
            'archivedResearchGenerations': archived,
            'readOnly': True, 'databaseWrites': 0, 'businessRowsChanged': 0}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('command', choices=('plan', 'apply', 'verify'))
    parser.add_argument('--config', required=True, help='Private unified connection configuration')
    source = parser.add_mutually_exclusive_group()
    source.add_argument('--registry')
    source.add_argument('--db')
    parser.add_argument('--monitoring-store')
    parser.add_argument('--expected-generation')
    parser.add_argument('--out', required=True, help='NEW private receipt/backup directory, outside source directories')
    args = parser.parse_args()
    applying = args.command == 'apply'
    try:
        output = Path(args.out).absolute()
        store.need(not output.exists() and not output.is_symlink(), 'NEW_OUTPUT_DIRECTORY_REQUIRED')
        if args.command == 'verify':
            result = verify(args.config)
        else:
            db, pointer, pointer_sha = source_snapshot(args.registry, args.db)
            monitor = Path(args.monitoring_store or '').resolve() / 'monitoring.sqlite'
            store.need(monitor.is_file(), 'CURRENT_MONITORING_DATABASE_REQUIRED')
            store.need(not any(output == p or output.is_relative_to(p) for p in (db.parent, monitor.parent)), 'OUTPUT_OVERLAPS_SOURCE')
            with closing(sqlite3.connect(db.as_uri() + '?mode=ro', uri=True)) as c:
                generation = dict(c.execute('SELECT key,value FROM meta'))['generation']
            store.need(args.expected_generation == generation, 'CURRENT_GENERATION_CONFIRMATION_REQUIRED')
            # Read-only target identity and both source inventories before any import.
            plans = [store.migrate(db, args.config, registry_pointer=pointer),
                     store.migrate(monitor, args.config, 'monitoring')]
            history = []
            seen = {generation}
            for entry in (pointer or {}).get('history', []):
                if entry['generation'] in seen:
                    continue
                path = (Path(args.registry).resolve() / entry['file']).resolve()
                store.need(path.is_relative_to(Path(args.registry).resolve()), 'INVALID_HISTORICAL_PATH')
                entry = {**entry, 'history': pointer.get('history', [])}
                plan = store.archive_research(path, args.config, entry)
                history.append((path, entry, plan))
                seen.add(entry['generation'])
            result = {'ok': True, 'schema': store.VERSION, 'readOnly': not applying,
                      'database': store.DATABASE, 'sourceGeneration': generation,
                      'sourcePointerSha256': pointer_sha, 'plans': plans,
                      'runningServicesSwitched': False, 'sourceDeleted': False, 'businessRowsChanged': 0}
            result['historicalGenerations'] = [p for _, _, p in history]
            if applying:
                output.mkdir(parents=True, mode=0o700)
                backup(db, output / 'research-before.sqlite')
                backup(monitor, output / 'monitoring-before.sqlite')
                if pointer:
                    write(output / 'pointer-before.json', pointer)
                # Backups are independent and include any SQLite WAL contents.
                for name, path, plan in [('research', output / 'research-before.sqlite', plans[0]),
                                         ('monitoring', output / 'monitoring-before.sqlite', plans[1])]:
                    with closing(sqlite3.connect(path)) as c:
                        store.need(store.inventory(c)['fingerprint'] == plan['fingerprint'], 'SOURCE_CHANGED_BEFORE_BACKUP')
                result['imports'] = [store.migrate(output / 'research-before.sqlite', args.config,
                                                  registry_pointer=pointer, apply=True),
                                     store.migrate(output / 'monitoring-before.sqlite', args.config,
                                                   'monitoring', apply=True)]
                result['historyImports'] = []
                for path, entry, plan in history:
                    imported = store.archive_research(path, args.config, entry, apply=True)
                    store.need(imported['fingerprint'] == plan['fingerprint'], 'HISTORICAL_SOURCE_CHANGED')
                    result['historyImports'].append(imported)
                if args.registry:
                    store.need(store.sha((Path(args.registry) / 'current.json').read_bytes()) == pointer_sha,
                               'SOURCE_GENERATION_CHANGED_DURING_IMPORT')
                # The operator pauses only source ingestion while taking and switching this baseline.
                result['verification'] = verify(args.config, {'research': db, 'monitoring': monitor})
                result['readyForServiceCutover'] = True
        write(output / 'migration-receipt.json', result)
        print(json.dumps({'ok': True, 'command': args.command, 'database': store.DATABASE,
                          'receipt': str(output / 'migration-receipt.json'),
                          'readyForServiceCutover': result.get('readyForServiceCutover', False)}, ensure_ascii=False))
        return 0
    except Exception as e:
        code = str(e) if isinstance(e, store.StoreError) and re.fullmatch('[A-Z][A-Z0-9_]+', str(e)) else 'MIGRATION_FAILED'
        print(json.dumps({'ok': False, 'error': {'code': code}, 'runningServicesSwitched': False,
                          'sourceDeleted': False, 'targetImportsMayExist': applying,
                          'note': 'Inspect the private backup and receipt directory; no credentials or raw error printed.'}))
        return 1


if __name__ == '__main__':
    sys.exit(main())