web-fetcher/web_fetcher/core.py

"""Core types: Observation record and provenance-tracked HTTP fetching."""

from __future__ import annotations

import hashlib
import json
import os
import time
import urllib.error
import urllib.request
from dataclasses import asdict, dataclass, field
from datetime import datetime, timezone
from email.utils import parsedate_to_datetime
from pathlib import Path
from typing import Any

USER_AGENT = "web-fetcher/0.2 (+https://forgejo.coulomb.social/coulomb/web-fetcher)"


@dataclass
class Observation:
    """One fetched value with its provenance. status: ok | not_found | error."""

    spec_id: str
    status: str
    value: Any = None
    unit: str | None = None
    period: str | None = None
    labels: dict[str, str] = field(default_factory=dict)
    source: str | None = None          # dataset or source name
    url: str | None = None
    published: str | None = None       # source-side update/publication timestamp
    retrieved_at: str | None = None
    payload_sha256: str | None = None
    note: str | None = None

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)

    def to_json(self) -> str:
        return json.dumps(self.to_dict(), ensure_ascii=False, sort_keys=True)


@dataclass
class Fetched:
    url: str
    body: bytes
    retrieved_at: str
    sha256: str
    content_type: str | None
    published: str | None = None       # HTTP Last-Modified, ISO 8601, when the server sends it


class FetchError(RuntimeError):
    pass


RETRY_STATUS = {429, 500, 502, 503, 504}
MAX_RETRY_AFTER = 120.0

# Indirection points for offline tests.
_urlopen = urllib.request.urlopen
_sleep = time.sleep


def _now() -> str:
    return datetime.now(timezone.utc).isoformat(timespec="seconds")


def _http_date(value: str | None) -> str | None:
    if not value:
        return None
    try:
        return parsedate_to_datetime(value).astimezone(timezone.utc).isoformat(timespec="seconds")
    except (TypeError, ValueError):
        return None


def _retry_after(value: str | None) -> float | None:
    """Seconds to wait from a Retry-After header (delta-seconds or HTTP date)."""
    if not value:
        return None
    try:
        return max(0.0, float(value))
    except ValueError:
        when = _http_date(value)
        if when is None:
            return None
        return max(0.0, (datetime.fromisoformat(when) - datetime.now(timezone.utc)).total_seconds())


def fetch_bytes(url: str, *, timeout: float = 60.0, retries: int = 2,
                cache_dir: str | os.PathLike | None = None,
                replay_dir: str | os.PathLike | None = None,
                headers: dict[str, str] | None = None) -> Fetched:
    """GET `url` and return body plus provenance. Optional content-addressed cache.

    The cache stores each payload under its SHA-256 and an index entry per URL.
    With `replay_dir`, no network request is made: the payload and its original
    provenance are served from such a cache (see `replay_bytes`).

    Only transient failures are retried (connection errors, timeouts, HTTP
    429/5xx), with exponential backoff or the server's Retry-After; other HTTP
    errors fail immediately. Provenance describes the attempt that succeeded.
    """
    if replay_dir:
        return replay_bytes(url, replay_dir)
    hdrs = {"User-Agent": USER_AGENT, **(headers or {})}
    last: Exception | None = None
    for attempt in range(retries + 1):
        wait = 1.5 * 2 ** attempt
        try:
            req = urllib.request.Request(url, headers=hdrs)
            with _urlopen(req, timeout=timeout) as resp:
                body = resp.read()
                ctype = resp.headers.get("Content-Type")
                modified = _http_date(resp.headers.get("Last-Modified"))
            digest = hashlib.sha256(body).hexdigest()
            fetched = Fetched(url=url, body=body, retrieved_at=_now(), sha256=digest,
                              content_type=ctype, published=modified)
            if cache_dir:
                _cache_store(Path(cache_dir), fetched)
            return fetched
        except urllib.error.HTTPError as exc:
            if exc.code not in RETRY_STATUS:
                raise FetchError(f"fetch failed for {url}: HTTP {exc.code} {exc.reason}") from exc
            last = exc
            hinted = _retry_after(exc.headers.get("Retry-After") if exc.headers else None)
            if hinted is not None:
                wait = min(hinted, MAX_RETRY_AFTER)
        except (urllib.error.URLError, TimeoutError, OSError) as exc:
            last = exc
        if attempt < retries:
            _sleep(wait)
    raise FetchError(f"fetch failed for {url}: {last}")


def replay_bytes(url: str, cache_dir: str | os.PathLike) -> Fetched:
    """Serve `url` from a cache written by `fetch_bytes(cache_dir=...)`, without network.

    Uses the most recent index entry for the URL and verifies the payload hash.
    Returns the original retrieval time, so a replayed run reproduces the
    observations of the run that filled the cache.
    """
    cache = Path(cache_dir)
    entry: dict[str, Any] | None = None
    try:
        with (cache / "index.jsonl").open(encoding="utf-8") as fh:
            for line in fh:
                if line.strip():
                    rec = json.loads(line)
                    if rec.get("url") == url:
                        entry = rec
    except FileNotFoundError:
        raise FetchError(f"replay cache {cache} has no index.jsonl") from None
    if entry is None:
        raise FetchError(f"not in replay cache: {url}")
    try:
        body = (cache / f"{entry['sha256']}.bin").read_bytes()
    except FileNotFoundError:
        raise FetchError(f"replay payload missing for {url}: {entry['sha256']}") from None
    if hashlib.sha256(body).hexdigest() != entry["sha256"]:
        raise FetchError(f"replay payload hash mismatch for {url}")
    return Fetched(url=url, body=body, retrieved_at=entry["retrieved_at"], sha256=entry["sha256"],
                   content_type=entry.get("content_type"), published=entry.get("published"))


def _cache_store(cache: Path, fetched: Fetched) -> None:
    cache.mkdir(parents=True, exist_ok=True)
    (cache / f"{fetched.sha256}.bin").write_bytes(fetched.body)
    with (cache / "index.jsonl").open("a", encoding="utf-8") as fh:
        fh.write(json.dumps({"url": fetched.url, "sha256": fetched.sha256,
                             "retrieved_at": fetched.retrieved_at,
                             "content_type": fetched.content_type,
                             "published": fetched.published}) + "\n")