#!/usr/bin/env python3 """ Download completed NSRDB county polygon archive ZIP files. This shared engine reads JSON acknowledgements produced by the polygon request workflow and downloads each outputs.downloadUrl. Metric-specific wrappers supply the response and archive directories. """ from __future__ import annotations import argparse import json import time import urllib.error import urllib.parse import urllib.request from dataclasses import dataclass, field from enum import Enum from pathlib import Path from typing import Any DEFAULT_TIMEOUT = 300 DEFAULT_CHUNK_SIZE = 1024 * 1024 class ArchiveDownloadState(Enum): """Distinct lifecycle phases for one NSRDB archive download.""" QUEUED = "queued" CHECKING = "checking" DOWNLOADING = "downloading" SKIPPED = "skipped" PENDING = "pending" DOWNLOADED = "downloaded" FAILED = "failed" class ArchiveDownloadEvent(Enum): """Events that may move an archive download to another lifecycle phase.""" START = "start" EXISTING_FILE_FOUND = "existing_file_found" DOWNLOAD_STARTED = "download_started" ARCHIVE_NOT_READY = "archive_not_ready" DOWNLOAD_SUCCEEDED = "download_succeeded" DOWNLOAD_FAILED = "download_failed" ARCHIVE_DOWNLOAD_TRANSITIONS = { (ArchiveDownloadState.QUEUED, ArchiveDownloadEvent.START): ArchiveDownloadState.CHECKING, ( ArchiveDownloadState.CHECKING, ArchiveDownloadEvent.EXISTING_FILE_FOUND, ): ArchiveDownloadState.SKIPPED, ( ArchiveDownloadState.CHECKING, ArchiveDownloadEvent.DOWNLOAD_STARTED, ): ArchiveDownloadState.DOWNLOADING, ( ArchiveDownloadState.DOWNLOADING, ArchiveDownloadEvent.ARCHIVE_NOT_READY, ): ArchiveDownloadState.PENDING, ( ArchiveDownloadState.DOWNLOADING, ArchiveDownloadEvent.DOWNLOAD_SUCCEEDED, ): ArchiveDownloadState.DOWNLOADED, ( ArchiveDownloadState.QUEUED, ArchiveDownloadEvent.DOWNLOAD_FAILED, ): ArchiveDownloadState.FAILED, ( ArchiveDownloadState.CHECKING, ArchiveDownloadEvent.DOWNLOAD_FAILED, ): ArchiveDownloadState.FAILED, ( ArchiveDownloadState.DOWNLOADING, ArchiveDownloadEvent.DOWNLOAD_FAILED, ): ArchiveDownloadState.FAILED, } class InvalidArchiveDownloadTransition(RuntimeError): """Reject an event that is not valid for the current download state.""" @dataclass class ArchiveDownloadStateMachine: """Track and validate the lifecycle of one archive download.""" label: str state: ArchiveDownloadState = ArchiveDownloadState.QUEUED history: list[ArchiveDownloadState] = field( default_factory=lambda: [ArchiveDownloadState.QUEUED] ) def transition(self, event: ArchiveDownloadEvent) -> ArchiveDownloadState: """Apply one event and return the resulting state.""" next_state = ARCHIVE_DOWNLOAD_TRANSITIONS.get((self.state, event)) if next_state is None: raise InvalidArchiveDownloadTransition( f"Invalid archive download transition for {self.label}: " f"{event.value} while {self.state.value}." ) self.state = next_state self.history.append(next_state) return next_state class PolygonArchiveDownloadError(RuntimeError): """Store a failed archive download with useful context.""" def __init__(self, message: str, status_code: int | None = None) -> None: super().__init__(message) self.status_code = status_code def is_pending_s3_archive_response(url: str, status_code: int) -> bool: """Return whether an S3 HTTP response likely means the archive is pending.""" if status_code != 403: return False host = urllib.parse.urlsplit(url).netloc.lower() return host == "s3.amazonaws.com" or host.endswith(".amazonaws.com") def read_response(response_path: Path) -> dict[str, Any]: """Read a saved NSRDB archive response JSON file.""" with response_path.open(encoding="utf-8") as handle: response = json.load(handle) if not isinstance(response, dict): raise ValueError("Response JSON must contain an object.") return response def download_url_from_response(response: dict[str, Any]) -> str: """Return the download URL from a saved NSRDB archive response.""" outputs = response.get("outputs") if not isinstance(outputs, dict): raise ValueError("Response JSON is missing an outputs object.") download_url = outputs.get("downloadUrl") if not isinstance(download_url, str) or not download_url.strip(): raise ValueError("Response JSON is missing outputs.downloadUrl.") return download_url.strip() def output_name_for_response(response_path: Path) -> str: """Convert a response JSON filename into the corresponding ZIP filename.""" stem = response_path.stem if stem.endswith("_response"): stem = stem[: -len("_response")] return f"{stem}.zip" def download_file( url: str, output_path: Path, timeout: int, overwrite: bool, machine: ArchiveDownloadStateMachine, ) -> ArchiveDownloadState: """Download a URL to disk and return the terminal download state.""" output_path.parent.mkdir(parents=True, exist_ok=True) if output_path.exists() and not overwrite: return machine.transition(ArchiveDownloadEvent.EXISTING_FILE_FOUND) machine.transition(ArchiveDownloadEvent.DOWNLOAD_STARTED) temp_path = output_path.with_suffix(f"{output_path.suffix}.part") if temp_path.exists(): temp_path.unlink() request = urllib.request.Request(url, headers={"User-Agent": "county-climate-explorer/0.1"}) try: with urllib.request.urlopen(request, timeout=timeout) as response: with temp_path.open("wb") as handle: while True: chunk = response.read(DEFAULT_CHUNK_SIZE) if not chunk: break handle.write(chunk) except urllib.error.HTTPError as error: if temp_path.exists(): temp_path.unlink() if is_pending_s3_archive_response(url, error.code): return machine.transition(ArchiveDownloadEvent.ARCHIVE_NOT_READY) machine.transition(ArchiveDownloadEvent.DOWNLOAD_FAILED) raise PolygonArchiveDownloadError(f"HTTP {error.code} {error.reason}", error.code) from error except urllib.error.URLError as error: if temp_path.exists(): temp_path.unlink() machine.transition(ArchiveDownloadEvent.DOWNLOAD_FAILED) raise PolygonArchiveDownloadError(str(error.reason)) from error temp_path.replace(output_path) return machine.transition(ArchiveDownloadEvent.DOWNLOAD_SUCCEEDED) def response_paths(response_dir: Path, start: int, limit: int | None) -> list[Path]: """Return saved response JSON files in deterministic order.""" paths = sorted(response_dir.glob("*.json")) if start < 0: raise ValueError("--start must be 0 or greater.") if limit is not None and limit < 1: raise ValueError("--limit must be 1 or greater.") if limit is None: return paths[start:] return paths[start : start + limit] def run(args: argparse.Namespace) -> None: """Download all requested NSRDB polygon archives.""" paths = response_paths(args.response_dir, args.start, args.limit) if not paths: print(f"No response JSON files found in {args.response_dir}") return counts = { ArchiveDownloadState.DOWNLOADED: 0, ArchiveDownloadState.SKIPPED: 0, ArchiveDownloadState.PENDING: 0, ArchiveDownloadState.FAILED: 0, } total = len(paths) for index, response_path in enumerate(paths, start=1): output_path = args.output_dir / output_name_for_response(response_path) label = response_path.name machine = ArchiveDownloadStateMachine(label) machine.transition(ArchiveDownloadEvent.START) try: response = read_response(response_path) download_url = download_url_from_response(response) state = download_file( download_url, output_path, args.timeout, args.overwrite, machine, ) except (OSError, ValueError, PolygonArchiveDownloadError) as error: if machine.state not in { ArchiveDownloadState.FAILED, ArchiveDownloadState.SKIPPED, ArchiveDownloadState.PENDING, ArchiveDownloadState.DOWNLOADED, }: machine.transition(ArchiveDownloadEvent.DOWNLOAD_FAILED) counts[machine.state] += 1 print(f"[{index}/{total}] Failed {label}: {error}") continue counts[state] += 1 if state is ArchiveDownloadState.SKIPPED: print(f"[{index}/{total}] Skipped existing {output_path}") elif state is ArchiveDownloadState.PENDING: print( f"[{index}/{total}] Pending {label}: archive URL returned " "S3 HTTP 403; retry after NSRDB finishes generating it." ) else: print(f"[{index}/{total}] Downloaded {output_path}") if args.delay > 0 and index < total: time.sleep(args.delay) print( f"Finished: downloaded={counts[ArchiveDownloadState.DOWNLOADED]}, " f"skipped={counts[ArchiveDownloadState.SKIPPED]}, " f"pending={counts[ArchiveDownloadState.PENDING]}, " f"failed={counts[ArchiveDownloadState.FAILED]}, output_dir={args.output_dir}" ) def parse_args( argv: list[str] | None = None, description: str | None = None, ) -> argparse.Namespace: """Parse command-line arguments.""" parser = argparse.ArgumentParser(description=description or __doc__) parser.add_argument("--response-dir", type=Path) parser.add_argument("--output-dir", type=Path) parser.add_argument("--start", type=int, default=0, help="Zero-based response file offset.") parser.add_argument("--limit", type=int, help="Number of response files to download.") parser.add_argument("--timeout", type=int, default=DEFAULT_TIMEOUT, help="Download timeout in seconds.") parser.add_argument("--delay", type=float, default=0.0, help="Seconds to wait between downloads.") parser.add_argument("--overwrite", action="store_true", help="Download even when the ZIP already exists.") args = parser.parse_args(argv) missing = [ option for option, value in { "--response-dir": args.response_dir, "--output-dir": args.output_dir, }.items() if not value ] if missing: parser.error(f"the following arguments are required: {', '.join(missing)}") if args.timeout <= 0: parser.error("--timeout must be greater than 0.") if args.delay < 0: parser.error("--delay must be 0 or greater.") return args def main( argv: list[str] | None = None, description: str | None = None, ) -> None: """Run the archive downloader.""" run(parse_args(argv, description)) if __name__ == "__main__": main()