paper-atlas/paper_atlas/postgres/schema.py

"""Inspect the complete business schema; rebuild SQLite engine internals natively."""
import hashlib
import json
import re

TABLES = ('sources', 'papers', 'identifiers', 'source_records', 'venues', 'editions',
          'venue_aliases', 'edition_sources', 'submissions', 'decisions',
          'submission_sources', 'submission_identifiers', 'decision_sources',
          'submission_merges', 'decision_merges', 'schema_migrations', 'tags',
          'paper_tags', 'match_candidates', 'paper_merges', 'ingest_runs')
CHECKS = {
    'venues': ["venue_type IN ('conference','journal','review_platform')"],
    'edition_sources': ["status IN ('pending','running','complete','partial','blocked','unavailable','not_yet_published')"],
    'decisions': ["status IN ('accepted','rejected','desk_rejected','withdrawn','under_review','unknown')"],
    'match_candidates': ["status IN ('pending','merged','distinct')"],
}
TYPES = {'INTEGER': 'BIGINT', 'TEXT': 'TEXT COLLATE "C"', 'REAL': 'DOUBLE PRECISION'}


def checks_in(ddl):
    result = []
    for match in re.finditer(r'\bCHECK\s*\(', ddl, re.I):
        start = position = match.end()
        depth, quoted = 1, False
        while position < len(ddl) and depth:
            char = ddl[position]
            if char == "'":
                if quoted and position + 1 < len(ddl) and ddl[position + 1] == "'":
                    position += 2; continue
                quoted = not quoted
            elif not quoted:
                depth += (char == '(') - (char == ')')
            position += 1
        if depth:
            raise ValueError('Unclosed CHECK constraint')
        result.append(ddl[start:position - 1])
    return result


def normalize_check(expression):
    # Normalize SQL layout/keywords, retaining case and spaces inside literals.
    tokens = re.findall(r"'(?:[^']|'')*'|[^']+", expression)
    return ''.join(token if token.startswith("'") else re.sub(r'\s+', '', token).lower() for token in tokens)


def quote(name):
    return '"' + name.replace('"', '""') + '"'


def inspect(connection):
    business = {name for schema, name, kind, *_ in connection.execute('PRAGMA table_list')
                if schema == 'main' and kind == 'table' and not name.startswith('sqlite_')}
    if business != set(TABLES):
        raise ValueError('Unexpected business tables: ' + repr(sorted(business ^ set(TABLES))))
    result = {}
    for table in TABLES:
        columns = [{'name': r[1], 'type': r[2].upper(), 'notnull': bool(r[3]),
                    'default': r[4], 'pk': r[5]}
                   for r in connection.execute('PRAGMA table_info(' + quote(table) + ')')]
        if any(c['type'] not in TYPES for c in columns):
            raise ValueError('Unsupported storage type in ' + table)
        pk = [c['name'] for c in sorted(columns, key=lambda c: c['pk']) if c['pk']]
        if not pk:
            raise ValueError('A deterministic primary key is required for ' + table)
        indexes = []
        for row in connection.execute('PRAGMA index_list(' + quote(table) + ')'):
            if row[3] == 'pk':
                continue
            if row[4]:
                raise ValueError('Partial indexes require explicit migration support')
            terms = [(r[2], bool(r[3]), r[4]) for r in connection.execute('PRAGMA index_xinfo(' + quote(row[1]) + ')') if r[5]]
            if any(name is None or coll != 'BINARY' for name, desc, coll in terms):
                raise ValueError('Unsupported expression or collation in ' + table)
            indexes.append({'name': row[1], 'unique': bool(row[2]), 'origin': row[3],
                            'terms': [(name, desc) for name, desc, coll in terms]})
        fks = {}
        for row in connection.execute('PRAGMA foreign_key_list(' + quote(table) + ')'):
            if row[5:7] != ('NO ACTION', 'NO ACTION'):
                raise ValueError('Unsupported foreign key actions')
            fk = fks.setdefault(row[0], {'table': row[2], 'columns': [], 'references': []})
            fk['columns'].append(row[3]); fk['references'].append(row[4])
        ddl = connection.execute("SELECT sql FROM sqlite_schema WHERE type='table' AND name=?", (table,)).fetchone()[0]
        checks = CHECKS.get(table, [])
        if sorted(map(normalize_check, checks_in(ddl))) != sorted(map(normalize_check, checks)):
            raise ValueError('Unrecognized CHECK constraint in ' + table)
        result[table] = {'columns': columns, 'pk': pk, 'indexes': indexes,
                         'foreign_keys': list(fks.values()), 'checks': checks}
    return result


def fingerprint(plan):
    return hashlib.sha256(json.dumps(plan, sort_keys=True, separators=(',', ':')).encode()).hexdigest()


def default_sql(value):
    if value is None:
        return ''
    if value.upper() == 'CURRENT_TIMESTAMP':
        return " DEFAULT to_char(CURRENT_TIMESTAMP AT TIME ZONE 'UTC','YYYY-MM-DD\"T\"HH24:MI:SS.MS\"Z\"')"
    if re.fullmatch(r'-?\d+|\'(?:[^\']|\'\')*\'', value):
        return ' DEFAULT ' + value
    raise ValueError('Unsupported column default')


def table_sql(schema, table, spec):
    columns = []
    for column in spec['columns']:
        generated = column['name'] == 'id' and spec['pk'] == ['id'] and column['type'] == 'INTEGER'
        value = quote(column['name']) + ' ' + TYPES[column['type']]
        value += ' GENERATED BY DEFAULT AS IDENTITY' if generated else default_sql(column['default'])
        if column['notnull']:
            value += ' NOT NULL'
        columns.append(value)
    return 'CREATE TABLE ' + quote(schema) + '.' + quote(table) + '(' + ','.join(columns) + ')'


def select_sql(schema, table, spec):
    prefix = quote(schema) + '.' if schema else ''
    return ('SELECT ' + ','.join(quote(c['name']) for c in spec['columns']) +
            ' FROM ' + prefix + quote(table) + ' ORDER BY ' + ','.join(quote(c) for c in spec['pk']))