Checkpointing for Interrupted Spatial Batch Jobs

#!/usr/bin/env python3
"""
Spatial batch processor with atomic checkpointing.

Usage:
    python checkpoint_batch.py /data/raw_gpkg /data/output_gpkg
    # Kill with Ctrl-C, then re-run — it resumes from the last completed file.

Requirements:
    pip install pyogrio
"""
import json
import logging
import os
import signal
import sys
from pathlib import Path
from typing import Optional

import pyogrio
import geopandas as gpd

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)-8s | %(message)s",
    stream=sys.stderr,
)
log = logging.getLogger(__name__)

CHECKPOINT_FILE = "spatial_batch_state.json"
TARGET_CRS = "EPSG:4326"


class SpatialCheckpoint:
    """Atomic, signal-aware checkpoint for spatial batch jobs."""

    def __init__(self, path: str = CHECKPOINT_FILE) -> None:
        self.path = Path(path)
        self.state: dict[str, bool] = self._load()
        self.interrupted = False
        # Flush on Ctrl-C (SIGINT) and orchestrator shutdown (SIGTERM)
        signal.signal(signal.SIGINT, self._handle_signal)
        signal.signal(signal.SIGTERM, self._handle_signal)

    def _load(self) -> dict[str, bool]:
        if not self.path.exists():
            return {}
        try:
            with open(self.path) as fh:
                return json.load(fh)
        except json.JSONDecodeError:
            # Truncated write from a previous crash — start fresh rather than
            # silently reprocessing all assets with wrong state
            log.warning("Checkpoint file corrupt; starting fresh: %s", self.path)
            return {}

    def _save(self) -> None:
        # Write to .tmp then os.replace() — atomic on POSIX and Windows 3.3+
        tmp = self.path.with_suffix(".tmp")
        with open(tmp, "w") as fh:
            json.dump(self.state, fh, indent=2)
        os.replace(str(tmp), str(self.path))   # ← crash-safe atomic rename

    def _handle_signal(self, signum: int, frame: object) -> None:
        log.info("Signal %d received — flushing checkpoint before exit.", signum)
        self.interrupted = True
        self._save()
        # Do NOT sys.exit() here; let the loop's `if checkpoint.interrupted` branch
        # exit cleanly so in-flight writes can finish their current iteration.

    def mark_complete(self, asset_id: str) -> None:
        """Call this AFTER the format-specific commit succeeds — never before."""
        self.state[asset_id] = True
        self._save()

    def get_pending(self, all_assets: list[str]) -> list[str]:
        """Return only assets that have not yet been successfully processed."""
        return [a for a in all_assets if not self.state.get(a, False)]

    @property
    def completed_count(self) -> int:
        return sum(1 for v in self.state.values() if v)


def process_one_asset(src_path: Path, dst_dir: Path) -> None:
    """
    Read a GeoPackage with pyogrio, reproject to EPSG:4326, write output.

    pyogrio is preferred over fiona for vector I/O: it uses Arrow-backed
    columnar reads that are ~5–10x faster on large feature sets.
    """
    gdf: gpd.GeoDataFrame = pyogrio.read_dataframe(
        str(src_path),
        use_arrow=True,    # Arrow-backed read; requires pyogrio >= 0.6
    )

    if gdf.crs is None:
        raise ValueError(f"No CRS on {src_path.name} — cannot reproject to {TARGET_CRS}")

    if str(gdf.crs) != TARGET_CRS:
        gdf = gdf.to_crs(TARGET_CRS)   # explicit EPSG:4326 coercion

    dst_path = dst_dir / src_path.name
    # pyogrio.write_dataframe wraps the underlying GPKG write in a single
    # transaction; on completion the SQLite COMMIT is issued before returning.
    pyogrio.write_dataframe(gdf, str(dst_path), driver="GPKG")


