pg2sqlite.py

#!/usr/bin/env python3
"""
Convert a PostgreSQL plain-text pg_dump (with COPY data) into a SQLite database.

Tailored to (but not limited to) Directus dumps from PG 13. Handles:
  - CREATE TABLE with PG types  -> SQLite affinity types
  - SERIAL columns (nextval defaults) -> INTEGER PRIMARY KEY AUTOINCREMENT
  - PRIMARY KEY / UNIQUE / FOREIGN KEY constraints (incl. ON DELETE actions)
  - DEFAULT clauses (cast-stripped; booleans -> 0/1)
  - COPY ... FROM stdin text data, including full escape decoding
  - booleans 't'/'f' -> 1/0 ; NULL (\\N) -> None ; JSON/uuid/timestamps -> TEXT

Data is inserted via parameterized executemany (no manual SQL escaping).
"""
import re
import sys
import sqlite3


# ---------------------------------------------------------------------------
# COPY text-format field decoder (PostgreSQL spec)
# ---------------------------------------------------------------------------
_SIMPLE_ESC = {'b': '\b', 'f': '\f', 'n': '\n', 'r': '\r',
               't': '\t', 'v': '\v', '\\': '\\'}


def decode_copy_field(s):
    """Decode one COPY text field. Returns None for the NULL marker (\\N)."""
    if s == r'\N':
        return None
    if '\\' not in s:
        return s
    out = []
    i, n = 0, len(s)
    while i < n:
        c = s[i]
        if c != '\\':
            out.append(c)
            i += 1
            continue
        if i + 1 >= n:                       # trailing lone backslash
            out.append('\\')
            break
        nxt = s[i + 1]
        if nxt in _SIMPLE_ESC:
            out.append(_SIMPLE_ESC[nxt])
            i += 2
        elif nxt == 'x':                     # \xHH hex (1-2 digits)
            j, hexd = i + 2, ''
            while j < n and len(hexd) < 2 and s[j] in '0123456789abcdefABCDEF':
                hexd += s[j]
                j += 1
            if hexd:
                out.append(chr(int(hexd, 16)))
                i = j
            else:
                out.append('x')
                i += 2
        elif nxt in '01234567':              # \OOO octal (1-3 digits)
            j, octd = i + 1, ''
            while j < n and len(octd) < 3 and s[j] in '01234567':
                octd += s[j]
                j += 1
            out.append(chr(int(octd, 8)))
            i = j
        else:                                # unknown escape -> literal char
            out.append(nxt)
            i += 2
    return ''.join(out)


# ---------------------------------------------------------------------------
# Type mapping  PG -> SQLite affinity
# ---------------------------------------------------------------------------
def pg_to_sqlite_type(pg_type):
    t = pg_type.strip().lower()
    if t in ('boolean', 'bool'):
        return 'INTEGER'                      # stored as 0/1
    if (t.startswith('int') or t in ('integer', 'bigint', 'smallint',
                                     'serial', 'bigserial', 'smallserial')):
        return 'INTEGER'
    if t in ('numeric', 'decimal') or t.startswith('numeric'):
        return 'NUMERIC'
    if t in ('real', 'double precision', 'float') or t.startswith('double'):
        return 'REAL'
    if t == 'bytea':
        return 'BLOB'
    # varchar/char/text/uuid/json/jsonb/date/time/timestamp* and anything else
    return 'TEXT'


def value_converter(pg_type):
    """Return a function converting a decoded COPY string to a Python value."""
    t = pg_type.strip().lower()
    if t in ('boolean', 'bool'):
        def conv(v):
            if v is None:
                return None
            return 1 if v in ('t', 'true', 'True', '1') else 0
        return conv
    if t.startswith('int') or t in ('integer', 'bigint', 'smallint'):
        return lambda v: None if v is None else int(v)
    if t in ('real', 'double precision', 'float') or t.startswith('double'):
        return lambda v: None if v is None else float(v)
    if t in ('numeric', 'decimal') or t.startswith('numeric'):
        # keep exact textual form to avoid float rounding
        return lambda v: v
    return lambda v: v                        # TEXT-like: keep decoded string


