pbt/stops_import.py

"""Load the checked-in DACH stop snapshot into the ``stops`` table.

The snapshot (data/stops_dach.csv.gz) is produced by scripts/fetch_stops.py
from OpenStreetMap data (© OpenStreetMap contributors, ODbL) for Austria,
Germany and Switzerland. Importing needs no network access.

At ~307k rows, processing this in one shot - the whole table loaded as ORM
objects, ~307k new Stop() instances held until a single final commit - is
enough to OOM a small host (confirmed: it did, on a 512 MB production LXC).
Everything below is batched (default 2000 rows) with a commit + expunge_all()
after each batch, so peak memory stays roughly flat regardless of table size
instead of growing with it.

    flask import-stops              # import the checked-in snapshot
    flask import-stops --file X     # import an alternative CSV(.gz)
"""
from __future__ import annotations

import csv
import gzip
import io
import unicodedata
from pathlib import Path

import click
from flask.cli import with_appcontext

from models import Stop, Trip, db

SNAPSHOT_PATH = Path(__file__).resolve().parent / "data" / "stops_dach.csv.gz"
BATCH_SIZE = 2000


def normalize_name(value: str) -> str:
    """Lowercase, fold accents and ß for diacritic-insensitive search.

    Mirrors scripts/fetch_stops.py:normalize_name - keep the two in sync.
    """
    value = value.strip().lower().replace("ß", "ss")
    decomposed = unicodedata.normalize("NFKD", value)
    return "".join(ch for ch in decomposed if not unicodedata.combining(ch))


def _open_csv(path: Path) -> io.TextIOBase:
    if str(path).endswith(".gz"):
        return gzip.open(path, "rt", newline="", encoding="utf-8")
    return open(path, "rt", newline="", encoding="utf-8")


def _row_fields(row: dict) -> dict:
    name = row["name"].strip()
    return dict(
        name=name,
        name_normalized=normalize_name(name),
        stop_type=row["stop_type"] or "other",
        country=(row.get("country") or "AT").strip().upper(),
        municipality=(row.get("municipality") or "").strip() or None,
        latitude=float(row["latitude"]),
        longitude=float(row["longitude"]),
        lines=(row.get("lines") or "").strip(),
    )


def _import_batch(batch: list[dict]) -> tuple[int, int]:
    """Upsert one batch. Only this batch's rows are ever loaded as ORM
    objects, so memory stays bounded by BATCH_SIZE, not table size."""
    ids = [row["osm_id"] for row in batch]
    existing = {
        (s.osm_type, s.osm_id): s
        for s in Stop.query.filter(Stop.osm_id.in_(ids))
    }
    inserted = updated = 0
    for row in batch:
        key = (row["osm_type"], row["osm_id"])
        fields = row["fields"]
        stop = existing.get(key)
        if stop is None:
            db.session.add(Stop(osm_type=key[0], osm_id=key[1], **fields))
            inserted += 1
        else:
            changed = False
            for attr, value in fields.items():
                if getattr(stop, attr) != value:
                    setattr(stop, attr, value)
                    changed = True
            updated += changed
    db.session.commit()
    db.session.expunge_all()  # drop the identity map's refs so GC can reclaim them
    return inserted, updated


def _remove_stale(seen: set[tuple[str, int]]) -> int:
    """Delete stops absent from the new snapshot, unless a trip still
    references them. Streams the id/key columns only (yield_per) rather than
    loading full Stop rows for the whole table."""
    referenced = {
        sid
        for (sid,) in db.session.query(Trip.origin_stop_id).distinct()
        if sid is not None
    } | {
        sid
        for (sid,) in db.session.query(Trip.destination_stop_id).distinct()
        if sid is not None
    }

    stale_ids = [
        sid
        for sid, osm_type, osm_id in db.session.query(
            Stop.id, Stop.osm_type, Stop.osm_id
        ).yield_per(5000)
        if (osm_type, osm_id) not in seen and sid not in referenced
    ]

    removed = 0
    for i in range(0, len(stale_ids), BATCH_SIZE):
        chunk = stale_ids[i : i + BATCH_SIZE]
        removed += Stop.query.filter(Stop.id.in_(chunk)).delete(synchronize_session=False)
        db.session.commit()
    return removed


def import_stops(path: Path | str = SNAPSHOT_PATH) -> dict[str, int]:
    """Upsert every row of the snapshot into ``stops`` by (osm_type, osm_id).

    Stops that vanished from the snapshot are removed unless a trip still
    references them, so foreign keys from ``trips`` stay valid.
    """
    path = Path(path)
    if not path.exists():
        raise click.ClickException(
            f"Snapshot not found: {path}\nRun scripts/fetch_stops.py first."
        )

    seen: set[tuple[str, int]] = set()
    inserted = updated = total = 0
    batch: list[dict] = []

    with _open_csv(path) as fh:
        for row in csv.DictReader(fh):
            key = (row["osm_type"], int(row["osm_id"]))
            seen.add(key)
            total += 1
            batch.append({"osm_type": key[0], "osm_id": key[1], "fields": _row_fields(row)})
            if len(batch) >= BATCH_SIZE:
                i, u = _import_batch(batch)
                inserted += i
                updated += u
                batch = []
        if batch:
            i, u = _import_batch(batch)
            inserted += i
            updated += u

    if total == 0:
        raise click.ClickException(f"{path} contains no rows.")

    removed = _remove_stale(seen)
    return {
        "total": total,
        "inserted": inserted,
        "updated": updated,
        "removed": removed,
    }


@click.command("import-stops")
@click.option(
    "--file",
    "file_path",
    type=click.Path(exists=True, dir_okay=False),
    default=None,
    help="CSV(.gz) to import instead of the checked-in snapshot.",
)
@with_appcontext
def import_stops_command(file_path: str | None) -> None:
    """Import DACH public-transport stops into the database."""
    stats = import_stops(file_path or SNAPSHOT_PATH)
    click.echo(
        "Stops imported: "
        f"{stats['total']} in snapshot, "
        f"{stats['inserted']} new, {stats['updated']} updated, "
        f"{stats['removed']} removed. "
        f"Total now: {Stop.query.count()}."
    )


def register_cli(app) -> None:
    app.cli.add_command(import_stops_command)