93 lines
3.0 KiB
Python
93 lines
3.0 KiB
Python
"""Append-only batch insert of dict rows into a PostgreSQL table via psycopg2."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from datetime import date
|
|
|
|
|
|
def pg_insert_rows(
|
|
dsn: str,
|
|
table: str,
|
|
rows: list[dict],
|
|
add_snapshot_date: bool = True,
|
|
) -> int:
|
|
"""Insert rows (list of dicts) into a PostgreSQL table, append-only.
|
|
|
|
Columns are derived from the dict keys. If rows have heterogeneous schemas,
|
|
the union of all keys is used and missing values are filled with None, so a
|
|
single parameterized statement covers every row. Insertion uses
|
|
psycopg2.extras.execute_values (no string formatting) to avoid SQL injection.
|
|
|
|
Args:
|
|
dsn: Connection string, e.g. "postgresql://user:pass@host:port/dbname".
|
|
table: Target table name (must already exist).
|
|
rows: List of dicts; each dict is one row, keys are column names.
|
|
add_snapshot_date: If True and a row lacks "snapshot_date", inject
|
|
snapshot_date = date.today() before inserting.
|
|
|
|
Returns:
|
|
Number of rows inserted.
|
|
|
|
Raises:
|
|
RuntimeError: If the connection or the insert fails. The original
|
|
psycopg2 exception is chained for debugging.
|
|
"""
|
|
if not rows:
|
|
return 0
|
|
|
|
# psycopg2 is imported lazily so the module imports without the dependency
|
|
# present (self-test / introspection) and fails clearly only when invoked.
|
|
try:
|
|
import psycopg2
|
|
from psycopg2 import extras as pg_extras
|
|
from psycopg2 import sql as pg_sql
|
|
except ImportError as exc: # pragma: no cover - exercised only without dep
|
|
raise RuntimeError(
|
|
"psycopg2 is required for pg_insert_rows; install psycopg2-binary"
|
|
) from exc
|
|
|
|
# Work on copies so we never mutate the caller's dicts.
|
|
prepared: list[dict] = [dict(row) for row in rows]
|
|
|
|
if add_snapshot_date:
|
|
today = date.today()
|
|
for row in prepared:
|
|
row.setdefault("snapshot_date", today)
|
|
|
|
# Stable union of columns across all rows (first-seen order).
|
|
columns: list[str] = []
|
|
seen: set[str] = set()
|
|
for row in prepared:
|
|
for key in row:
|
|
if key not in seen:
|
|
seen.add(key)
|
|
columns.append(key)
|
|
|
|
if not columns:
|
|
return 0
|
|
|
|
# Build the value tuples in column order, filling absent keys with None.
|
|
values = [tuple(row.get(col) for col in columns) for row in prepared]
|
|
|
|
insert_stmt = pg_sql.SQL("INSERT INTO {table} ({cols}) VALUES %s").format(
|
|
table=pg_sql.Identifier(table),
|
|
cols=pg_sql.SQL(", ").join(pg_sql.Identifier(c) for c in columns),
|
|
)
|
|
|
|
conn = None
|
|
try:
|
|
conn = psycopg2.connect(dsn)
|
|
with conn.cursor() as cur:
|
|
pg_extras.execute_values(cur, insert_stmt, values)
|
|
conn.commit()
|
|
return len(values)
|
|
except Exception as exc:
|
|
if conn is not None:
|
|
conn.rollback()
|
|
raise RuntimeError(
|
|
f"pg_insert_rows failed inserting into {table!r}: {exc}"
|
|
) from exc
|
|
finally:
|
|
if conn is not None:
|
|
conn.close()
|