gendesign/data/sql/61_import_nspd_batch.py

"""Импорт NSPD-сырого батча (cad_quarters + cad_buildings) из JSONL в Postgres.

Поддерживает 2 формата JSONL:
- v1 (старый): {cad, q_wkt, buildings: [{cn, p, n, a, ar, f, yb, yc, c, rd, w}]}
- v2 (новый): {cad, q_wkt, q_props, buildings: [{cn, w, props}]} — props это полный
  options-объект из NSPD (сохраняется как raw_props в БД).

Подключение через SSH-тунель:
    ssh -N -L 15432:localhost:5432 gendesign
    psql порт: localhost:15432, db gendesign, user gendesign, pwd <см .env>

Usage:
    cd backend && uv run python ../data/sql/61_import_nspd_batch.py <path/to/jsonl>
"""

from __future__ import annotations

import json
import os
import sys
from pathlib import Path

import psycopg

PG_DSN = os.environ.get(
    "TUNNEL_DSN",
    "postgresql://gendesign:2J2SBPMKuS998fiwhtQqDhMI@localhost:15432/gendesign",
)


def _to_int(x):
    if x is None or x == "":
        return None
    try:
        return int(x)
    except (TypeError, ValueError):
        return None


def _to_num(x):
    if x is None or x == "":
        return None
    try:
        return float(x)
    except (TypeError, ValueError):
        return None


def _to_date(x):
    if not x:
        return None
    s = str(x)
    return s[:10] if len(s) >= 10 else None


def insert_quarters(cur, quarters: list[dict]) -> int:
    inserted = 0
    for q in quarters:
        if not q.get("q_wkt"):
            continue
        raw = q.get("q_props")  # full options dict from NSPD (v2)
        cur.execute(
            """
            INSERT INTO cad_quarters_geom (cad_number, geom, raw_props, source)
            VALUES (
                %s,
                ST_Multi(ST_Transform(
                    ST_SetSRID(ST_GeomFromText(%s), 3857),
                    4326
                ))::geometry(MultiPolygon, 4326),
                CAST(%s AS jsonb),
                'nspd'
            )
            ON CONFLICT (cad_number) DO UPDATE SET
                geom = EXCLUDED.geom,
                raw_props = COALESCE(EXCLUDED.raw_props, cad_quarters_geom.raw_props),
                fetched_at = NOW()
            """,
            (
                q["cad"],
                q["q_wkt"],
                json.dumps(raw, ensure_ascii=False) if raw else None,
            ),
        )
        inserted += 1
    return inserted


def insert_buildings(cur, buildings: list[tuple[dict, str]]) -> int:
    """buildings: [(building_dict, parent_quarter_cad), ...]"""
    inserted = 0
    for b, qcad in buildings:
        if not b.get("w"):
            continue
        # v2 has full props; v1 has cherry-picked top-level keys
        props = b.get("props") or {}
        # Prefer v2 (props), fallback to v1 (top-level)
        cad_num = props.get("cad_num") or b.get("cn")
        purpose = props.get("purpose") or b.get("p")
        name = props.get("building_name") or b.get("n")
        addr = props.get("readable_address") or b.get("a")
        area = _to_num(props.get("area") if props else b.get("ar"))
        floors = props.get("floors") or b.get("f")
        year_built = _to_int(props.get("year_built") if props else b.get("yb"))
        year_comm = _to_int(props.get("year_commisioning") if props else b.get("yc"))
        cost = _to_num(props.get("cost_value") if props else b.get("c"))
        reg_date = _to_date(props.get("registration_date") if props else b.get("rd"))
        # New v2-only fields
        status = props.get("status")
        ownership = props.get("ownership_type")
        cultural = props.get("cultural_heritage_object") or props.get("cultural_heritage_val")
        if cultural is not None and not isinstance(cultural, str):
            cultural = json.dumps(cultural, ensure_ascii=False)
        underground = props.get("underground_floors")
        build_rec_area = _to_num(props.get("build_record_area"))
        build_rec_type = props.get("build_record_type_value")
        common_status = props.get("common_data_status")
        obj_type = props.get("type")
        raw_props_json = json.dumps(props, ensure_ascii=False) if props else None

        cur.execute(
            """
            INSERT INTO cad_buildings (
                cad_num, quarter_cad_num, geom, purpose, building_name,
                readable_address, area, floors, year_built, year_commisioning,
                cost_value, registration_date,
                status, ownership_type, cultural_heritage, underground_floors,
                build_record_area, build_record_type, common_data_status, obj_type,
                raw_props
            )
            VALUES (
                %s, %s,
                ST_Transform(ST_SetSRID(ST_GeomFromText(%s), 3857), 4326),
                %s, %s, %s, %s, %s, %s, %s, %s, %s,
                %s, %s, %s, %s, %s, %s, %s, %s,
                CAST(%s AS jsonb)
            )
            ON CONFLICT (cad_num) DO UPDATE SET
                cost_value = EXCLUDED.cost_value,
                area = EXCLUDED.area,
                year_built = EXCLUDED.year_built,
                year_commisioning = EXCLUDED.year_commisioning,
                status = EXCLUDED.status,
                ownership_type = EXCLUDED.ownership_type,
                cultural_heritage = EXCLUDED.cultural_heritage,
                underground_floors = EXCLUDED.underground_floors,
                build_record_area = EXCLUDED.build_record_area,
                build_record_type = EXCLUDED.build_record_type,
                common_data_status = EXCLUDED.common_data_status,
                obj_type = EXCLUDED.obj_type,
                raw_props = COALESCE(EXCLUDED.raw_props, cad_buildings.raw_props),
                fetched_at = NOW()
            """,
            (
                cad_num, qcad, b["w"],
                purpose, name, addr,
                area, floors,
                year_built, year_comm,
                cost, reg_date,
                status, ownership, cultural, underground,
                build_rec_area, build_rec_type, common_status, obj_type,
                raw_props_json,
            ),
        )
        inserted += 1
    return inserted


def main() -> int:
    if len(sys.argv) < 2:
        print("usage: 61_import_nspd_batch.py <jsonl_path>")
        return 1
    fp = Path(sys.argv[1])
    if not fp.exists():
        print(f"file not found: {fp}")
        return 1

    quarters: list[dict] = []
    buildings: list[tuple[dict, str]] = []
    with fp.open(encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            q = json.loads(line)
            quarters.append(q)
            for b in q.get("buildings", []):
                buildings.append((b, q["cad"]))

    print(f"Parsed: {len(quarters)} quarters, {len(buildings)} buildings")
    conn = psycopg.connect(PG_DSN, connect_timeout=10)
    try:
        cur = conn.cursor()
        n_q = insert_quarters(cur, quarters)
        n_b = insert_buildings(cur, buildings)
        conn.commit()
        print(f"Inserted: {n_q} quarters, {n_b} buildings")
        cur.execute("SELECT COUNT(*) FROM cad_quarters_geom")
        print(f"DB total cad_quarters_geom: {cur.fetchone()[0]}")
        cur.execute("SELECT COUNT(*) FROM cad_buildings")
        print(f"DB total cad_buildings: {cur.fetchone()[0]}")
    finally:
        conn.close()
    return 0


if __name__ == "__main__":
    sys.exit(main())