905 lines
30 KiB
Python
905 lines
30 KiB
Python
"""Per-file SAS-to-Postgres loader.
|
|
|
|
Library-style functions plus a thin CLI wrapper. Designed so an orchestrator
|
|
can wrap the library for directory/batch mode; orchestration is out of scope
|
|
here.
|
|
|
|
Python 3.9 compatible (target is an air-gapped host that currently only has
|
|
3.9). ``from __future__ import annotations`` lets us use PEP 585 generics
|
|
as annotations; runtime-resolved type uses (dataclass defaults, etc.) stick
|
|
to ``typing``.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import datetime as dt
|
|
import io
|
|
import json
|
|
import os
|
|
import sys
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
from typing import Any, Dict, List, Optional, Tuple
|
|
|
|
import pandas as pd
|
|
import psycopg2
|
|
import psycopg2.extensions
|
|
import pyreadstat
|
|
import yaml
|
|
from dotenv import load_dotenv
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Top-level tunables
|
|
# ---------------------------------------------------------------------------
|
|
|
|
COERCE_CHAR_COLUMNS = True
|
|
"""If True, try to promote object (string) columns to numeric/date/timestamp
|
|
when every non-empty value parses cleanly."""
|
|
|
|
CHAR_INFERENCE_MIN_VALUES = 3
|
|
"""Don't attempt character-column coercion with fewer than this many non-empty
|
|
values; too small a sample is easy to mis-infer."""
|
|
|
|
NUMERIC_INT_RANGE = (-2_147_483_648, 2_147_483_647)
|
|
"""INTEGER bounds; anything outside becomes BIGINT."""
|
|
|
|
|
|
VALID_IF_EXISTS = ("fail", "replace", "append")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Dataclasses
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@dataclass
|
|
class LoaderConfig:
|
|
filename: Path
|
|
schemaname: str
|
|
tablename: str
|
|
if_exists: str = "fail"
|
|
include: Optional[List[str]] = None
|
|
exclude: Optional[List[str]] = None
|
|
|
|
|
|
@dataclass
|
|
class ColumnSpec:
|
|
name: str
|
|
postgres_type: str
|
|
nullable: bool
|
|
sas_format: Optional[str] = None
|
|
source_dtype: Optional[str] = None
|
|
notes: List[str] = field(default_factory=list)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Custom exceptions
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TableExistsError(RuntimeError):
|
|
"""Raised when if_exists=fail and the target table already exists."""
|
|
|
|
|
|
class SchemaCompatibilityError(RuntimeError):
|
|
"""Raised when if_exists=append and the incoming schema is not
|
|
compatible with the existing table."""
|
|
|
|
|
|
class ValidationError(RuntimeError):
|
|
"""Raised when --validate detects a mismatch against the manifest."""
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Connection
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def connect() -> psycopg2.extensions.connection:
|
|
"""Open a psycopg2 connection using standard libpq env vars.
|
|
|
|
Assumes `.env` has already been loaded (the CLI does this before calling).
|
|
Orchestrators that wrap this module should either call ``load_dotenv()``
|
|
themselves or ensure the env vars are set.
|
|
"""
|
|
conn = psycopg2.connect(
|
|
host=os.environ.get("PGHOST"),
|
|
port=os.environ.get("PGPORT"),
|
|
user=os.environ.get("PGUSER"),
|
|
password=os.environ.get("PGPASSWORD"),
|
|
dbname=os.environ.get("PGDATABASE"),
|
|
)
|
|
return conn
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Config loading
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def load_config(path: Path) -> LoaderConfig:
|
|
"""Parse and validate the YAML config at ``path``."""
|
|
path = Path(path)
|
|
with path.open("r", encoding="utf-8") as f:
|
|
raw = yaml.safe_load(f)
|
|
|
|
if not isinstance(raw, dict):
|
|
raise ValueError(f"Config at {path} must be a YAML mapping at the top level.")
|
|
|
|
missing = [k for k in ("filename", "schemaname", "tablename") if k not in raw]
|
|
if missing:
|
|
raise ValueError(f"Config {path} missing required keys: {', '.join(missing)}")
|
|
|
|
filename = Path(raw["filename"])
|
|
if not filename.is_absolute():
|
|
filename = (path.parent / filename).resolve() if (path.parent / filename).exists() else Path(raw["filename"])
|
|
|
|
schemaname = str(raw["schemaname"])
|
|
tablename = str(raw["tablename"])
|
|
|
|
if_exists = str(raw.get("if_exists", "fail")).lower()
|
|
if if_exists not in VALID_IF_EXISTS:
|
|
raise ValueError(
|
|
f"Config {path}: if_exists={if_exists!r} is not one of {VALID_IF_EXISTS}"
|
|
)
|
|
|
|
include = raw.get("include")
|
|
exclude = raw.get("exclude")
|
|
if include is not None and exclude is not None:
|
|
raise ValueError(
|
|
f"Config {path}: 'include' and 'exclude' are mutually exclusive; set at most one."
|
|
)
|
|
if include is not None and not isinstance(include, list):
|
|
raise ValueError(f"Config {path}: 'include' must be a list of column names.")
|
|
if exclude is not None and not isinstance(exclude, list):
|
|
raise ValueError(f"Config {path}: 'exclude' must be a list of column names.")
|
|
|
|
return LoaderConfig(
|
|
filename=filename,
|
|
schemaname=schemaname,
|
|
tablename=tablename,
|
|
if_exists=if_exists,
|
|
include=[str(c) for c in include] if include is not None else None,
|
|
exclude=[str(c) for c in exclude] if exclude is not None else None,
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Reader
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def read_sas(path: Path) -> Tuple[pd.DataFrame, Any]:
|
|
"""Dispatch to the right pyreadstat reader by extension.
|
|
|
|
Invariants (learned the hard way while building the sample generator):
|
|
|
|
* ``.xpt`` / ``.xport`` - no encoding arg; pyreadstat is flaky about
|
|
encoding on XPORT files it wrote itself.
|
|
* ``.sas7bdat`` - explicit ``encoding="latin-1"`` per colleague guidance.
|
|
"""
|
|
path = Path(path)
|
|
suffix = path.suffix.lower()
|
|
if suffix in (".xpt", ".xport"):
|
|
return pyreadstat.read_xport(str(path))
|
|
if suffix == ".sas7bdat":
|
|
return pyreadstat.read_sas7bdat(str(path), encoding="latin-1")
|
|
raise ValueError(f"Unsupported SAS file extension: {suffix}")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Column filtering
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def apply_column_filter(
|
|
df: pd.DataFrame,
|
|
include: Optional[List[str]],
|
|
exclude: Optional[List[str]],
|
|
) -> pd.DataFrame:
|
|
"""Restrict ``df`` to the requested columns. Names missing from the frame
|
|
raise a clear error rather than silently dropping."""
|
|
if include is not None and exclude is not None:
|
|
raise ValueError("include and exclude are mutually exclusive.")
|
|
|
|
if include is not None:
|
|
missing = [c for c in include if c not in df.columns]
|
|
if missing:
|
|
raise ValueError(f"include references unknown columns: {missing}")
|
|
return df.loc[:, list(include)].copy()
|
|
|
|
if exclude is not None:
|
|
missing = [c for c in exclude if c not in df.columns]
|
|
if missing:
|
|
raise ValueError(f"exclude references unknown columns: {missing}")
|
|
return df.drop(columns=list(exclude)).copy()
|
|
|
|
return df.copy()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Type inference
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
_DATE_FORMAT_PREFIXES = ("DATE", "YYMMDD", "MMDDYY", "DDMMYY", "JULIAN")
|
|
|
|
|
|
def _format_driven_type(sas_format: Optional[str]) -> Optional[str]:
|
|
"""Return a Postgres type inferred from the SAS format string, or None
|
|
if the format doesn't pin it down."""
|
|
if not sas_format:
|
|
return None
|
|
fmt = sas_format.upper().lstrip()
|
|
# DATETIME must be checked before DATE since "DATETIME20." starts with "DATE".
|
|
if fmt.startswith("DATETIME"):
|
|
return "TIMESTAMP"
|
|
if fmt.startswith("TIME"):
|
|
return "TIME"
|
|
for prefix in _DATE_FORMAT_PREFIXES:
|
|
if fmt.startswith(prefix):
|
|
return "DATE"
|
|
return None
|
|
|
|
|
|
def _all_null(series: pd.Series) -> bool:
|
|
if pd.api.types.is_object_dtype(series):
|
|
return bool(series.map(lambda v: v is None or (isinstance(v, str) and v == "") or (isinstance(v, float) and pd.isna(v))).all())
|
|
return bool(series.isna().all())
|
|
|
|
|
|
def _char_missing_mask(series: pd.Series) -> pd.Series:
|
|
return series.map(lambda v: v is None or (isinstance(v, float) and pd.isna(v)) or (isinstance(v, str) and v == ""))
|
|
|
|
|
|
def _is_nullable(series: pd.Series) -> bool:
|
|
"""True if the column has at least one missing value."""
|
|
if pd.api.types.is_object_dtype(series):
|
|
return bool(_char_missing_mask(series).any())
|
|
return bool(series.isna().any())
|
|
|
|
|
|
def _numeric_int_target(series: pd.Series) -> Optional[str]:
|
|
"""Given a numeric (float64) series, if every non-null value is a whole
|
|
number, return INTEGER or BIGINT depending on range; else None."""
|
|
nonnull = series.dropna()
|
|
if nonnull.empty:
|
|
return None
|
|
# Whole-number test. Guard against inf.
|
|
try:
|
|
whole = ((nonnull % 1) == 0).all()
|
|
except TypeError:
|
|
return None
|
|
if not whole:
|
|
return None
|
|
lo, hi = NUMERIC_INT_RANGE
|
|
vmin = nonnull.min()
|
|
vmax = nonnull.max()
|
|
if lo <= vmin and vmax <= hi:
|
|
return "INTEGER"
|
|
return "BIGINT"
|
|
|
|
|
|
def _object_is_dates(series: pd.Series) -> Tuple[bool, bool]:
|
|
"""Return (all-date-like, any-datetime). If every non-null value is a
|
|
``datetime.date`` / ``datetime.datetime`` / ``pd.Timestamp``, return True
|
|
plus whether at least one carries a time component."""
|
|
nonnull = series.dropna()
|
|
if nonnull.empty:
|
|
return False, False
|
|
any_datetime = False
|
|
for v in nonnull:
|
|
if isinstance(v, dt.datetime) or isinstance(v, pd.Timestamp):
|
|
any_datetime = True
|
|
continue
|
|
if isinstance(v, dt.date):
|
|
continue
|
|
return False, False
|
|
return True, any_datetime
|
|
|
|
|
|
def _try_int_coerce(values: List[str]) -> Optional[str]:
|
|
"""If every value parses as an int, return INTEGER/BIGINT, else None."""
|
|
ints: List[int] = []
|
|
for v in values:
|
|
s = v.strip()
|
|
try:
|
|
ints.append(int(s))
|
|
except ValueError:
|
|
return None
|
|
if not ints:
|
|
return None
|
|
lo, hi = NUMERIC_INT_RANGE
|
|
if all(lo <= i <= hi for i in ints):
|
|
return "INTEGER"
|
|
return "BIGINT"
|
|
|
|
|
|
def _try_float_coerce(values: List[str]) -> bool:
|
|
for v in values:
|
|
try:
|
|
float(v)
|
|
except ValueError:
|
|
return False
|
|
return True
|
|
|
|
|
|
def _try_date_coerce(values: List[str]) -> bool:
|
|
for v in values:
|
|
try:
|
|
dt.date.fromisoformat(v)
|
|
except (ValueError, TypeError):
|
|
return False
|
|
return True
|
|
|
|
|
|
def _try_datetime_coerce(values: List[str]) -> bool:
|
|
for v in values:
|
|
try:
|
|
dt.datetime.fromisoformat(v)
|
|
except (ValueError, TypeError):
|
|
return False
|
|
return True
|
|
|
|
|
|
def _infer_char_type(series: pd.Series) -> str:
|
|
"""Object/string column inference. Returns a Postgres type string."""
|
|
mask = _char_missing_mask(series)
|
|
nonempty = [str(v) for v in series[~mask].tolist()]
|
|
if not COERCE_CHAR_COLUMNS or len(nonempty) < CHAR_INFERENCE_MIN_VALUES:
|
|
return "TEXT"
|
|
|
|
int_guess = _try_int_coerce(nonempty)
|
|
if int_guess is not None:
|
|
return int_guess
|
|
if _try_float_coerce(nonempty):
|
|
return "DOUBLE PRECISION"
|
|
if _try_date_coerce(nonempty):
|
|
return "DATE"
|
|
if _try_datetime_coerce(nonempty):
|
|
return "TIMESTAMP"
|
|
return "TEXT"
|
|
|
|
|
|
def infer_schema(
|
|
df: pd.DataFrame,
|
|
meta: Any,
|
|
*,
|
|
coerce_chars: bool = COERCE_CHAR_COLUMNS,
|
|
) -> Dict[str, ColumnSpec]:
|
|
"""Infer a Postgres column spec for each column in ``df``.
|
|
|
|
``meta`` is the pyreadstat metadata object; we read
|
|
``meta.original_variable_types`` (a dict keyed by column name) for
|
|
format-driven date/time/timestamp inference.
|
|
|
|
The ``coerce_chars`` kwarg lets callers override the module-level
|
|
``COERCE_CHAR_COLUMNS`` without mutating global state. Internally the
|
|
char-inference helpers still read the constant - a full override would
|
|
thread the flag through, but the one-knob story here is intentional.
|
|
"""
|
|
original_formats: Dict[str, str] = dict(getattr(meta, "original_variable_types", {}) or {})
|
|
|
|
# Temporarily flip the module-level flag if the caller asked us to.
|
|
global COERCE_CHAR_COLUMNS
|
|
saved = COERCE_CHAR_COLUMNS
|
|
COERCE_CHAR_COLUMNS = coerce_chars
|
|
try:
|
|
out: Dict[str, ColumnSpec] = {}
|
|
for col in df.columns:
|
|
series = df[col]
|
|
sas_format = original_formats.get(col)
|
|
notes: List[str] = []
|
|
|
|
pg_type = _format_driven_type(sas_format)
|
|
|
|
if pg_type is None:
|
|
if _all_null(series):
|
|
pg_type = "TEXT"
|
|
notes.append("all-null column; defaulting to TEXT")
|
|
elif pd.api.types.is_datetime64_any_dtype(series):
|
|
pg_type = "TIMESTAMP"
|
|
elif pd.api.types.is_object_dtype(series):
|
|
is_dates, any_dt = _object_is_dates(series)
|
|
if is_dates:
|
|
pg_type = "TIMESTAMP" if any_dt else "DATE"
|
|
else:
|
|
pg_type = _infer_char_type(series)
|
|
elif pd.api.types.is_numeric_dtype(series):
|
|
int_target = _numeric_int_target(series)
|
|
if int_target is not None:
|
|
pg_type = int_target
|
|
else:
|
|
pg_type = "DOUBLE PRECISION"
|
|
else:
|
|
pg_type = "TEXT"
|
|
notes.append(f"unhandled dtype {series.dtype}; defaulting to TEXT")
|
|
|
|
nullable = _is_nullable(series)
|
|
|
|
out[col] = ColumnSpec(
|
|
name=col,
|
|
postgres_type=pg_type,
|
|
nullable=nullable,
|
|
sas_format=sas_format,
|
|
source_dtype=str(series.dtype),
|
|
notes=notes,
|
|
)
|
|
return out
|
|
finally:
|
|
COERCE_CHAR_COLUMNS = saved
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Table management
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _quote_ident(ident: str) -> str:
|
|
"""Quote a Postgres identifier. psycopg2 doesn't expose this directly
|
|
until 2.8+ with sql.Identifier; we do it by hand to stay driver-simple."""
|
|
return '"' + ident.replace('"', '""') + '"'
|
|
|
|
|
|
def _qualified(schema: str, table: str) -> str:
|
|
return f"{_quote_ident(schema)}.{_quote_ident(table)}"
|
|
|
|
|
|
def _table_exists(conn, schema: str, table: str) -> bool:
|
|
with conn.cursor() as cur:
|
|
cur.execute(
|
|
"SELECT 1 FROM information_schema.tables "
|
|
"WHERE table_schema = %s AND table_name = %s",
|
|
(schema, table),
|
|
)
|
|
return cur.fetchone() is not None
|
|
|
|
|
|
def render_create_table(schema: str, table: str, columns: Dict[str, ColumnSpec]) -> str:
|
|
lines = []
|
|
for spec in columns.values():
|
|
null_clause = "" if spec.nullable else " NOT NULL"
|
|
lines.append(f" {_quote_ident(spec.name)} {spec.postgres_type}{null_clause}")
|
|
body = ",\n".join(lines)
|
|
return f"CREATE TABLE {_qualified(schema, table)} (\n{body}\n);"
|
|
|
|
|
|
def _create_table_sql(conn, schema: str, table: str, columns: Dict[str, ColumnSpec]) -> None:
|
|
sql = render_create_table(schema, table, columns)
|
|
with conn.cursor() as cur:
|
|
cur.execute(sql)
|
|
|
|
|
|
def _drop_table(conn, schema: str, table: str) -> None:
|
|
with conn.cursor() as cur:
|
|
cur.execute(f"DROP TABLE {_qualified(schema, table)}")
|
|
|
|
|
|
# Normalization table: map both loader-emitted and Postgres-reported type
|
|
# strings to a single canonical family name. Ignore length/precision
|
|
# modifiers like VARCHAR(n) and NUMERIC(p,s).
|
|
_TYPE_NORMALIZATION: Dict[str, str] = {
|
|
"INTEGER": "integer",
|
|
"INT": "integer",
|
|
"INT4": "integer",
|
|
"BIGINT": "bigint",
|
|
"INT8": "bigint",
|
|
"SMALLINT": "smallint",
|
|
"INT2": "smallint",
|
|
"DOUBLE PRECISION": "double precision",
|
|
"FLOAT8": "double precision",
|
|
"REAL": "real",
|
|
"FLOAT4": "real",
|
|
"NUMERIC": "numeric",
|
|
"DECIMAL": "numeric",
|
|
"TEXT": "text",
|
|
"VARCHAR": "character varying",
|
|
"CHARACTER VARYING": "character varying",
|
|
"CHAR": "character",
|
|
"CHARACTER": "character",
|
|
"BPCHAR": "character",
|
|
"BOOLEAN": "boolean",
|
|
"BOOL": "boolean",
|
|
"DATE": "date",
|
|
"TIMESTAMP": "timestamp without time zone",
|
|
"TIMESTAMP WITHOUT TIME ZONE": "timestamp without time zone",
|
|
"TIMESTAMPTZ": "timestamp with time zone",
|
|
"TIMESTAMP WITH TIME ZONE": "timestamp with time zone",
|
|
"TIME": "time without time zone",
|
|
"TIME WITHOUT TIME ZONE": "time without time zone",
|
|
"TIMETZ": "time with time zone",
|
|
"TIME WITH TIME ZONE": "time with time zone",
|
|
}
|
|
|
|
|
|
def _normalize_type(pg_type: str) -> str:
|
|
"""Strip length/precision modifiers and map to canonical family."""
|
|
stripped = pg_type.strip().upper()
|
|
# Remove trailing (n) / (p,s) before the space-separated tail.
|
|
# Examples: "VARCHAR(10)" -> "VARCHAR"; "TIMESTAMP(6) WITHOUT TIME ZONE" -> "TIMESTAMP WITHOUT TIME ZONE"
|
|
import re
|
|
|
|
stripped = re.sub(r"\(\s*\d+\s*(?:,\s*\d+\s*)?\)", "", stripped).strip()
|
|
# Collapse doubled whitespace after paren removal.
|
|
stripped = re.sub(r"\s+", " ", stripped)
|
|
return _TYPE_NORMALIZATION.get(stripped, stripped.lower())
|
|
|
|
|
|
def _assert_schema_compatible(
|
|
conn, schema: str, table: str, columns: Dict[str, ColumnSpec]
|
|
) -> None:
|
|
"""Pre-flight check for if_exists=append. See plan section on option B."""
|
|
with conn.cursor() as cur:
|
|
cur.execute(
|
|
"SELECT column_name, data_type, is_nullable "
|
|
"FROM information_schema.columns "
|
|
"WHERE table_schema = %s AND table_name = %s",
|
|
(schema, table),
|
|
)
|
|
existing = {row[0]: (row[1], row[2]) for row in cur.fetchall()}
|
|
|
|
mismatches: List[str] = []
|
|
warnings: List[str] = []
|
|
|
|
for name, spec in columns.items():
|
|
if name not in existing:
|
|
mismatches.append(
|
|
f"column {name!r} not present in target {schema}.{table}"
|
|
)
|
|
continue
|
|
target_type, target_nullable = existing[name]
|
|
inferred_norm = _normalize_type(spec.postgres_type)
|
|
target_norm = _normalize_type(target_type)
|
|
if inferred_norm != target_norm:
|
|
mismatches.append(
|
|
f"column {name!r}: inferred {spec.postgres_type} "
|
|
f"(normalized {inferred_norm!r}) but target is {target_type} "
|
|
f"(normalized {target_norm!r})"
|
|
)
|
|
target_is_notnull = (target_nullable == "NO")
|
|
if spec.nullable and target_is_notnull:
|
|
warnings.append(
|
|
f"column {name!r}: incoming allows NULLs but target is NOT NULL; "
|
|
"COPY will fail if any NULLs appear"
|
|
)
|
|
|
|
for w in warnings:
|
|
print(f"[warn] {w}", file=sys.stderr)
|
|
|
|
if mismatches:
|
|
raise SchemaCompatibilityError(
|
|
"append-mode schema compatibility check failed:\n - "
|
|
+ "\n - ".join(mismatches)
|
|
)
|
|
|
|
|
|
def create_table(
|
|
conn,
|
|
schema_name: str,
|
|
table_name: str,
|
|
columns: Dict[str, ColumnSpec],
|
|
if_exists: str,
|
|
) -> None:
|
|
"""Create (or verify) the target table according to ``if_exists``."""
|
|
if if_exists not in VALID_IF_EXISTS:
|
|
raise ValueError(f"if_exists must be one of {VALID_IF_EXISTS}, got {if_exists!r}")
|
|
|
|
exists = _table_exists(conn, schema_name, table_name)
|
|
if exists:
|
|
if if_exists == "fail":
|
|
raise TableExistsError(
|
|
f"Table {schema_name}.{table_name} already exists and if_exists=fail"
|
|
)
|
|
if if_exists == "replace":
|
|
_drop_table(conn, schema_name, table_name)
|
|
_create_table_sql(conn, schema_name, table_name, columns)
|
|
return
|
|
if if_exists == "append":
|
|
_assert_schema_compatible(conn, schema_name, table_name, columns)
|
|
return
|
|
else:
|
|
_create_table_sql(conn, schema_name, table_name, columns)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# COPY loading
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _seconds_to_time(v: Any) -> Optional[dt.time]:
|
|
if v is None:
|
|
return None
|
|
if isinstance(v, float) and pd.isna(v):
|
|
return None
|
|
if isinstance(v, dt.time):
|
|
return v
|
|
if isinstance(v, (dt.datetime, pd.Timestamp)):
|
|
return v.time() if not pd.isna(v) else None
|
|
try:
|
|
total = int(round(float(v)))
|
|
except (TypeError, ValueError):
|
|
return None
|
|
h, rem = divmod(total, 3600)
|
|
m, s = divmod(rem, 60)
|
|
# Clamp; TIME8. is always within a day.
|
|
h = max(0, min(h, 23))
|
|
return dt.time(h, m, s)
|
|
|
|
|
|
def _prepare_for_copy(df: pd.DataFrame, columns: Dict[str, ColumnSpec]) -> pd.DataFrame:
|
|
"""Materialize a copy of ``df`` with each column in the right shape for
|
|
``to_csv`` so the CSV lands as valid input for the target Postgres type.
|
|
"""
|
|
out = pd.DataFrame(index=df.index)
|
|
for name, spec in columns.items():
|
|
series = df[name]
|
|
pg = spec.postgres_type.upper()
|
|
|
|
if pg in ("INTEGER", "BIGINT", "SMALLINT"):
|
|
if pd.api.types.is_object_dtype(series):
|
|
series = pd.to_numeric(
|
|
series.replace({"": None}), errors="coerce"
|
|
)
|
|
out[name] = series.astype("Int64")
|
|
elif pg in ("DOUBLE PRECISION", "REAL", "NUMERIC"):
|
|
if pd.api.types.is_object_dtype(series):
|
|
series = pd.to_numeric(
|
|
series.replace({"": None}), errors="coerce"
|
|
)
|
|
out[name] = series.astype("float64")
|
|
elif pg == "DATE":
|
|
if pd.api.types.is_datetime64_any_dtype(series):
|
|
out[name] = series.dt.date
|
|
elif pd.api.types.is_object_dtype(series):
|
|
def _to_date(v: Any) -> Optional[dt.date]:
|
|
if v is None or (isinstance(v, float) and pd.isna(v)):
|
|
return None
|
|
if isinstance(v, dt.datetime):
|
|
return v.date()
|
|
if isinstance(v, dt.date):
|
|
return v
|
|
if isinstance(v, str):
|
|
if v == "":
|
|
return None
|
|
try:
|
|
return dt.date.fromisoformat(v)
|
|
except ValueError:
|
|
return None
|
|
return None
|
|
out[name] = series.map(_to_date)
|
|
else:
|
|
out[name] = series
|
|
elif pg in ("TIMESTAMP", "TIMESTAMP WITHOUT TIME ZONE", "TIMESTAMP WITH TIME ZONE"):
|
|
if pd.api.types.is_datetime64_any_dtype(series):
|
|
out[name] = series
|
|
elif pd.api.types.is_object_dtype(series):
|
|
def _to_dt(v: Any) -> Optional[dt.datetime]:
|
|
if v is None or (isinstance(v, float) and pd.isna(v)):
|
|
return None
|
|
if isinstance(v, dt.datetime):
|
|
return v
|
|
if isinstance(v, dt.date):
|
|
return dt.datetime(v.year, v.month, v.day)
|
|
if isinstance(v, pd.Timestamp):
|
|
return v.to_pydatetime() if not pd.isna(v) else None
|
|
if isinstance(v, str):
|
|
if v == "":
|
|
return None
|
|
try:
|
|
return dt.datetime.fromisoformat(v)
|
|
except ValueError:
|
|
return None
|
|
return None
|
|
out[name] = series.map(_to_dt)
|
|
else:
|
|
out[name] = series
|
|
elif pg in ("TIME", "TIME WITHOUT TIME ZONE", "TIME WITH TIME ZONE"):
|
|
out[name] = series.map(_seconds_to_time)
|
|
elif pg in ("TEXT", "VARCHAR", "CHARACTER VARYING", "CHAR", "CHARACTER"):
|
|
# Leave empty strings as "" so `NULL ''` in COPY turns them into NULL.
|
|
def _to_str(v: Any) -> Any:
|
|
if v is None:
|
|
return ""
|
|
if isinstance(v, float) and pd.isna(v):
|
|
return ""
|
|
return str(v)
|
|
out[name] = series.map(_to_str)
|
|
elif pg == "BOOLEAN":
|
|
out[name] = series.astype("boolean") if series.dtype != object else series
|
|
else:
|
|
out[name] = series
|
|
return out
|
|
|
|
|
|
def copy_dataframe(
|
|
conn,
|
|
schema_name: str,
|
|
table_name: str,
|
|
df: pd.DataFrame,
|
|
columns: Dict[str, ColumnSpec],
|
|
) -> int:
|
|
"""Stream ``df`` into Postgres via ``COPY ... FROM STDIN``.
|
|
|
|
Returns the number of rows inserted.
|
|
"""
|
|
prepared = _prepare_for_copy(df, columns)
|
|
|
|
buf = io.StringIO()
|
|
prepared.to_csv(
|
|
buf,
|
|
index=False,
|
|
header=False,
|
|
na_rep="",
|
|
date_format="%Y-%m-%d %H:%M:%S",
|
|
)
|
|
buf.seek(0)
|
|
|
|
col_list = ", ".join(_quote_ident(name) for name in columns.keys())
|
|
sql = (
|
|
f"COPY {_qualified(schema_name, table_name)} ({col_list}) "
|
|
f"FROM STDIN WITH (FORMAT csv, NULL '')"
|
|
)
|
|
|
|
with conn.cursor() as cur:
|
|
cur.copy_expert(sql, buf)
|
|
rowcount = cur.rowcount
|
|
|
|
return int(rowcount) if rowcount is not None else len(prepared)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Manifest validation
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _match_manifest_type(inferred: str, manifest_entry: Dict[str, Any]) -> bool:
|
|
inferred_norm = _normalize_type(inferred)
|
|
if "postgres_type" in manifest_entry:
|
|
return inferred_norm == _normalize_type(manifest_entry["postgres_type"])
|
|
if "acceptable_types" in manifest_entry:
|
|
return any(
|
|
inferred_norm == _normalize_type(t)
|
|
for t in manifest_entry["acceptable_types"]
|
|
)
|
|
return False
|
|
|
|
|
|
def validate_against_manifest(
|
|
inferred: Dict[str, ColumnSpec],
|
|
manifest_path: Path,
|
|
) -> List[str]:
|
|
"""Compare the inferred schema against the expected-types manifest.
|
|
|
|
Returns a list of human-readable problem strings; empty list means OK.
|
|
"""
|
|
manifest_path = Path(manifest_path)
|
|
if not manifest_path.exists():
|
|
return [f"manifest not found: {manifest_path}"]
|
|
|
|
with manifest_path.open("r", encoding="utf-8") as f:
|
|
manifest = json.load(f)
|
|
|
|
problems: List[str] = []
|
|
|
|
only_in_inferred = set(inferred) - set(manifest)
|
|
only_in_manifest = set(manifest) - set(inferred)
|
|
if only_in_inferred:
|
|
problems.append(
|
|
f"columns in inferred but not manifest: {sorted(only_in_inferred)}"
|
|
)
|
|
if only_in_manifest:
|
|
problems.append(
|
|
f"columns in manifest but not inferred: {sorted(only_in_manifest)}"
|
|
)
|
|
|
|
for name, spec in inferred.items():
|
|
entry = manifest.get(name)
|
|
if entry is None:
|
|
continue
|
|
if not _match_manifest_type(spec.postgres_type, entry):
|
|
expected = entry.get("postgres_type") or entry.get("acceptable_types")
|
|
problems.append(
|
|
f"column {name!r}: inferred {spec.postgres_type!r}, "
|
|
f"manifest expected {expected!r}"
|
|
)
|
|
manifest_nullable = bool(entry.get("nullable", True))
|
|
if spec.nullable and not manifest_nullable:
|
|
problems.append(
|
|
f"column {name!r}: inferred nullable, manifest expects NOT NULL "
|
|
f"(loosening nullability is never allowed)"
|
|
)
|
|
|
|
return problems
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# CLI
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _build_argparser() -> argparse.ArgumentParser:
|
|
p = argparse.ArgumentParser(
|
|
description="Load a single SAS file (XPT or sas7bdat) into Postgres.",
|
|
)
|
|
p.add_argument("--config", required=True, type=Path, help="Path to YAML config")
|
|
p.add_argument(
|
|
"--validate",
|
|
action="store_true",
|
|
help=(
|
|
"Compare inferred schema against <filename-stem>.expected.json "
|
|
"next to the SAS file; exits nonzero on mismatch."
|
|
),
|
|
)
|
|
p.add_argument(
|
|
"--dry-run",
|
|
action="store_true",
|
|
help="Print inferred CREATE TABLE and stop; don't touch Postgres.",
|
|
)
|
|
return p
|
|
|
|
|
|
def _format_columns_summary(columns: Dict[str, ColumnSpec]) -> str:
|
|
lines = []
|
|
for spec in columns.values():
|
|
null = "" if spec.nullable else " NOT NULL"
|
|
lines.append(f" {spec.name}: {spec.postgres_type}{null}")
|
|
return "\n".join(lines)
|
|
|
|
|
|
def main(argv: Optional[List[str]] = None) -> int:
|
|
args = _build_argparser().parse_args(argv)
|
|
|
|
load_dotenv()
|
|
|
|
cfg = load_config(args.config)
|
|
|
|
if not cfg.filename.exists():
|
|
print(f"error: SAS file not found: {cfg.filename}", file=sys.stderr)
|
|
return 2
|
|
|
|
df, meta = read_sas(cfg.filename)
|
|
df = apply_column_filter(df, cfg.include, cfg.exclude)
|
|
columns = infer_schema(df, meta)
|
|
|
|
if args.validate:
|
|
manifest_path = cfg.filename.with_suffix("").with_suffix(".expected.json")
|
|
# The above strips .xpt then appends .expected.json, e.g.
|
|
# "sample_kitchensink.xpt" -> "sample_kitchensink.expected.json".
|
|
problems = validate_against_manifest(columns, manifest_path)
|
|
if problems:
|
|
print("validation failed:", file=sys.stderr)
|
|
for p in problems:
|
|
print(f" - {p}", file=sys.stderr)
|
|
return 1
|
|
print(f"validation OK ({len(columns)} columns match {manifest_path.name})")
|
|
|
|
if args.dry_run:
|
|
print(render_create_table(cfg.schemaname, cfg.tablename, columns))
|
|
return 0
|
|
|
|
conn = connect()
|
|
conn.autocommit = False
|
|
try:
|
|
create_table(conn, cfg.schemaname, cfg.tablename, columns, cfg.if_exists)
|
|
inserted = copy_dataframe(conn, cfg.schemaname, cfg.tablename, df, columns)
|
|
conn.commit()
|
|
except Exception:
|
|
conn.rollback()
|
|
raise
|
|
finally:
|
|
conn.close()
|
|
|
|
print(
|
|
f"loaded {inserted} rows into {cfg.schemaname}.{cfg.tablename} "
|
|
f"({len(columns)} columns)"
|
|
)
|
|
print("final schema:")
|
|
print(_format_columns_summary(columns))
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|