Streaming Large GeoPackage Loads into PostGIS

from __future__ import annotations

from itertools import islice
from pathlib import Path
from typing import Iterable, Iterator

import fiona
import psycopg
import shapely
from shapely.geometry import shape

CHUNK = 50_000


def chunks(iterable: Iterable, size: int) -> Iterator[list]:
    it = iter(iterable)
    while batch := list(islice(it, size)):
        yield batch


def create_staging(conn: psycopg.Connection, table: str, srid: int) -> None:
    """Unlogged, index-free staging: the fastest shape a table can have for writes."""
    with conn.cursor() as cur:
        cur.execute(f"DROP TABLE IF EXISTS {table}_stage")
        cur.execute(
            f"""
            CREATE UNLOGGED TABLE {table}_stage (
                feature_id text,
                geom       geometry(MultiPolygon, {srid}),
                properties jsonb
            )
            """
        )


def stream_load(
    gpkg: Path, conn: psycopg.Connection, table: str, srid: int = 25832, layer: str = "features"
) -> int:
    """Load a GeoPackage into a staging table in bounded memory. Returns the row count."""
    create_staging(conn, table, srid)
    loaded = 0

    with fiona.open(gpkg, layer=layer) as src:
        if src.crs.to_epsg() != srid:
            # Refuse rather than transform: the chain reprojects at exactly one step,
            # and a surprise transform here would hide a contract violation upstream.
            raise ValueError(f"expected EPSG:{srid}, file declares {src.crs.to_epsg()}")

        for batch in chunks(src, CHUNK):
            with conn.cursor().copy(
                f"COPY {table}_stage (feature_id, geom, properties) FROM STDIN (FORMAT BINARY)"
            ) as copy:
                copy.set_types(["text", "geometry", "jsonb"])
                for feature in batch:
                    geom = shape(feature["geometry"])
                    if geom.geom_type == "Polygon":
                        geom = shapely.MultiPolygon([geom])   # match the column's type
                    copy.write_row((
                        feature["properties"].get("gml_id") or str(feature["id"]),
                        shapely.to_wkb(geom, include_srid=False),
                        json.dumps(dict(feature["properties"]), default=str),
                    ))
            loaded += len(batch)
            conn.commit()          # commit per chunk: bounded WAL, resumable progress

    return loaded


def finalise(conn: psycopg.Connection, table: str) -> None:
    """Index once, then swap. Both steps are far cheaper here than during the load."""
    with conn.cursor() as cur:
        cur.execute(f"CREATE INDEX ON {table}_stage USING gist (geom)")
        cur.execute(f"CREATE UNIQUE INDEX ON {table}_stage (feature_id)")
        cur.execute(f"ANALYZE {table}_stage")
        # One transaction: consumers see the old table or the new one, never neither.
        cur.execute("BEGIN")
        cur.execute(f"DROP TABLE IF EXISTS {table}_old")
        cur.execute(f"ALTER TABLE IF EXISTS {table} RENAME TO {table}_old")
        cur.execute(f"ALTER TABLE {table}_stage RENAME TO {table}")
        cur.execute(f"ALTER TABLE {table} SET LOGGED")
        cur.execute("COMMIT")

import resource


def test_memory_stays_bounded(large_gpkg, conn) -> None:
    before = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
    stream_load(large_gpkg, conn, table="parcels")
    after = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
    # A 6 GB file must not push peak RSS up by more than a few hundred megabytes.
    assert (after - before) < 800_000, "peak memory grew with the file — not streaming"


def test_failed_load_leaves_live_table_intact(large_gpkg, conn, monkeypatch) -> None:
    baseline = row_count(conn, "parcels")
    monkeypatch.setattr("mymodule.finalise", lambda *a, **k: 1 / 0)
    with pytest.raises(ZeroDivisionError):
        run_load(large_gpkg, conn, table="parcels")
    # The staging table may be half-built; the live table must be untouched.
    assert row_count(conn, "parcels") == baseline

SELECT count(*) AS features, ST_SRID(geom) AS srid FROM parcels GROUP BY 2;

SELECT indexname, indexdef FROM pg_indexes WHERE tablename = 'parcels';

-- Cheap sanity on the geometry itself: anything outside the CRS's valid range
-- means the contract upstream was wrong, whatever the SRID column says.
SELECT min(ST_XMin(geom)), max(ST_XMax(geom)),
       min(ST_YMin(geom)), max(ST_YMax(geom))
FROM   parcels;