playground/export_pg.py

#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.10"
# dependencies = ["psycopg[binary]"]
# ///
"""
export_pg.py — run an arbitrary Postgres query via psycopg and stream the
results into a SQLite database in batches, without pulling the whole result
set into memory.

This uses a server-side (named) cursor: the query is planned and executed
by Postgres exactly once, and rows are then fetched from that single result
set in --page-size batches. This is the same thing a GUI SQL client does
when you "load more rows" — it is NOT the same as re-running the query with
a different LIMIT/OFFSET each time, which would redo the scan/filter/sort
from scratch on every page and be dramatically slower on a big query.

Example (the motivating case):

    ./export_pg.py \\
        --dsn "service=mydb" \\
        --query 'SELECT id, item_identifier FROM order_item_details' \\
        --table order_item_details \\
        -o export.sqlite3

Connection: --dsn is passed straight to psycopg.connect() as a conninfo
string (e.g. "host=... dbname=... user=..." or a postgresql:// URL). If
omitted, psycopg falls back to the standard PG* environment variables /
libpq defaults.
"""

import argparse
import datetime
import decimal
import json
import sqlite3
import sys
import time
import uuid

import psycopg


def sqlite_type_for(value) -> str:
    if isinstance(value, bool) or isinstance(value, int):
        return "INTEGER"
    if isinstance(value, float):
        return "REAL"
    if isinstance(value, (bytes, bytearray, memoryview)):
        return "BLOB"
    return "TEXT"


def infer_sqlite_type(values: list) -> str:
    for v in values:
        if v is not None:
            return sqlite_type_for(v)
    return "TEXT"


def adapt(value):
    if value is None or isinstance(value, (int, float, str, bytes)):
        return value
    if isinstance(value, memoryview):
        return bytes(value)
    if isinstance(value, decimal.Decimal):
        return str(value)
    if isinstance(value, (datetime.date, datetime.time, datetime.datetime)):
        return value.isoformat()
    if isinstance(value, uuid.UUID):
        return str(value)
    if isinstance(value, (list, dict)):
        return json.dumps(value)
    return str(value)


def main():
    class Formatter(argparse.RawDescriptionHelpFormatter,
                     argparse.ArgumentDefaultsHelpFormatter):
        pass

    ap = argparse.ArgumentParser(description=__doc__, formatter_class=Formatter)
    ap.add_argument("--query", required=True,
                     help="query to run (a trailing ';' is stripped automatically)")
    ap.add_argument("-o", "--output", required=True, help="path to SQLite db file")
    ap.add_argument("--table", required=True, help="destination table name")
    ap.add_argument("--dsn", help="psycopg conninfo string / postgresql:// URL")
    ap.add_argument("--page-size", type=int, default=2_000,
                     help="rows fetched from the server-side cursor per batch")
    ap.add_argument("--if-exists", choices=["fail", "replace", "append"],
                     default="fail", help="behavior when --table already exists")
    args = ap.parse_args()

    query = args.query.strip().rstrip(";").strip()

    con = sqlite3.connect(args.output)
    cur = con.cursor()

    exists = cur.execute(
        "SELECT 1 FROM sqlite_master WHERE type='table' AND name=?",
        (args.table,),
    ).fetchone()
    if exists:
        if args.if_exists == "fail":
            print(f"table {args.table!r} already exists in {args.output} "
                  "(use --if-exists replace|append)", file=sys.stderr)
            sys.exit(1)
        elif args.if_exists == "replace":
            cur.execute(f'DROP TABLE "{args.table}"')
            exists = None

    pg = psycopg.connect(args.dsn) if args.dsn else psycopg.connect()

    total = 0
    columns = None
    col_types = None
    start = time.monotonic()
    show_progress = sys.stderr.isatty()

    def report(done: bool):
        elapsed = time.monotonic() - start
        rate = total / elapsed if elapsed > 0 else 0
        msg = f"{total} rows written ({rate:,.0f} rows/s, {elapsed:.0f}s elapsed)"
        if show_progress and not done:
            print(f"\r{msg}", end="", file=sys.stderr, flush=True)
        else:
            print(msg, file=sys.stderr)

    try:
        print(query, file=sys.stderr)
        with pg.cursor(name="export_pg") as pg_cur:
            pg_cur.itersize = args.page_size
            pg_cur.execute(query)

            while True:
                rows = pg_cur.fetchmany(args.page_size)
                if not rows:
                    break

                if columns is None:
                    columns = [d.name for d in pg_cur.description]
                    col_types = [
                        infer_sqlite_type([r[i] for r in rows])
                        for i in range(len(columns))
                    ]
                    if not exists:
                        cols_sql = ", ".join(
                            f'"{c}" {t}' for c, t in zip(columns, col_types)
                        )
                        cur.execute(f'CREATE TABLE "{args.table}" ({cols_sql})')

                placeholders = ", ".join("?" for _ in columns)
                converted = [[adapt(v) for v in row] for row in rows]
                cur.executemany(
                    f'INSERT INTO "{args.table}" VALUES ({placeholders})', converted
                )
                con.commit()

                total += len(rows)
                report(done=len(rows) < args.page_size)
    finally:
        pg.close()
        con.close()

    print(f"done: {total} rows written to {args.output}:{args.table}",
          file=sys.stderr)


if __name__ == "__main__":
    main()