162 lines
6.1 KiB
Python
162 lines
6.1 KiB
Python
"""
|
|
SQL rendering: DataFrames into schema-agnostic PostgreSQL seed statements.
|
|
|
|
Three shapes, the same rows:
|
|
|
|
insert one INSERT per row, ON CONFLICT DO NOTHING. Runs in any client and can
|
|
be re-run; the column list repeats on every row, so it is the largest.
|
|
batch one INSERT per batch_rows rows, still ON CONFLICT DO NOTHING. Same
|
|
guarantees as insert, a fraction of the size, and far faster to load.
|
|
copy COPY ... FROM stdin, tab-separated. Smallest and fastest, but only psql
|
|
runs it, and COPY has no ON CONFLICT: loading it twice fails on the
|
|
first duplicate key.
|
|
|
|
A block is one transaction: open(), units, close(). Units are the pieces a file
|
|
may be split between (a row, or a batch), so every part of a split file is a
|
|
complete, loadable script. Getting them onto disk is output.py's business, and
|
|
nothing here holds a whole table's text, except a capped sample, which is small.
|
|
"""
|
|
|
|
import math
|
|
import re
|
|
|
|
import pandas as pd
|
|
|
|
from progress import Ticker
|
|
|
|
FORMATS = ("insert", "batch", "copy")
|
|
DEFAULT_BATCH_ROWS = 500
|
|
# Rows measured to estimate a full table's size, spread evenly through it: the
|
|
# first rows of a table are often shorter (small ids, early dates) than the rest.
|
|
MEASURE_ROWS = 1000
|
|
|
|
|
|
def sanitize_identifier(identifier: str) -> str:
|
|
"""Sanitize names for SQL tables, columns, and filenames."""
|
|
clean = re.sub(r"[^\w]", "_", str(identifier).strip().lower())
|
|
clean = re.sub(r"_+", "_", clean)
|
|
return clean.strip("_")
|
|
|
|
|
|
def sql_value(v) -> str:
|
|
if pd.isna(v):
|
|
return "NULL"
|
|
if isinstance(v, (bool, int)):
|
|
return str(v)
|
|
if isinstance(v, float):
|
|
return str(int(v)) if v.is_integer() else str(v)
|
|
escaped = str(v).replace("'", "''")
|
|
return f"'{escaped}'"
|
|
|
|
|
|
def copy_value(v) -> str:
|
|
"""The same value in COPY's text format: \\N for NULL, backslash escapes."""
|
|
if pd.isna(v):
|
|
return r"\N"
|
|
if isinstance(v, (bool, int)):
|
|
return str(v)
|
|
if isinstance(v, float):
|
|
return str(int(v)) if v.is_integer() else str(v)
|
|
return (str(v).replace("\\", "\\\\").replace("\t", "\\t")
|
|
.replace("\n", "\\n").replace("\r", "\\r"))
|
|
|
|
|
|
def utf8_len(text: str) -> int:
|
|
return len(text.encode("utf-8"))
|
|
|
|
|
|
class Shape:
|
|
"""How one table's rows are framed in one format."""
|
|
|
|
def __init__(self, fmt, table_name, columns, batch_rows=DEFAULT_BATCH_ROWS):
|
|
if fmt not in FORMATS:
|
|
raise ValueError(f"unknown sql format: {fmt}")
|
|
self.fmt = fmt
|
|
self.table_name = table_name
|
|
self.batch_rows = batch_rows
|
|
table_ref = f'"{table_name}"'
|
|
cols = ", ".join(f'"{sanitize_identifier(c)}"' for c in columns)
|
|
self.head = f"-- Generated seed data for table: {table_ref}\n"
|
|
if fmt == "copy":
|
|
self.head += "-- COPY FROM stdin: load with psql. COPY has no ON CONFLICT, so load it once.\n"
|
|
self.insert = f"INSERT INTO {table_ref} ({cols}) VALUES"
|
|
self.copy = f"COPY {table_ref} ({cols}) FROM stdin;\n"
|
|
|
|
def open(self, note="") -> str:
|
|
text = self.head + note + "BEGIN;\n\n"
|
|
return text + self.copy if self.fmt == "copy" else text
|
|
|
|
def close(self) -> str:
|
|
return ("\\.\n" if self.fmt == "copy" else "") + "\nCOMMIT;\n"
|
|
|
|
def _row(self, values) -> str:
|
|
if self.fmt == "copy":
|
|
return "\t".join(copy_value(v) for v in values) + "\n"
|
|
vals = ", ".join(sql_value(v) for v in values)
|
|
if self.fmt == "insert":
|
|
return f"{self.insert} ({vals}) ON CONFLICT DO NOTHING;\n"
|
|
return f"({vals})"
|
|
|
|
def _batch(self, rows) -> str:
|
|
return f"{self.insert}\n" + ",\n".join(rows) + "\nON CONFLICT DO NOTHING;\n\n"
|
|
|
|
def units(self, df: pd.DataFrame):
|
|
"""The rows as splittable pieces: one per row, or one per batch."""
|
|
ticker = Ticker(self.table_name, len(df))
|
|
batch = []
|
|
# iterrows, not itertuples: it hands values over the way the original
|
|
# tool did, and seed files people already load depend on that quoting.
|
|
for i, (_, row) in enumerate(df.iterrows(), 1):
|
|
text = self._row(row)
|
|
if self.fmt == "batch":
|
|
batch.append(text)
|
|
if len(batch) == self.batch_rows:
|
|
yield self._batch(batch)
|
|
batch = []
|
|
else:
|
|
yield text
|
|
if i % Ticker.CHECK == 0:
|
|
ticker.tick(i)
|
|
if batch:
|
|
yield self._batch(batch)
|
|
|
|
def estimate(self, df: pd.DataFrame) -> int:
|
|
"""Bytes one unsplit file of the whole table would take, from rows spread through it."""
|
|
total = len(df)
|
|
n = min(total, MEASURE_ROWS)
|
|
picks = sorted({round(i * (total - 1) / (n - 1)) for i in range(n)}) if n > 1 else [0]
|
|
measured = df.iloc[picks]
|
|
per_row = sum(utf8_len(self._row(row)) for _, row in measured.iterrows()) / len(measured)
|
|
fixed = utf8_len(self.open() + self.close())
|
|
if self.fmt == "batch":
|
|
batches = math.ceil(total / self.batch_rows)
|
|
framing = utf8_len(f"{self.insert}\n") + utf8_len("\nON CONFLICT DO NOTHING;\n\n")
|
|
return fixed + round((per_row + 2) * total) + batches * framing
|
|
return fixed + round(per_row * total)
|
|
|
|
|
|
def render_sample(shape: Shape, df: pd.DataFrame, max_rows: int):
|
|
"""
|
|
A capped table as one block of text, with a note of what was left out.
|
|
|
|
Return (text, full_bytes). Small by definition, so it is built in memory:
|
|
the note at the top needs the estimate, which needs the rows.
|
|
"""
|
|
shown = df.head(max_rows)
|
|
body = "".join(shape.units(shown))
|
|
full_bytes = shape.estimate(df)
|
|
note = (
|
|
f"-- SAMPLE: first {len(shown)} of {len(df)} rows. The full file would be "
|
|
f"~{human_bytes(full_bytes)}; run without --max-rows for all of it.\n"
|
|
)
|
|
return shape.open(note) + body + shape.close(), full_bytes
|
|
|
|
|
|
def human_bytes(n: int) -> str:
|
|
size = float(n)
|
|
for unit in ("B", "K", "M", "G"):
|
|
if size < 1000 or unit == "G":
|
|
return f"{size:.0f}{unit}" if unit == "B" else f"{size:.1f}{unit}"
|
|
size /= 1000
|
|
return f"{n}B"
|