netx/scripts/_import_ume_csv_snapshot.py

"""One-shot: ensure schema, then load field CSVs from ume/ into local DB."""

from __future__ import annotations

import csv
import json
import sys
from datetime import datetime
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

import netx_api.models  # noqa: F401
from netx_api.db import Base, SessionLocal, engine
from netx_api.schema_patches import apply_all_legacy_startup_ddl
from sqlalchemy import text

UME_DIR = ROOT / "ume"

FILES = [
    ("ume_inventory_ne", "ume_inventory_ne_202608061843.csv"),
    ("ume_topo_node", "ume_topo_node_202608061841.csv"),
    ("ume_topo_link", "ume_topo_link_202608061841.csv"),
    ("topo_fabric_node", "topo_fabric_node_202608061847.csv"),
    ("topo_fabric_edge", "topo_fabric_edge_202608061847.csv"),
]

INT_COLS = {"x_pos", "y_pos", "layer_rate"}
JSON_COLS = {"attrs"}
DT_COLS = {
    "first_seen_at",
    "last_seen_at",
    "discovered_at",
    "created_at",
    "updated_at",
}
NULLABLE_EMPTY = {
    "managed_ne_id",
    "region_folder_id",
    "discovered_at",
    "x_pos",
    "y_pos",
    "layer_rate",
}
# Unique nullable on fabric nodes: empty string would violate UNIQUE; use NULL.
FABRIC_NULL_EMPTY = {"ume_ne_id", "managed_ne_id", "region_folder_id"}


def _parse_dt(v: str):
    s = (v or "").strip()
    if not s:
        return None
    for fmt in (
        "%Y-%m-%d %H:%M:%S.%f",
        "%Y-%m-%d %H:%M:%S",
        "%Y-%m-%dT%H:%M:%S.%f",
        "%Y-%m-%dT%H:%M:%S",
    ):
        try:
            return datetime.strptime(s, fmt)
        except ValueError:
            continue
    return s


def _coerce(col: str, raw: str, *, table: str):
    if raw is None:
        return None
    v = raw
    if isinstance(v, str):
        v = v.strip()
    null_empty = set(NULLABLE_EMPTY)
    if table == "topo_fabric_node":
        null_empty |= FABRIC_NULL_EMPTY
    if v == "" and col in null_empty:
        return None
    if col in INT_COLS:
        if v == "":
            return None
        return int(v)
    if col in JSON_COLS:
        if v == "":
            return {}
        if isinstance(v, (dict, list)):
            return v
        try:
            return json.loads(v)
        except Exception:
            return {}
    if col in DT_COLS:
        return _parse_dt(v if isinstance(v, str) else str(v))
    return v if v is not None else ""


def _load_csv(path: Path, *, table: str) -> tuple[list[str], list[dict]]:
    with path.open("r", encoding="utf-8-sig", newline="") as f:
        reader = csv.DictReader(f)
        cols = list(reader.fieldnames or [])
        rows = []
        for row in reader:
            rows.append({c: _coerce(c, row.get(c, ""), table=table) for c in cols})
    return cols, rows


def main() -> None:
    print("ensure schema…")
    Base.metadata.create_all(bind=engine)
    apply_all_legacy_startup_ddl(engine)

    # FK-safe truncate order
    with engine.begin() as conn:
        conn.execute(
            text(
                "TRUNCATE TABLE topo_fabric_edge, topo_fabric_node, "
                "ume_topo_link, ume_topo_node, ume_inventory_ne RESTART IDENTITY CASCADE"
            )
        )
        print("truncated ok")

    db = SessionLocal()
    try:
        for table, fname in FILES:
            path = UME_DIR / fname
            if not path.exists():
                raise FileNotFoundError(path)
            cols, rows = _load_csv(path, table=table)
            print(f"load {table}: {len(rows)} rows from {fname}")
            if not rows:
                continue
            # chunked insert via core
            chunk = 1000
            meta = Base.metadata.tables[table]
            for i in range(0, len(rows), chunk):
                part = rows[i : i + chunk]
                # drop unknown cols just in case
                clean = [{k: r[k] for k in cols if k in meta.c} for r in part]
                db.execute(meta.insert(), clean)
                db.commit()
                print(f"  … {min(i + chunk, len(rows))}/{len(rows)}")
        print("counts:")
        for table, _ in FILES:
            n = db.execute(text(f"SELECT count(*) FROM {table}")).scalar()
            print(f"  {table}: {n}")
    finally:
        db.close()


if __name__ == "__main__":
    main()