252 lines
9.7 KiB
Python
252 lines
9.7 KiB
Python
"""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
|