69 lines
2.3 KiB
Python
69 lines
2.3 KiB
Python
"""Apply a .sql script file against a PostgreSQL database via psycopg2."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
|
|
|
|
def pg_apply_sql(dsn: str, sql_path: str) -> int:
|
|
"""Read a .sql file and execute its full contents against PostgreSQL.
|
|
|
|
The whole script is sent in a single cursor.execute call. psycopg2 runs
|
|
multi-statement scripts in one execute when there are no bound parameters,
|
|
so DDL files with several statements separated by ";" apply atomically in
|
|
one transaction. Designed for idempotent migrations (the SQL itself uses
|
|
"IF NOT EXISTS"). Commits on success.
|
|
|
|
Args:
|
|
dsn: Connection string, e.g. "postgresql://user:pass@host:port/dbname".
|
|
sql_path: Path to the .sql file to apply (e.g. db/migrations/001_init.sql).
|
|
|
|
Returns:
|
|
Number of non-empty statements applied (counted by splitting on ";").
|
|
At minimum 1 when the script is non-empty.
|
|
|
|
Raises:
|
|
RuntimeError: If the file cannot be read, or the connection / execution
|
|
fails. The original exception is chained for debugging.
|
|
"""
|
|
path = Path(sql_path)
|
|
try:
|
|
script = path.read_text(encoding="utf-8")
|
|
except OSError as exc:
|
|
raise RuntimeError(
|
|
f"pg_apply_sql could not read {sql_path!r}: {exc}"
|
|
) from exc
|
|
|
|
if not script.strip():
|
|
return 0
|
|
|
|
# Lazy import so the module loads even without psycopg2 installed.
|
|
try:
|
|
import psycopg2
|
|
except ImportError as exc: # pragma: no cover - exercised only without dep
|
|
raise RuntimeError(
|
|
"psycopg2 is required for pg_apply_sql; install psycopg2-binary"
|
|
) from exc
|
|
|
|
# Best-effort statement count (informational return value only). Strip
|
|
# blank fragments produced by a trailing semicolon.
|
|
statement_count = sum(1 for part in script.split(";") if part.strip())
|
|
statement_count = max(statement_count, 1)
|
|
|
|
conn = None
|
|
try:
|
|
conn = psycopg2.connect(dsn)
|
|
with conn.cursor() as cur:
|
|
cur.execute(script)
|
|
conn.commit()
|
|
return statement_count
|
|
except Exception as exc:
|
|
if conn is not None:
|
|
conn.rollback()
|
|
raise RuntimeError(
|
|
f"pg_apply_sql failed applying {sql_path!r}: {exc}"
|
|
) from exc
|
|
finally:
|
|
if conn is not None:
|
|
conn.close()
|