Initial Climate Mood Analysis project
This commit is contained in:
Justin Fisher
@@ -0,0 +1,325 @@
|
||||
#!/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()
|
||||
Reference in New Issue
Block a user