from __future__ import annotations """Orchestrate the OSM station/track import and load pipeline.""" import argparse import sys from dataclasses import dataclass from pathlib import Path from typing import Callable, Sequence from backend.app.core.osm_config import DEFAULT_REGIONS from backend.scripts import ( stations_import, stations_load, tracks_import, tracks_load, ) @dataclass(slots=True) class Stage: label: str runner: Callable[[list[str] | None], int] args: list[str] input_path: Path | None = None output_path: Path | None = None def build_argument_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser( description="Run the station and track import/load workflow in sequence.", ) parser.add_argument( "--region", choices=[region.name for region in DEFAULT_REGIONS] + ["all"], default="all", help="Region selector forwarded to the import scripts (default: all).", ) parser.add_argument( "--output-dir", type=Path, default=Path("data"), help="Directory where intermediate JSON payloads are stored (default: data/).", ) parser.add_argument( "--stations-json", type=Path, help="Existing station JSON file to load; defaults to /osm_stations.json.", ) parser.add_argument( "--tracks-json", type=Path, help="Existing track JSON file to load; defaults to /osm_tracks.json.", ) parser.add_argument( "--skip-station-import", action="store_true", help="Skip the station import step (expects --stations-json to point to data).", ) parser.add_argument( "--skip-station-load", action="store_true", help="Skip loading stations into PostGIS.", ) parser.add_argument( "--skip-track-import", action="store_true", help="Skip the track import step (expects --tracks-json to point to data).", ) parser.add_argument( "--skip-track-load", action="store_true", help="Skip loading tracks into PostGIS.", ) parser.add_argument( "--dry-run", action="store_true", help="Print the planned stages without invoking Overpass or mutating the database.", ) parser.add_argument( "--commit", dest="commit", action="store_true", default=True, help="Commit database changes produced by the load steps (default).", ) parser.add_argument( "--no-commit", dest="commit", action="store_false", help="Rollback database changes after load steps (dry run).", ) return parser def _build_stage_plan(args: argparse.Namespace) -> list[Stage]: station_json = args.stations_json or args.output_dir / "osm_stations.json" track_json = args.tracks_json or args.output_dir / "osm_tracks.json" stages: list[Stage] = [] if not args.skip_station_import: stages.append( Stage( label="Import stations", runner=stations_import.main, args=["--output", str(station_json), "--region", args.region], output_path=station_json, ) ) if not args.skip_station_load: load_args = [str(station_json)] if not args.commit: load_args.append("--no-commit") stages.append( Stage( label="Load stations", runner=stations_load.main, args=load_args, input_path=station_json, ) ) if not args.skip_track_import: stages.append( Stage( label="Import tracks", runner=tracks_import.main, args=["--output", str(track_json), "--region", args.region], output_path=track_json, ) ) if not args.skip_track_load: load_args = [str(track_json)] if not args.commit: load_args.append("--no-commit") stages.append( Stage( label="Load tracks", runner=tracks_load.main, args=load_args, input_path=track_json, ) ) return stages def _describe_plan(stages: Sequence[Stage]) -> None: if not stages: print("No stages selected; nothing to do.") return print("Selected stages:") for stage in stages: detail = " ".join(stage.args) if stage.args else "" print(f" - {stage.label}: {detail}") def _execute_stage(stage: Stage) -> None: print(f"\n>>> {stage.label}") if stage.output_path is not None: stage.output_path.parent.mkdir(parents=True, exist_ok=True) if stage.input_path is not None and not stage.input_path.exists(): raise RuntimeError( f"Expected input file {stage.input_path} for {stage.label}; run the import step first or provide an existing file." ) try: exit_code = stage.runner(stage.args) except SystemExit as exc: # argparse.error exits via SystemExit exit_code = int(exc.code or 0) if exit_code: raise RuntimeError(f"{stage.label} failed with exit code {exit_code}.") def main(argv: list[str] | None = None) -> int: parser = build_argument_parser() args = parser.parse_args(argv) stages = _build_stage_plan(args) if args.dry_run: print("Dry run: the following stages would run in order.") _describe_plan(stages) return 0 for stage in stages: _execute_stage(stage) print("\nOSM refresh pipeline completed successfully.") return 0 if __name__ == "__main__": sys.exit(main())