feat(export-pg): Add batched Postgres to SQLite export

2026-09-05 09:29:29 +02:00
parent 3e5d58ea25
commit f3c20d8ed5
1 changed files with 170 additions and 0 deletions
@@ -0,0 +1,170 @@
+#!/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()