politics-tracking/backend/politrack/core/fetch.py

"""The Fetcher: the only component that talks to the network. It enforces the crawling policy:

  * honest User-Agent with contact address,
  * robots.txt (cached per origin); exemptions only for documented APIs listed in SourceSpec.robots_exempt,
  * per-host throttling (max of SourceSpec.min_interval and robots.txt crawl-delay),
  * conditional requests (ETag / Last-Modified from the raw archive),
  * retries with exponential back-off for 429/5xx, honouring Retry-After,
  * TDM reservations (W3C TDMRep: `tdm-reservation` header, /.well-known/tdmrep.json, /tdmrep.json),
  * archiving of raw bodies (unless SourceSpec.keep_raw is False or TDM is reserved).
"""

from __future__ import annotations

import logging
import threading
import time
from datetime import datetime, timezone
from urllib.parse import urlsplit

import httpx
from protego import Protego
from psycopg import Connection

from politrack.config import get_settings
from politrack.core.adapter import Request, Response, redact_url
from politrack.core.archive import Archive
from politrack.core.spec import SourceSpec

log = logging.getLogger(__name__)
RETRY_STATUS = {429, 500, 502, 503, 504}


class Disallowed(Exception):
    """robots.txt forbids fetching this URL."""


class Fetcher:
    def __init__(self, archive: Archive | None = None, client: httpx.Client | None = None) -> None:
        settings = get_settings()
        self.user_agent = settings.user_agent
        self.default_interval = settings.min_request_interval
        self.archive = archive or Archive()
        self.client = client or httpx.Client(
            headers={"User-Agent": self.user_agent, "Accept-Language": "de-AT,de;q=0.9,en;q=0.5"},
            timeout=settings.http_timeout,
            follow_redirects=True,
        )
        self._robots: dict[str, Protego | None] = {}
        self._tdm: dict[str, bool] = {}
        self._last: dict[str, float] = {}
        self._lock = threading.Lock()

    # -- policy -------------------------------------------------------------------------

    @staticmethod
    def origin(url: str) -> str:
        parts = urlsplit(url)
        return f"{parts.scheme}://{parts.netloc}"

    def robots(self, url: str) -> Protego | None:
        origin = self.origin(url)
        if origin not in self._robots:
            try:
                resp = self.client.get(f"{origin}/robots.txt")
                self._robots[origin] = Protego.parse(resp.text) if resp.status_code == 200 else None
            except httpx.HTTPError:
                self._robots[origin] = None
        return self._robots[origin]

    def allowed(self, url: str, spec: SourceSpec) -> bool:
        if any(url.startswith(prefix) for prefix in spec.robots_exempt):
            return True
        rules = self.robots(url)
        return rules is None or rules.can_fetch(url, self.user_agent)

    def tdm_reserved(self, url: str) -> bool:
        origin = self.origin(url)
        if origin not in self._tdm:
            reserved = False
            for path in ("/.well-known/tdmrep.json", "/tdmrep.json"):
                try:
                    resp = self.client.get(origin + path)
                except httpx.HTTPError:
                    continue
                if resp.status_code == 200:
                    try:
                        rules = resp.json()
                    except ValueError:
                        continue
                    reserved = any(str(r.get("tdm-reservation")) == "1" for r in rules if isinstance(r, dict))
                    break
            self._tdm[origin] = reserved
        return self._tdm[origin]

    def _throttle(self, url: str, spec: SourceSpec) -> None:
        rules = self._robots.get(self.origin(url))
        delay = rules.crawl_delay(self.user_agent) if rules else None
        interval = max(spec.min_interval or self.default_interval, float(delay or 0))
        host = urlsplit(url).netloc
        with self._lock:
            wait = self._last.get(host, 0) + interval - time.monotonic()
            if wait > 0:
                time.sleep(wait)
            self._last[host] = time.monotonic()

    # -- fetching ---------------------------------------------------------------------

    def fetch(self, conn: Connection, req: Request, *, adapter: str, spec: SourceSpec, source_id: int | None,
              attempts: int = 5) -> Response:
        if not self.allowed(req.url, spec):
            raise Disallowed(req.url)
        last = self.archive.last(conn, req.fingerprint) if req.conditional else None
        headers = dict(req.headers or {})
        if last and last["stored"]:
            if last["etag"]:
                headers["If-None-Match"] = last["etag"]
            if last["last_modified"]:
                headers["If-Modified-Since"] = last["last_modified"]

        resp: httpx.Response | None = None
        for attempt in range(attempts):
            self._throttle(req.url, spec)
            try:
                resp = self.client.request(req.method, req.url, params=req.params, json=req.json, headers=headers)
            except httpx.TransportError as exc:
                if attempt == attempts - 1:
                    raise
                log.warning("transport error %s (%s), retrying", req.url, exc)
                time.sleep(min(5 * 2 ** attempt, 120))
                continue
            if resp.status_code in RETRY_STATUS and attempt < attempts - 1:
                retry_after = resp.headers.get("Retry-After", "")
                time.sleep(min(float(retry_after) if retry_after.isdigit() else 5 * 2 ** attempt, 300))
                continue
            break
        assert resp is not None
        tdm = resp.headers.get("tdm-reservation") == "1" or self.tdm_reserved(req.url)
        now = datetime.now(timezone.utc)

        if resp.status_code == 304 and last:
            body = self.archive.read_blob(last["sha256"]) or b""
            return Response(request=req, status=200, url=last["url"], headers={"content-type": last["content_type"] or ""},
                            body=body, fetched_at=now, sha256=last["sha256"], fetch_id=last["id"],
                            tdm_reserved=last["tdm_reserved"], from_archive=True)

        lowered = {k.lower(): v for k, v in resp.headers.items()}
        fetch_id, sha = self.archive.log(
            conn, adapter=adapter, source_id=source_id, request_fp=req.fingerprint, url=redact_url(str(resp.url)),
            method=req.method, request=req.describe(), status=resp.status_code, body=resp.content, headers=lowered,
            keep=spec.keep_raw, tdm_reserved=tdm,
        )
        conn.commit()
        resp.raise_for_status()
        return Response(request=req, status=resp.status_code, url=redact_url(str(resp.url)), headers=lowered, body=resp.content,
                        fetched_at=now, sha256=sha or "", fetch_id=fetch_id, tdm_reserved=tdm)

    def close(self) -> None:
        self.client.close()