"""Shared resumable download helpers for collectors.""" from __future__ import annotations import hashlib import json import os import time from datetime import UTC, datetime from pathlib import Path from typing import Any, Awaitable, Callable import httpx ProgressCallback = Callable[[int, int | None], Awaitable[None]] ValidateCallback = Callable[[Path], bool] class DownloadHTTPStatusError(RuntimeError): """HTTP status error that keeps the upstream response body for caller-specific handling.""" def __init__(self, *, url: str, status_code: int, body: str) -> None: self.url = url self.status_code = status_code self.body = body preview = body.strip().replace("\r", " ").replace("\n", " ")[:240] suffix = f": {preview}" if preview else "" super().__init__(f"HTTP {status_code} while downloading {url}{suffix}") def default_download_cache_root() -> Path: configured = os.getenv("PLANET_DOWNLOAD_CACHE_DIR") if configured: return Path(configured).expanduser() planet_cache = os.getenv("PLANET_CACHE_DIR") if planet_cache: return Path(planet_cache).expanduser() / "downloads" xdg_cache = os.getenv("XDG_CACHE_HOME") if xdg_cache: return Path(xdg_cache).expanduser() / "planet" / "downloads" return Path.home() / ".cache" / "planet" / "downloads" class ResumableFileDownloader: """Download files with cache validators and byte-range resume support.""" def __init__( self, *, cache_namespace: str, user_agent: str = "Planet-Intelligence-System/1.0 (Python/collector)", default_accept: str = "*/*", cache_root: Path | None = None, ) -> None: self._cache_dir = (cache_root or default_download_cache_root()) / cache_namespace self._user_agent = user_agent self._default_accept = default_accept @staticmethod def _cache_key(url: str) -> str: return hashlib.sha1(url.encode("utf-8")).hexdigest()[:16] def _cache_paths(self, url: str, extension: str) -> tuple[Path, Path, Path]: key = self._cache_key(url) normalized_ext = extension if extension.startswith(".") else f".{extension}" final_path = self._cache_dir / f"{key}{normalized_ext}" part_path = self._cache_dir / f"{key}{normalized_ext}.part" meta_path = self._cache_dir / f"{key}.meta.json" return final_path, part_path, meta_path def cached_file_path(self, url: str, extension: str) -> Path: final_path, _, _ = self._cache_paths(url, extension) return final_path def get_cached_file( self, url: str, extension: str, *, validate_existing: ValidateCallback | None = None, ) -> Path | None: final_path = self.cached_file_path(url, extension) if not final_path.exists(): return None if validate_existing and not validate_existing(final_path): final_path.unlink(missing_ok=True) return None return final_path @staticmethod def _load_meta(meta_path: Path) -> dict[str, Any]: if not meta_path.exists(): return {} try: return json.loads(meta_path.read_text(encoding="utf-8")) except (json.JSONDecodeError, OSError): return {} @staticmethod def _save_meta(meta_path: Path, payload: dict[str, Any]) -> None: meta_path.write_text(json.dumps(payload, ensure_ascii=False), encoding="utf-8") @staticmethod def _validators_match(meta: dict[str, Any], remote: dict[str, Any]) -> bool: etag = str(remote.get("etag") or "").strip() last_modified = str(remote.get("last_modified") or "").strip() if etag: return etag == str(meta.get("etag") or "").strip() if last_modified: return last_modified == str(meta.get("last_modified") or "").strip() return True async def fetch_remote_info(self, client: httpx.AsyncClient, url: str) -> dict[str, Any]: try: response = await client.head(url) if response.status_code >= 400: return {} content_length_raw = response.headers.get("content-length") content_length = int(content_length_raw) if content_length_raw else None return { "etag": response.headers.get("etag"), "last_modified": response.headers.get("last-modified"), "content_length": content_length, "accept_ranges": (response.headers.get("accept-ranges") or "").lower(), } except (httpx.HTTPError, ValueError): return {} async def download_file( self, client: httpx.AsyncClient, url: str, *, extension: str, accept: str | None = None, progress_callback: ProgressCallback | None = None, validate_existing: ValidateCallback | None = None, ) -> Path: self._cache_dir.mkdir(parents=True, exist_ok=True) final_path, part_path, meta_path = self._cache_paths(url, extension) meta = self._load_meta(meta_path) remote = await self.fetch_remote_info(client, url) expected_size = remote.get("content_length") if final_path.exists(): local_size = final_path.stat().st_size size_match = expected_size is None or local_size == expected_size if self._validators_match(meta, remote) and size_match: if validate_existing and not validate_existing(final_path): final_path.unlink(missing_ok=True) else: if progress_callback and expected_size and expected_size > 0: await progress_callback(expected_size, expected_size) return final_path can_resume = (remote.get("accept_ranges") or "") == "bytes" resume_from = part_path.stat().st_size if part_path.exists() else 0 if expected_size is not None and resume_from > expected_size: part_path.unlink(missing_ok=True) resume_from = 0 if not self._validators_match(meta, remote): part_path.unlink(missing_ok=True) resume_from = 0 headers = { "User-Agent": self._user_agent, "Accept": accept or self._default_accept, } if final_path.exists(): if meta.get("etag"): headers["If-None-Match"] = str(meta.get("etag")) elif meta.get("last_modified"): headers["If-Modified-Since"] = str(meta.get("last_modified")) if can_resume and resume_from > 0: headers["Range"] = f"bytes={resume_from}-" if remote.get("etag"): headers["If-Range"] = str(remote.get("etag")) elif remote.get("last_modified"): headers["If-Range"] = str(remote.get("last_modified")) async with client.stream("GET", url, headers=headers) as response: if response.status_code == 304 and final_path.exists(): if progress_callback and expected_size and expected_size > 0: await progress_callback(expected_size, expected_size) return final_path if response.status_code >= 400: body = (await response.aread()).decode("utf-8", errors="replace") raise DownloadHTTPStatusError(url=url, status_code=response.status_code, body=body) if response.status_code == 206 and resume_from > 0: mode = "ab" else: mode = "wb" resume_from = 0 downloaded = resume_from last_emit_bytes = 0 last_emit_time = time.monotonic() min_emit_bytes = ( max(expected_size // 150, 512 * 1024) if expected_size and expected_size > 0 else 1024 * 1024 ) with part_path.open(mode) as f: if progress_callback and downloaded > 0: await progress_callback(downloaded, expected_size) async for chunk in response.aiter_bytes(): if not chunk: continue f.write(chunk) downloaded += len(chunk) if not progress_callback: continue now = time.monotonic() should_emit = ( expected_size is None or downloaded >= expected_size or downloaded - last_emit_bytes >= min_emit_bytes or now - last_emit_time >= 2.0 ) if should_emit: last_emit_bytes = downloaded last_emit_time = now await progress_callback(downloaded, expected_size) final_size = part_path.stat().st_size if part_path.exists() else 0 if expected_size is not None and final_size != expected_size: raise RuntimeError( f"Resumable download incomplete for {url}: expected={expected_size}, got={final_size}" ) part_path.replace(final_path) self._save_meta( meta_path, { "url": url, "etag": remote.get("etag"), "last_modified": remote.get("last_modified"), "content_length": expected_size, "updated_at": datetime.now(UTC).isoformat(), }, ) if validate_existing and not validate_existing(final_path): raise RuntimeError(f"Downloaded file validation failed for {url}") if progress_callback and expected_size and expected_size > 0: await progress_callback(expected_size, expected_size) return final_path