Justin Fisher
326 lines
11 KiB
Python
326 lines
11 KiB
Python
#!/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()
|