# ---------------------------------------------------------------------------
# Default-clause translation
# ---------------------------------------------------------------------------
_CAST_RE = re.compile(r"::[\w ]+(\([\d, ]+\))?\s*$")


def translate_default(expr):
    """Translate a PG column DEFAULT expression to a SQLite-safe one, or None
    to drop it (e.g. nextval)."""
    e = expr.strip()
    if e.lower().startswith('nextval('):
        return None                           # handled by AUTOINCREMENT
    e = _CAST_RE.sub('', e).strip()           # strip ::type cast
    low = e.lower()
    if low == 'true':
        return '1'
    if low == 'false':
        return '0'
    if low == 'null':
        return 'NULL'
    return e                                   # quoted literal / number / CURRENT_TIMESTAMP


# ---------------------------------------------------------------------------
# Identifier helpers
# ---------------------------------------------------------------------------
def unquote_ident(tok):
    tok = tok.strip()
    if tok.startswith('"') and tok.endswith('"'):
        return tok[1:-1].replace('""', '"')
    return tok


def q(ident):
    """Quote an identifier for SQLite output."""
    return '"' + ident.replace('"', '""') + '"'


def split_cols(s):
    """Split a parenthesised, comma-separated column list, honoring quotes."""
    return [unquote_ident(c) for c in re.findall(r'"[^"]*"|[^,\s]+', s)]


# ---------------------------------------------------------------------------
# Column-definition parser
# ---------------------------------------------------------------------------
def parse_column_def(line):
    """Parse one CREATE TABLE column line. Returns dict or None (table constraint)."""
    s = line.strip().rstrip(',').strip()
    if not s:
        return None
    upper = s.upper()
    # skip inline table constraints (none expected in these dumps, but be safe)
    if upper.startswith(('PRIMARY KEY', 'UNIQUE', 'FOREIGN KEY', 'CONSTRAINT',
                         'CHECK', 'EXCLUDE')):
        return None
    # column name (quoted or bare)
    if s.startswith('"'):
        m = re.match(r'"((?:[^"]|"")*)"\s*(.*)', s)
        name, rest = unquote_ident('"' + m.group(1) + '"'), m.group(2)
    else:
        m = re.match(r'(\S+)\s+(.*)', s)
        name, rest = m.group(1), m.group(2)
    not_null = False
    if re.search(r'\bNOT\s+NULL\s*$', rest, re.IGNORECASE):
        not_null = True
        rest = re.sub(r'\bNOT\s+NULL\s*$', '', rest, flags=re.IGNORECASE).strip()
    default = None
    mdef = re.search(r'\bDEFAULT\b', rest, re.IGNORECASE)
    if mdef:
        pg_type = rest[:mdef.start()].strip()
        default = rest[mdef.end():].strip()
    else:
        pg_type = rest.strip()
    return {'name': name, 'pg_type': pg_type, 'not_null': not_null,
            'default': default}


# ---------------------------------------------------------------------------
# Dump parser
# ---------------------------------------------------------------------------
class Table:
    def __init__(self, name):
        self.name = name
        self.columns = []           # list of column dicts (CREATE TABLE order)
        self.pk = []                # pk column names
        self.uniques = []           # list of (constraint_name, [cols])
        self.fks = []               # list of dict(cols, reftable, refcols, action)
        self.serial_cols = set()    # columns with nextval default


