Reduce RAM consumtion during import

stops_import.py

@@ -4,6 +4,13 @@ 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)
 """
@@ -16,12 +23,12 @@ import unicodedata
 from pathlib import Path
 
 import click
-from flask import current_app
 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:
@@ -40,6 +47,78 @@ def _open_csv(path: Path) -> io.TextIOBase:
     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).
 
@@ -52,59 +131,32 @@ def import_stops(path: Path | str = SNAPSHOT_PATH) -> dict[str, int]:
             f"Snapshot not found: {path}\nRun scripts/fetch_stops.py first."
         )
 
-    with _open_csv(path) as fh:
-        rows = list(csv.DictReader(fh))
-    if not rows:
-        raise click.ClickException(f"{path} contains no rows.")
-
-    existing = {(s.osm_type, s.osm_id): s for s in Stop.query.all()}
     seen: set[tuple[str, int]] = set()
-    inserted = updated = 0
+    inserted = updated = total = 0
+    batch: list[dict] = []
 
-    for row in rows:
+    with _open_csv(path) as fh:
+        for row in csv.DictReader(fh):
             key = (row["osm_type"], int(row["osm_id"]))
             seen.add(key)
-        name = row["name"].strip()
-        fields = 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(),
-        )
-        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
+            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
 
-    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
-    }
-    removed = 0
-    for key, stop in existing.items():
-        if key not in seen and stop.id not in referenced:
-            db.session.delete(stop)
-            removed += 1
+    if total == 0:
+        raise click.ClickException(f"{path} contains no rows.")
 
-    db.session.commit()
+    removed = _remove_stale(seen)
     return {
-        "total": len(rows),
+        "total": total,
         "inserted": inserted,
         "updated": updated,
         "removed": removed,