def process_spatial_batch(src_dir: Path, dst_dir: Path) -> int:
    dst_dir.mkdir(parents=True, exist_ok=True)

    all_assets = sorted(str(p) for p in src_dir.glob("*.gpkg"))
    if not all_assets:
        log.error("No .gpkg files found under %s", src_dir)
        return 2   # POSIX exit 2 = bad arguments / no input

    checkpoint = SpatialCheckpoint()
    pending = checkpoint.get_pending(all_assets)

    log.info(
        "Batch: %d total, %d already done, %d pending.",
        len(all_assets),
        checkpoint.completed_count,
        len(pending),
    )

    for asset in pending:
        if checkpoint.interrupted:
            log.info("Graceful shutdown — stopping before %s.", Path(asset).name)
            return 1   # POSIX exit 1 = interrupted / non-fatal error

        src_path = Path(asset)
        try:
            log.info("Processing %s", src_path.name)
            process_one_asset(src_path, dst_dir)
            # mark_complete AFTER the write transaction commits
            checkpoint.mark_complete(asset)
            log.info("Done: %s", src_path.name)
        except Exception as exc:
            # Isolate failures: one bad file does not abort the batch
            log.error("Failed %s: %s", src_path.name, exc)
            # Asset stays pending; it will be retried on the next run

    log.info("Batch finished. %d/%d assets processed.", checkpoint.completed_count, len(all_assets))
    return 0


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Resumable spatial batch processor")
    parser.add_argument("src_dir", type=Path, help="Directory of source .gpkg files")
    parser.add_argument("dst_dir", type=Path, help="Destination directory for output .gpkg files")
    args = parser.parse_args()

    sys.exit(process_spatial_batch(args.src_dir, args.dst_dir))

import shutil
import tempfile
from pathlib import Path
import pyogrio
import geopandas as gpd

def write_shapefile_safe(gdf: gpd.GeoDataFrame, dst_path: Path) -> None:
    """
    Write a Shapefile via a staging directory to avoid orphaned sidecars.
    dst_path should be the .shp file; all sidecars share the same stem.
    """
    stem = dst_path.stem
    with tempfile.TemporaryDirectory() as staging:
        staging_path = Path(staging) / f"{stem}.shp"
        pyogrio.write_dataframe(gdf, str(staging_path), driver="ESRI Shapefile")

        # Verify the staged output before moving
        staged_gdf = pyogrio.read_dataframe(str(staging_path))
        if len(staged_gdf) != len(gdf):
            raise RuntimeError(
                f"Staged shapefile feature count mismatch: "
                f"expected {len(gdf)}, got {len(staged_gdf)}"
            )

        # Move all sidecars (.shp, .shx, .dbf, .cpg, .prj) atomically
        for sidecar in Path(staging).glob(f"{stem}.*"):
            shutil.move(str(sidecar), str(dst_path.parent / sidecar.name))

# Inspect the checkpoint file — all completed assets should show true
python3 - <<'EOF'
import json, pathlib, sys

cp = pathlib.Path("spatial_batch_state.json")
if not cp.exists():
    print("No checkpoint file found.")
    sys.exit(0)

state = json.loads(cp.read_text())
done  = [k for k, v in state.items() if v]
total = len(state)
print(f"Checkpoint: {len(done)}/{total} assets marked complete")
for path in done:
    p = pathlib.Path(path)
    if not p.exists():
        print(f"  WARNING: checkpoint entry exists but file missing: {p.name}")
EOF

# Cross-check: count output files vs checkpoint entries
python3 - <<'EOF'
import json, pathlib, sys

state   = json.loads(pathlib.Path("spatial_batch_state.json").read_text())
done    = {k for k, v in state.items() if v}
outputs = set(str(p) for p in pathlib.Path("/data/output_gpkg").glob("*.gpkg"))

missing_outputs = done - outputs
extra_outputs   = outputs - done

if missing_outputs:
    print("Checkpoint says done but output file absent:")
    for p in sorted(missing_outputs):
        print(f"  {p}")
if extra_outputs:
    print("Output file exists but not in checkpoint (may be from a previous run):")
    for p in sorted(extra_outputs):
        print(f"  {p}")
if not missing_outputs and not extra_outputs:
    print("Checkpoint and output directory are in sync.")
EOF