def parse_dump(text):
    lines = text.split('\n')
    tables = {}                     # name -> Table
    order = []                      # table creation order
    copies = []                     # list of (table, [cols], [rows])

    i, n = 0, len(lines)
    cur_alter = None                # current ALTER TABLE ONLY target

    while i < n:
        line = lines[i]
        stripped = line.strip()

        # ---- CREATE TABLE ------------------------------------------------
        m = re.match(r'CREATE TABLE (?:"?public"?\.)?"?([^"\s(]+)"?\s*\(', stripped)
        if m:
            tname = m.group(1)
            tbl = tables.setdefault(tname, Table(tname))
            if tname not in order:
                order.append(tname)
            i += 1
            while i < n and lines[i].strip() != ');':
                col = parse_column_def(lines[i])
                if col:
                    tbl.columns.append(col)
                i += 1
            i += 1
            continue

        # ---- ALTER TABLE ONLY public.X  (capture target) ----------------
        m = re.match(r'ALTER TABLE (?:ONLY )?(?:"?public"?\.)?"?([^"\s]+)"?'
                     r'(?:\s+(.*))?$', stripped)
        if m:
            cur_alter = m.group(1)
            tail = (m.group(2) or '').strip()
            # inline ALTER COLUMN ... SET DEFAULT nextval (serial detection)
            ms = re.search(r'ALTER COLUMN "?([^"\s]+)"? SET DEFAULT nextval', tail)
            if ms and cur_alter in tables:
                tables[cur_alter].serial_cols.add(ms.group(1))
            # the ADD CONSTRAINT may be on this same line or following lines
            i += 1
            # peek subsequent ADD CONSTRAINT line(s) belonging to this ALTER
            while i < n:
                cline = lines[i].strip()
                if cline.startswith('ADD CONSTRAINT'):
                    _parse_constraint(cline, cur_alter, tables)
                    i += 1
                    # constraint statement may span until ';'
                    while i < n and not lines[i - 1].rstrip().endswith(';'):
                        # continuation lines (rare) — append handled in _parse
                        i += 1
                    continue
                break
            continue

        # ---- COPY ... FROM stdin ----------------------------------------
        m = re.match(r'COPY (?:"?public"?\.)?"?([^"\s(]+)"?\s*\(([^)]*)\)\s+FROM stdin;',
                     stripped)
        if m:
            tname = m.group(1)
            cols = split_cols(m.group(2))
            rows = []
            i += 1
            while i < n and lines[i] != r'\.':
                rows.append(lines[i].split('\t'))
                i += 1
            i += 1  # skip the \.
            copies.append((tname, cols, rows))
            continue

        i += 1

    return tables, order, copies


def _parse_constraint(cline, table_name, tables):
    if table_name not in tables:
        tables[table_name] = Table(table_name)
    tbl = tables[table_name]
    m = re.match(r'ADD CONSTRAINT "?([^"\s]+)"?\s+(.*?);?\s*$', cline)
    if not m:
        return
    cname, body = m.group(1), m.group(2)
    mp = re.match(r'PRIMARY KEY\s*\(([^)]*)\)', body, re.IGNORECASE)
    if mp:
        tbl.pk = split_cols(mp.group(1))
        return
    mu = re.match(r'UNIQUE\s*\(([^)]*)\)', body, re.IGNORECASE)
    if mu:
        tbl.uniques.append((cname, split_cols(mu.group(1))))
        return
    mf = re.match(r'FOREIGN KEY\s*\(([^)]*)\)\s+REFERENCES\s+(?:"?public"?\.)?'
                  r'"?([^"\s(]+)"?\s*\(([^)]*)\)\s*(.*)$', body, re.IGNORECASE)
    if mf:
        tbl.fks.append({
            'cols': split_cols(mf.group(1)),
            'reftable': mf.group(2),
            'refcols': split_cols(mf.group(3)),
            'action': mf.group(4).strip(),
        })


