280 lines
17 KiB
Python
280 lines
17 KiB
Python
"""Resumable, unpublished COPY followed by complete ordered content verification."""
|
|
import argparse
|
|
from contextlib import closing
|
|
import fcntl
|
|
import hashlib
|
|
import json
|
|
from pathlib import Path
|
|
import re
|
|
import shutil
|
|
import sqlite3
|
|
import uuid
|
|
import zlib
|
|
|
|
from paper_atlas.db.submissions import timestamp
|
|
from paper_atlas.postgres.client import connect, SEARCH_VECTOR
|
|
from paper_atlas.postgres.schema import TABLES, inspect, fingerprint, quote, select_sql, table_sql
|
|
from paper_atlas.postgres import text
|
|
|
|
|
|
def digest_file(path):
|
|
digest = hashlib.sha256()
|
|
with path.open('rb') as stream:
|
|
while block := stream.read(8 * 1024 ** 2):
|
|
digest.update(block)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def update_hash(digest, row):
|
|
# Raw JSON columns remain text: whitespace, escapes and source evidence survive.
|
|
data = json.dumps(tuple(row), ensure_ascii=False, separators=(',', ':')).encode()
|
|
digest.update(len(data).to_bytes(8, 'big')); digest.update(data)
|
|
|
|
|
|
def save(path, result):
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
part = path.with_suffix(path.suffix + '.part')
|
|
part.write_text(json.dumps(result, ensure_ascii=False, indent=2) + '\n')
|
|
part.replace(path)
|
|
|
|
|
|
def free_space(path, minimum):
|
|
if path and shutil.disk_usage(path).free < minimum * 1024 ** 3:
|
|
raise RuntimeError('Migration storage headroom fell below the configured minimum')
|
|
|
|
|
|
def qualify(schema, table):
|
|
return quote(schema) + '.' + quote(table)
|
|
|
|
|
|
def add_constraints(pg, stage, plan):
|
|
def constraint(table, name, clause):
|
|
exists = pg.execute('''SELECT 1 FROM pg_constraint co JOIN pg_class cl ON cl.oid=co.conrelid
|
|
JOIN pg_namespace ns ON ns.oid=cl.relnamespace WHERE ns.nspname=%s AND cl.relname=%s AND co.conname=%s''',
|
|
(stage, table, name)).fetchone()
|
|
if not exists:
|
|
pg.execute('ALTER TABLE ' + qualify(stage, table) + ' ADD CONSTRAINT ' + quote(name) + ' ' + clause)
|
|
|
|
for table, spec in plan.items():
|
|
constraint(table, table + '_pkey', 'PRIMARY KEY (' + ','.join(map(quote, spec['pk'])) + ')')
|
|
for number, index in enumerate(spec['indexes']):
|
|
if index['origin'] == 'u':
|
|
constraint(table, table + '_unique_' + str(number), 'UNIQUE (' + ','.join(quote(n) for n, desc in index['terms']) + ')')
|
|
else:
|
|
terms = ','.join(quote(n) + (' DESC' if desc else '') for n, desc in index['terms'])
|
|
pg.execute('CREATE ' + ('UNIQUE ' if index['unique'] else '') + 'INDEX IF NOT EXISTS ' +
|
|
quote(index['name']) + ' ON ' + qualify(stage, table) + '(' + terms + ')')
|
|
for table, spec in plan.items():
|
|
for number, fk in enumerate(spec['foreign_keys']):
|
|
clause = ('FOREIGN KEY (' + ','.join(map(quote, fk['columns'])) + ') REFERENCES ' +
|
|
qualify(stage, fk['table']) + '(' + ','.join(map(quote, fk['references'])) + ') NOT VALID')
|
|
constraint(table, table + '_fk_' + str(number), clause)
|
|
for number, check in enumerate(spec['checks']):
|
|
constraint(table, table + '_check_' + str(number), 'CHECK (' + check + ') NOT VALID')
|
|
constraints = pg.execute('''SELECT cl.relname,co.conname FROM pg_constraint co
|
|
JOIN pg_class cl ON cl.oid=co.conrelid JOIN pg_namespace ns ON ns.oid=cl.relnamespace
|
|
WHERE ns.nspname=%s AND NOT co.convalidated ORDER BY cl.relname,co.conname''', (stage,)).fetchall()
|
|
for table, name in constraints:
|
|
pg.execute('ALTER TABLE ' + qualify(stage, table) + ' VALIDATE CONSTRAINT ' + quote(name))
|
|
|
|
|
|
def sequences(pg, source, stage, plan):
|
|
old = dict(source.execute('SELECT name,seq FROM sqlite_sequence'))
|
|
history = {'papers': ('paper_merges', 'former_paper_id'),
|
|
'submissions': ('submission_merges', 'former_submission_id'),
|
|
'decisions': ('decision_merges', 'former_decision_id')}
|
|
result = {}
|
|
for table, spec in plan.items():
|
|
if spec['pk'] != ['id']:
|
|
continue
|
|
value = max(source.execute('SELECT coalesce(max(id),0) FROM ' + quote(table)).fetchone()[0], old.get(table, 0))
|
|
if table in history:
|
|
audit, column = history[table]
|
|
value = max(value, source.execute('SELECT coalesce(max(' + quote(column) + '),0) FROM ' + quote(audit)).fetchone()[0])
|
|
sequence = pg.execute('SELECT pg_get_serial_sequence(%s,%s)', (qualify(stage, table), 'id')).fetchone()[0]
|
|
pg.execute('SELECT setval(%s,%s,%s)', (sequence, max(value, 1), value > 0))
|
|
result[table] = value
|
|
return result
|
|
|
|
|
|
def upgrade_text_codec(pg, control, stage, plan):
|
|
"""Upgrade committed legacy staging rows once; never recopy their source."""
|
|
changed = {}
|
|
with pg.transaction():
|
|
copied = {r[0] for r in pg.execute('SELECT name FROM ' + qualify(control, 'tables'))}
|
|
if not copied <= plan.keys():
|
|
raise ValueError('Unknown committed staging tables')
|
|
# An older run may have reached constraint validation. Replacing SOH
|
|
# can temporarily equal another legacy key even though the final mapping
|
|
# is bijective. Rebuild unpublished constraints/indexes after the upgrade.
|
|
constraints = pg.execute('''SELECT cl.relname,co.conname FROM pg_constraint co
|
|
JOIN pg_class cl ON cl.oid=co.conrelid JOIN pg_namespace ns ON ns.oid=cl.relnamespace
|
|
WHERE ns.nspname=%s AND co.contype IN ('f','p','u')
|
|
ORDER BY CASE WHEN co.contype='f' THEN 0 ELSE 1 END''', (stage,)).fetchall()
|
|
for table, name in constraints:
|
|
pg.execute('ALTER TABLE ' + qualify(stage, table) + ' DROP CONSTRAINT ' + quote(name))
|
|
indexes = pg.execute('''SELECT cl.relname FROM pg_index i JOIN pg_class cl ON cl.oid=i.indexrelid
|
|
JOIN pg_namespace ns ON ns.oid=cl.relnamespace WHERE ns.nspname=%s AND i.indisunique''', (stage,)).fetchall()
|
|
for name, in indexes:
|
|
pg.execute('DROP INDEX ' + qualify(stage, name))
|
|
for table in plan:
|
|
if table not in copied:
|
|
continue
|
|
columns = [quote(c['name']) for c in plan[table]['columns'] if c['type'] == 'TEXT']
|
|
if columns:
|
|
assignments = ','.join(c + '=replace(' + c + ',%s,%s)' for c in columns)
|
|
predicate = ' OR '.join('strpos(' + c + ',%s)>0' for c in columns)
|
|
params = [v for c in columns for v in (text.ESCAPE, text.encode(text.ESCAPE))]
|
|
params += [text.ESCAPE] * len(columns)
|
|
changed[table] = pg.execute('UPDATE ' + qualify(stage, table) + ' SET ' + assignments +
|
|
' WHERE ' + predicate, params).rowcount
|
|
pg.execute('UPDATE ' + qualify(control, 'run') + ' SET text_codec=%s WHERE id=1', (text.CODEC,))
|
|
return changed
|
|
|
|
|
|
def migrate(db, service, out, *, schema='research', storage_path=None, min_free_gib=8):
|
|
if not re.fullmatch(r'[a-z][a-z0-9_]{0,35}', schema):
|
|
raise ValueError('Use a short lowercase target schema name')
|
|
db, out = Path(db).resolve(), Path(out).resolve()
|
|
if not db.is_file() or out == db:
|
|
raise ValueError('Existing source and a distinct report are required')
|
|
if out.exists() and json.loads(out.read_text()).get('status') == 'complete':
|
|
raise ValueError('Completed report must be preserved')
|
|
free_space(storage_path, min_free_gib)
|
|
control = schema + '_migration'
|
|
report = {'status': 'preflight', 'source': str(db), 'service': service, 'schema': schema,
|
|
'started_at': timestamp(), 'tables': {}, 'scope': 'All 21 business tables; FTS5 internals rebuilt as native PostgreSQL GIN.'}
|
|
with Path(str(db) + '.maintenance.lock').open('a') as lock:
|
|
fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
|
|
if Path(str(db) + '-wal').exists() and Path(str(db) + '-wal').stat().st_size:
|
|
raise ValueError('Freeze writers and checkpoint SQLite WAL before migrating')
|
|
with closing(sqlite3.connect(db.as_uri() + '?mode=ro', uri=True)) as source, closing(connect(service)) as pg:
|
|
source.execute('PRAGMA query_only=ON'); source.execute('PRAGMA cache_size=-16384'); source.execute('BEGIN')
|
|
plan = inspect(source)
|
|
stat = db.stat()
|
|
identity = {'file_sha256': digest_file(db), 'bytes': stat.st_size, 'mtime_ns': stat.st_mtime_ns,
|
|
'schema_sha256': fingerprint(plan)}
|
|
report['source_identity'] = identity
|
|
advisory = zlib.crc32(schema.encode()) & 0x7fffffff
|
|
if not pg.execute('SELECT pg_try_advisory_lock(%s,%s)', (615173, advisory)).fetchone()[0]:
|
|
raise RuntimeError('Another migration holds the target lock')
|
|
if pg.execute('SELECT 1 FROM pg_namespace WHERE nspname=%s', (schema,)).fetchone():
|
|
raise ValueError('Target schema already exists; published data will never be overwritten')
|
|
pg.execute('CREATE SCHEMA IF NOT EXISTS ' + quote(control))
|
|
pg.execute('CREATE TABLE IF NOT EXISTS ' + qualify(control, 'run') +
|
|
'(id INTEGER PRIMARY KEY CHECK(id=1), source_json TEXT NOT NULL, stage TEXT NOT NULL, status TEXT NOT NULL)')
|
|
pg.execute('CREATE TABLE IF NOT EXISTS ' + qualify(control, 'tables') +
|
|
'(name TEXT PRIMARY KEY, rows BIGINT NOT NULL, sha256 TEXT NOT NULL, verified BOOLEAN NOT NULL DEFAULT false)')
|
|
pg.execute('ALTER TABLE ' + qualify(control, 'run') + " ADD COLUMN IF NOT EXISTS text_codec TEXT NOT NULL DEFAULT 'legacy'")
|
|
run = pg.execute('SELECT source_json,stage,status,text_codec FROM ' + qualify(control, 'run') + ' WHERE id=1').fetchone()
|
|
if run:
|
|
if json.loads(run[0]) != identity or run[2] == 'complete':
|
|
raise ValueError('Resume requires the exact same unpublished SQLite snapshot')
|
|
stage = run[1]
|
|
if run[3] == 'legacy':
|
|
report['legacy_text_upgrade_rows'] = upgrade_text_codec(pg, control, stage, plan)
|
|
elif run[3] != text.CODEC:
|
|
raise ValueError('Unsupported staging text storage codec')
|
|
else:
|
|
stage = schema + '_stage_' + uuid.uuid4().hex[:12]
|
|
with pg.transaction():
|
|
pg.execute('CREATE SCHEMA ' + quote(stage))
|
|
for table, spec in plan.items():
|
|
pg.execute(table_sql(stage, table, spec))
|
|
pg.execute('INSERT INTO ' + qualify(control, 'run') +
|
|
'(id,source_json,stage,status,text_codec) VALUES(1,%s,%s,%s,%s)',
|
|
(json.dumps(identity, sort_keys=True), stage, 'copying', text.CODEC))
|
|
report.update(status='copying', stage=stage, text_codec=text.CODEC,
|
|
content_hash_scope='Decoded original SQLite values; storage escapes are reversible.')
|
|
save(out, report)
|
|
try:
|
|
for table, spec in plan.items():
|
|
free_space(storage_path, min_free_gib)
|
|
previous = pg.execute('SELECT rows,sha256 FROM ' + qualify(control, 'tables') + ' WHERE name=%s', (table,)).fetchone()
|
|
if previous:
|
|
rows, digest = previous
|
|
else:
|
|
digest = hashlib.sha256(); rows = 0
|
|
with pg.transaction():
|
|
# Only this run's unpublished, incomplete table can be reset.
|
|
pg.execute('TRUNCATE ' + qualify(stage, table))
|
|
columns = ','.join(quote(c['name']) for c in spec['columns'])
|
|
with pg.cursor().copy('COPY ' + qualify(stage, table) + '(' + columns + ') FROM STDIN') as copy:
|
|
for row in source.execute(select_sql(None, table, spec)):
|
|
copy.write_row(tuple(map(text.encode, row))); update_hash(digest, row); rows += 1
|
|
if rows % 10000 == 0:
|
|
free_space(storage_path, min_free_gib)
|
|
if rows % 100000 == 0:
|
|
print(json.dumps({'stage': 'copying', 'table': table, 'rows': rows}), flush=True)
|
|
digest = digest.hexdigest()
|
|
pg.execute('INSERT INTO ' + qualify(control, 'tables') + '(name,rows,sha256) VALUES(%s,%s,%s)', (table, rows, digest))
|
|
report['tables'][table] = {'source_rows': rows, 'source_sha256': digest}
|
|
save(out, report)
|
|
print(json.dumps({'stage': 'copied', 'table': table, 'rows': rows}), flush=True)
|
|
report['status'] = 'validating_constraints'; save(out, report)
|
|
add_constraints(pg, stage, plan)
|
|
report['sequences'] = sequences(pg, source, stage, plan)
|
|
report['status'] = 'verifying_contents'; save(out, report)
|
|
for table, spec in plan.items():
|
|
digest = hashlib.sha256(); storage_digest = hashlib.sha256(); rows = 0
|
|
with pg.transaction(), pg.cursor(name='verify_' + table) as cursor:
|
|
cursor.itersize = 512
|
|
cursor.execute(select_sql(stage, table, spec))
|
|
for row in cursor:
|
|
update_hash(storage_digest, row)
|
|
update_hash(digest, tuple(map(text.decode, row))); rows += 1
|
|
entry = report['tables'][table]
|
|
entry.update(target_rows=rows, target_sha256=digest.hexdigest(),
|
|
target_storage_sha256=storage_digest.hexdigest())
|
|
if rows != entry['source_rows'] or digest.hexdigest() != entry['source_sha256']:
|
|
raise ValueError('Complete content verification failed for ' + table)
|
|
pg.execute('UPDATE ' + qualify(control, 'tables') + ' SET verified=true WHERE name=%s', (table,))
|
|
save(out, report)
|
|
print(json.dumps({'stage': 'verified', 'table': table, 'rows': rows}), flush=True)
|
|
report['status'] = 'building_search_index'; save(out, report)
|
|
free_space(storage_path, min_free_gib)
|
|
pg.execute('CREATE INDEX IF NOT EXISTS papers_search_gin ON ' + qualify(stage, 'papers') + ' USING gin (' + SEARCH_VECTOR + ')')
|
|
search = pg.execute('''SELECT am.amname,i.indisvalid,i.indisready,t.relname
|
|
FROM pg_index i JOIN pg_class ix ON ix.oid=i.indexrelid
|
|
JOIN pg_class t ON t.oid=i.indrelid JOIN pg_namespace ns ON ns.oid=ix.relnamespace
|
|
JOIN pg_am am ON am.oid=ix.relam WHERE ns.nspname=%s AND ix.relname='papers_search_gin' ''', (stage,)).fetchone()
|
|
if search != ('gin', True, True, 'papers'):
|
|
raise ValueError('Native full-text index is not valid and ready')
|
|
for table in TABLES:
|
|
pg.execute('ANALYZE ' + qualify(stage, table))
|
|
invalid = pg.execute('''SELECT count(*) FROM pg_constraint co JOIN pg_class cl ON cl.oid=co.conrelid
|
|
JOIN pg_namespace ns ON ns.oid=cl.relnamespace WHERE ns.nspname=%s AND NOT co.convalidated''', (stage,)).fetchone()[0]
|
|
if invalid:
|
|
raise ValueError('Unvalidated constraints remain')
|
|
if (db.stat().st_size, db.stat().st_mtime_ns) != (identity['bytes'], identity['mtime_ns']):
|
|
raise ValueError('Source changed during migration')
|
|
if Path(str(db) + '-wal').exists() and Path(str(db) + '-wal').stat().st_size:
|
|
raise ValueError('A source writer committed during migration')
|
|
report.update(status='prepared', foreign_keys='all_validated', checks='all_validated', native_search_index=True)
|
|
save(out, report)
|
|
with pg.transaction():
|
|
pg.execute('COMMENT ON SCHEMA ' + quote(stage) + " IS 'paper-atlas:text-codec=" + text.CODEC + "'")
|
|
pg.execute('ALTER SCHEMA ' + quote(stage) + ' RENAME TO ' + quote(schema))
|
|
pg.execute('UPDATE ' + qualify(control, 'run') + " SET status='complete' WHERE id=1")
|
|
report.update(status='complete', finished_at=timestamp())
|
|
save(out, report)
|
|
return report
|
|
except BaseException as error:
|
|
report.update(status='error', error_type=type(error).__name__, error=str(error), failed_at=timestamp())
|
|
save(out, report)
|
|
raise
|
|
|
|
|
|
def main(argv=None, *, prog=None):
|
|
parser = argparse.ArgumentParser(prog=prog)
|
|
parser.add_argument('--db', required=True, help='Frozen SQLite source')
|
|
parser.add_argument('--pg-service', required=True, help='Protected libpq service name')
|
|
parser.add_argument('--schema', default='research')
|
|
parser.add_argument('--out', required=True)
|
|
parser.add_argument('--storage-path', help='Target filesystem path for free-space checks')
|
|
parser.add_argument('--min-free-gib', type=float, default=8)
|
|
args = parser.parse_args(argv)
|
|
result = migrate(args.db, args.pg_service, args.out, schema=args.schema,
|
|
storage_path=args.storage_path, min_free_gib=args.min_free_gib)
|
|
print(json.dumps(result, ensure_ascii=False), flush=True)
|