# ---------------------------------------------------------------------------
# SQLite DDL generation
# ---------------------------------------------------------------------------
def build_create_table(tbl):
    lines = []
    pk_single = tbl.pk[0] if len(tbl.pk) == 1 else None
    coltypes = {c['name']: c['pg_type'] for c in tbl.columns}
    inline_pk_done = False

    for c in tbl.columns:
        name = c['name']
        sqlite_type = pg_to_sqlite_type(c['pg_type'])
        parts = [q(name), sqlite_type]
        is_int_pk = (name == pk_single and sqlite_type == 'INTEGER')
        if is_int_pk:
            parts.append('PRIMARY KEY')
            if name in tbl.serial_cols:
                parts.append('AUTOINCREMENT')
            inline_pk_done = True
            # NOT NULL implied; skip default (sequence-driven)
        else:
            if c['not_null']:
                parts.append('NOT NULL')
            if c['default'] is not None:
                d = translate_default(c['default'])
                if d is not None:
                    parts.append('DEFAULT ' + d)
        lines.append('    ' + ' '.join(parts))

    # table-level PRIMARY KEY (composite, or single non-integer)
    if tbl.pk and not inline_pk_done:
        lines.append('    PRIMARY KEY (%s)' % ', '.join(q(c) for c in tbl.pk))

    # UNIQUE constraints
    for cname, cols in tbl.uniques:
        lines.append('    UNIQUE (%s)' % ', '.join(q(c) for c in cols))

    # FOREIGN KEY constraints
    for fk in tbl.fks:
        cols = ', '.join(q(c) for c in fk['cols'])
        refcols = ', '.join(q(c) for c in fk['refcols'])
        clause = '    FOREIGN KEY (%s) REFERENCES %s (%s)' % (
            cols, q(fk['reftable']), refcols)
        if fk['action']:
            clause += ' ' + fk['action']
        lines.append(clause)

    return 'CREATE TABLE %s (\n%s\n);' % (q(tbl.name), ',\n'.join(lines)), coltypes


# ---------------------------------------------------------------------------
# Main conversion
# ---------------------------------------------------------------------------
def convert(dump_path, db_path):
    with open(dump_path, 'r', encoding='utf-8', newline='') as f:
        text = f.read()

    tables, order, copies = parse_dump(text)

    conn = sqlite3.connect(db_path)
    conn.execute('PRAGMA foreign_keys = OFF;')   # don't enforce during load
    conn.execute('PRAGMA journal_mode = OFF;')
    conn.execute('PRAGMA synchronous = OFF;')
    cur = conn.cursor()

    report = {'tables': [], 'total_rows': 0}

    # create all tables
    for tname in order:
        ddl, _ = build_create_table(tables[tname])
        cur.execute(ddl)
    conn.commit()

    # load data
    expected = {}
    for tname, cols, rows in copies:
        tbl = tables[tname]
        coltype = {c['name']: c['pg_type'] for c in tbl.columns}
        convs = [value_converter(coltype.get(col, 'text')) for col in cols]
        decoded_rows = []
        for raw in rows:
            if len(raw) == 1 and raw[0] == '':       # stray blank line guard
                continue
            vals = [convs[k](decode_copy_field(raw[k])) for k in range(len(cols))]
            decoded_rows.append(vals)
        placeholders = ', '.join('?' * len(cols))
        collist = ', '.join(q(c) for c in cols)
        sql = 'INSERT INTO %s (%s) VALUES (%s)' % (q(tname), collist, placeholders)
        cur.executemany(sql, decoded_rows)
        expected[tname] = len(decoded_rows)
        report['total_rows'] += len(decoded_rows)
    conn.commit()

    # verify row counts
    mismatches = []
    for tname in order:
        cur.execute('SELECT COUNT(*) FROM %s' % q(tname))
        got = cur.fetchone()[0]
        exp = expected.get(tname, 0)
        report['tables'].append((tname, exp, got))
        if got != exp:
            mismatches.append((tname, exp, got))

    conn.commit()
    conn.execute('PRAGMA foreign_keys = ON;')
    conn.close()

    return report, mismatches


if __name__ == '__main__':
    dump_path, db_path = sys.argv[1], sys.argv[2]
    report, mismatches = convert(dump_path, db_path)
    print('Converted %s -> %s' % (dump_path, db_path))
    print('Tables: %d   Total rows: %d' % (len(report['tables']), report['total_rows']))
    nonempty = [(t, e, g) for (t, e, g) in report['tables'] if e or g]
    for t, e, g in sorted(nonempty, key=lambda x: -x[1]):
        flag = '' if e == g else '  <-- MISMATCH'
        print('  %-32s expected=%-7d got=%-7d%s' % (t, e, g, flag))
    if mismatches:
        print('\nROW COUNT MISMATCHES:', mismatches)
        sys.exit(1)
    else:
        print('\nAll row counts match.')