dataapp_crawler 0.1.0__py3-none-any.whl

dataapp_crawler/checkpoint.py

@@ -0,0 +1,487 @@
+"""Versioned, bounded crawler state for Runtime-managed persistence.
+
+Checkpoint documents contain target URLs and developer-owned request data, so every
+export must be stored as a private artifact. Request headers and bodies require the
+additional ``include_sensitive=True`` opt-in. The SDK only models state; storing the
+returned document is a Runtime responsibility.
+"""
+
+from __future__ import annotations
+
+import base64
+import binascii
+import json
+import math
+from collections.abc import Iterable, Mapping
+from dataclasses import dataclass
+from itertools import islice
+from typing import Any
+
+from .errors import ConfigurationError
+from .models import Request, Stats
+
+_SCHEMA = "dataapp_crawler-checkpoint"
+_VERSION = 1
+_MAX_REQUESTS = 10_000
+_MAX_KEYS = 20_000
+_MAX_DOCUMENT_BYTES = 16 * 1024 * 1024
+_MAX_STRING_LENGTH = 4 * 1024 * 1024
+_MAX_JSON_DEPTH = 32
+_MAX_JSON_NODES = 100_000
+
+_REQUEST_FIELDS = frozenset(
+    {
+        "url",
+        "method",
+        "headers",
+        "params",
+        "content",
+        "json",
+        "label",
+        "user_data",
+        "key",
+        "replayable",
+        "mode",
+        "site_key",
+        "session_key",
+        "deduplicate",
+        "wait_until",
+        "snapshot",
+    }
+)
+_STATS_FIELDS = frozenset(
+    {
+        "requests",
+        "pages",
+        "captured",
+        "results",
+        "retries",
+        "duplicates",
+        "failed",
+        "stop_reason",
+    }
+)
+_DOCUMENT_FIELDS = frozenset(
+    {
+        "schema",
+        "version",
+        "pending",
+        "inflight",
+        "dedup_keys",
+        "seen_keys",
+        "stats",
+        "elapsed_seconds",
+    }
+)
+
+
+@dataclass(frozen=True)
+class CheckpointStats:
+    """Immutable counters consumed before the checkpoint was captured."""
+
+    requests: int = 0
+    pages: int = 0
+    captured: int = 0
+    results: int = 0
+    retries: int = 0
+    duplicates: int = 0
+    failed: int = 0
+
+    def to_stats(self) -> Stats:
+        return Stats(
+            requests=self.requests,
+            pages=self.pages,
+            captured=self.captured,
+            results=self.results,
+            retries=self.retries,
+            duplicates=self.duplicates,
+            failed=self.failed,
+            stop_reason=None,
+        )
+
+
+class CrawlCheckpoint:
+    """An immutable in-memory checkpoint with at-least-once inflight recovery."""
+
+    __slots__ = (
+        "_pending_records",
+        "_inflight_records",
+        "_dedup_keys",
+        "_seen_keys",
+        "_stats",
+        "_elapsed_seconds",
+    )
+
+    def __init__(
+        self,
+        *,
+        pending_records: tuple[str, ...],
+        inflight_records: tuple[str, ...],
+        dedup_keys: frozenset[str],
+        seen_keys: frozenset[str],
+        stats: CheckpointStats,
+        elapsed_seconds: float,
+    ) -> None:
+        self._pending_records = pending_records
+        self._inflight_records = inflight_records
+        self._dedup_keys = dedup_keys
+        self._seen_keys = seen_keys
+        self._stats = stats
+        self._elapsed_seconds = elapsed_seconds
+
+    def __repr__(self) -> str:
+        return (
+            "CrawlCheckpoint("
+            f"pending={len(self._pending_records)}, inflight={len(self._inflight_records)}, "
+            f"dedup_keys={len(self._dedup_keys)}, seen_keys={len(self._seen_keys)}, "
+            f"elapsed_seconds={self._elapsed_seconds!r})"
+        )
+
+    @classmethod
+    def capture(
+        cls,
+        *,
+        pending: Iterable[Request],
+        inflight: Iterable[Request],
+        dedup_keys: Iterable[str],
+        seen_keys: Iterable[str],
+        stats: Stats,
+        elapsed_seconds: float,
+    ) -> CrawlCheckpoint:
+        pending_records = _bounded_request_records(pending, _MAX_REQUESTS)
+        inflight_records = _bounded_request_records(inflight, _MAX_REQUESTS - len(pending_records))
+        for request in (_request_from_record(record) for record in inflight_records):
+            if request.replayable is False or (
+                request.method not in {"GET", "HEAD"} and request.replayable is not True
+            ):
+                raise ConfigurationError(
+                    "Inflight requests must use GET/HEAD or explicitly set replayable=True"
+                )
+        checkpoint_stats = _checkpoint_stats(stats)
+        elapsed = _nonnegative_finite(elapsed_seconds, "elapsed_seconds")
+        dedup = _keys(
+            set(_keys(dedup_keys, "dedup_keys"))
+            | {
+                request.unique_key
+                for record in (*pending_records, *inflight_records)
+                if (request := _request_from_record(record)).deduplicate
+            },
+            "dedup_keys",
+        )
+        seen = _keys(seen_keys, "seen_keys")
+        return cls(
+            pending_records=pending_records,
+            inflight_records=inflight_records,
+            dedup_keys=dedup,
+            seen_keys=seen,
+            stats=checkpoint_stats,
+            elapsed_seconds=elapsed,
+        )
+
+    @property
+    def pending_requests(self) -> tuple[Request, ...]:
+        return tuple(_request_from_record(record) for record in self._pending_records)
+
+    @property
+    def inflight_requests(self) -> tuple[Request, ...]:
+        return tuple(_request_from_record(record) for record in self._inflight_records)
+
+    @property
+    def requeued_requests(self) -> tuple[Request, ...]:
+        """Return inflight first: restored work is intentionally at-least-once."""
+        return self.inflight_requests + self.pending_requests
+
+    @property
+    def dedup_keys(self) -> frozenset[str]:
+        return self._dedup_keys
+
+    @property
+    def seen_keys(self) -> frozenset[str]:
+        return self._seen_keys
+
+    @property
+    def stats(self) -> Stats:
+        """Return a fresh mutable Stats value with the prior stop reason cleared."""
+        return self._stats.to_stats()
+
+    @property
+    def elapsed(self) -> float:
+        return self._elapsed_seconds
+
+    @property
+    def elapsed_seconds(self) -> float:
+        return self._elapsed_seconds
+
+    def to_dict(self, *, include_sensitive: bool = False) -> dict[str, Any]:
+        _validate_sensitive_opt_in(include_sensitive)
+        requests = [
+            *(_request_from_record(record) for record in self._pending_records),
+            *(_request_from_record(record) for record in self._inflight_records),
+        ]
+        if not include_sensitive and any(
+            _has_sensitive_request_state(request) for request in requests
+        ):
+            raise ConfigurationError(
+                "Checkpoint contains request headers or body data; "
+                "pass include_sensitive=True and store the artifact privately"
+            )
+        document = {
+            "schema": _SCHEMA,
+            "version": _VERSION,
+            "pending": [_load_record(record) for record in self._pending_records],
+            "inflight": [_load_record(record) for record in self._inflight_records],
+            "dedup_keys": sorted(self._dedup_keys),
+            "seen_keys": sorted(self._seen_keys),
+            "stats": {**self._stats.__dict__, "stop_reason": None},
+            "elapsed_seconds": self._elapsed_seconds,
+        }
+        _document_size(document)
+        return document
+
+    @classmethod
+    def from_dict(
+        cls, document: Mapping[str, Any], *, include_sensitive: bool = False
+    ) -> CrawlCheckpoint:
+        if not isinstance(document, Mapping):
+            raise ConfigurationError("Checkpoint must be a JSON object")
+        copied = _json_copy(dict(document), "checkpoint")
+        _document_size(copied)
+        _exact_fields(copied, _DOCUMENT_FIELDS, "checkpoint")
+        if copied["schema"] != _SCHEMA:
+            raise ConfigurationError("Unknown checkpoint schema")
+        if type(copied["version"]) is not int or copied["version"] != _VERSION:
+            raise ConfigurationError("Unsupported checkpoint version")
+        pending = _request_list(copied["pending"], "pending")
+        inflight = _request_list(copied["inflight"], "inflight")
+        _validate_sensitive_opt_in(include_sensitive)
+        if not include_sensitive and any(
+            _has_sensitive_request_state(request) for request in (*pending, *inflight)
+        ):
+            raise ConfigurationError(
+                "Checkpoint contains request headers or body data; "
+                "pass include_sensitive=True and store the artifact privately"
+            )
+        return cls.capture(
+            pending=pending,
+            inflight=inflight,
+            dedup_keys=_key_list(copied["dedup_keys"], "dedup_keys"),
+            seen_keys=_key_list(copied["seen_keys"], "seen_keys"),
+            stats=_stats_from_document(copied["stats"]),
+            elapsed_seconds=copied["elapsed_seconds"],
+        )
+
+
+def _validate_sensitive_opt_in(include_sensitive: bool) -> None:
+    if not isinstance(include_sensitive, bool):
+        raise ConfigurationError("include_sensitive must be a boolean")
+
+
+def _has_sensitive_request_state(request: Request) -> bool:
+    return (
+        bool(request.headers) or request.content not in {None, "", b""} or request.json is not None
+    )
+
+
+def _request_record(request: Request) -> str:
+    if not isinstance(request, Request):
+        raise ConfigurationError("Checkpoint queues must contain Request values")
+    content: Any
+    if isinstance(request.content, bytes):
+        content = {"encoding": "base64", "data": base64.b64encode(request.content).decode("ascii")}
+    elif request.content is None or isinstance(request.content, str):
+        content = request.content
+    else:
+        raise ConfigurationError("Request content must be bytes, text or None")
+    value = {
+        "url": request.url,
+        "method": request.method,
+        "headers": request.headers,
+        "params": request.params,
+        "content": content,
+        "json": request.json,
+        "label": request.label,
+        "user_data": request.user_data,
+        "key": request.key,
+        "replayable": request.replayable,
+        "mode": request.mode,
+        "site_key": request.site_key,
+        "session_key": request.session_key,
+        "deduplicate": request.deduplicate,
+        "wait_until": request.wait_until,
+        "snapshot": request.snapshot,
+    }
+    copied = _json_copy(value, "request")
+    return json.dumps(copied, ensure_ascii=True, separators=(",", ":"), sort_keys=True)
+
+
+def _request_from_record(record: str) -> Request:
+    value = _load_record(record)
+    _exact_fields(value, _REQUEST_FIELDS, "request")
+    for name in ("url", "method", "label", "mode", "wait_until"):
+        if not isinstance(value[name], str):
+            raise ConfigurationError(f"request.{name} must be a string")
+    for name in ("key", "site_key", "session_key"):
+        if value[name] is not None and not isinstance(value[name], str):
+            raise ConfigurationError(f"request.{name} must be a string or null")
+    if value["replayable"] is not None and not isinstance(value["replayable"], bool):
+        raise ConfigurationError("request.replayable must be a boolean or null")
+    for name in ("deduplicate", "snapshot"):
+        if not isinstance(value[name], bool):
+            raise ConfigurationError(f"request.{name} must be a boolean")
+    if not isinstance(value["headers"], dict) or any(
+        not isinstance(key, str) or not isinstance(item, str)
+        for key, item in value["headers"].items()
+    ):
+        raise ConfigurationError("request.headers must map strings to strings")
+    if value["params"] is not None and not isinstance(value["params"], dict):
+        raise ConfigurationError("request.params must be an object or null")
+    if not isinstance(value["user_data"], dict):
+        raise ConfigurationError("request.user_data must be an object")
+    content = value["content"]
+    if isinstance(content, dict):
+        _exact_fields(content, frozenset({"encoding", "data"}), "request content")
+        if content["encoding"] != "base64" or not isinstance(content["data"], str):
+            raise ConfigurationError("Invalid encoded request content")
+        try:
+            content = base64.b64decode(content["data"], validate=True)
+        except (ValueError, binascii.Error):
+            raise ConfigurationError("Invalid base64 request content") from None
+    elif content is not None and not isinstance(content, str):
+        raise ConfigurationError("Request content must be encoded bytes, text or None")
+    try:
+        return Request(**{**value, "content": content})
+    except (TypeError, ValueError, ConfigurationError) as error:
+        if isinstance(error, ConfigurationError):
+            raise
+        raise ConfigurationError("Invalid request in checkpoint") from None
+
+
+def _request_list(value: Any, name: str) -> tuple[Request, ...]:
+    if not isinstance(value, list):
+        raise ConfigurationError(f"{name} must be a JSON array")
+    if len(value) > _MAX_REQUESTS:
+        raise ConfigurationError("Checkpoint contains too many requests")
+    requests: list[Request] = []
+    for item in value:
+        if not isinstance(item, dict):
+            raise ConfigurationError(f"{name} entries must be JSON objects")
+        requests.append(_request_from_record(json.dumps(item, separators=(",", ":"))))
+    return tuple(requests)
+
+
+def _bounded_request_records(values: Iterable[Request], limit: int) -> tuple[str, ...]:
+    records = tuple(_request_record(request) for request in islice(values, limit + 1))
+    if len(records) > limit:
+        raise ConfigurationError("Checkpoint contains too many requests")
+    return records
+
+
+def _checkpoint_stats(stats: Stats) -> CheckpointStats:
+    if not isinstance(stats, Stats):
+        raise ConfigurationError("stats must be a Stats value")
+    values = {}
+    for name in CheckpointStats.__dataclass_fields__:
+        value = getattr(stats, name)
+        if type(value) is not int or value < 0:
+            raise ConfigurationError(f"stats.{name} must be a nonnegative integer")
+        values[name] = value
+    return CheckpointStats(**values)
+
+
+def _stats_from_document(value: Any) -> Stats:
+    if not isinstance(value, dict):
+        raise ConfigurationError("stats must be a JSON object")
+    _exact_fields(value, _STATS_FIELDS, "stats")
+    if value["stop_reason"] is not None:
+        raise ConfigurationError("Restored checkpoint stop_reason must be null")
+    fields = {name: value[name] for name in CheckpointStats.__dataclass_fields__}
+    for name, item in fields.items():
+        if type(item) is not int or item < 0:
+            raise ConfigurationError(f"stats.{name} must be a nonnegative integer")
+    return Stats(**fields, stop_reason=None)
+
+
+def _keys(values: Iterable[str], name: str) -> frozenset[str]:
+    collected: list[str] = []
+    for value in islice(values, _MAX_KEYS + 1):
+        if not isinstance(value, str) or not value or len(value) > _MAX_STRING_LENGTH:
+            raise ConfigurationError(f"{name} must contain bounded non-empty strings")
+        collected.append(value)
+    if len(collected) > _MAX_KEYS:
+        raise ConfigurationError(f"{name} contains too many values")
+    return frozenset(collected)
+
+
+def _key_list(value: Any, name: str) -> tuple[str, ...]:
+    if not isinstance(value, list):
+        raise ConfigurationError(f"{name} must be a JSON array")
+    if len(value) != len(set(item for item in value if isinstance(item, str))):
+        raise ConfigurationError(f"{name} must not contain duplicate values")
+    return tuple(_keys(value, name))
+
+
+def _nonnegative_finite(value: Any, name: str) -> float:
+    if isinstance(value, bool) or not isinstance(value, (int, float)):
+        raise ConfigurationError(f"{name} must be a finite nonnegative number")
+    result = float(value)
+    if not math.isfinite(result) or result < 0:
+        raise ConfigurationError(f"{name} must be a finite nonnegative number")
+    return result
+
+
+def _json_copy(value: Any, name: str) -> Any:
+    nodes = 0
+
+    def validate(item: Any, depth: int) -> None:
+        nonlocal nodes
+        nodes += 1
+        if nodes > _MAX_JSON_NODES or depth > _MAX_JSON_DEPTH:
+            raise ConfigurationError(f"{name} exceeds checkpoint JSON limits")
+        if item is None or isinstance(item, (bool, int)):
+            return
+        if isinstance(item, float):
+            if not math.isfinite(item):
+                raise ConfigurationError(f"{name} contains a non-finite number")
+            return
+        if isinstance(item, str):
+            if len(item) > _MAX_STRING_LENGTH:
+                raise ConfigurationError(f"{name} contains an oversized string")
+            return
+        if isinstance(item, list):
+            for child in item:
+                validate(child, depth + 1)
+            return
+        if isinstance(item, dict):
+            for key, child in item.items():
+                if not isinstance(key, str) or len(key) > _MAX_STRING_LENGTH:
+                    raise ConfigurationError(f"{name} object keys must be bounded strings")
+                validate(child, depth + 1)
+            return
+        raise ConfigurationError(f"{name} contains a non-JSON value")
+
+    validate(value, 0)
+    try:
+        return json.loads(json.dumps(value, allow_nan=False))
+    except (TypeError, ValueError):
+        raise ConfigurationError(f"{name} must be strict JSON") from None
+
+
+def _load_record(record: str) -> dict[str, Any]:
+    value = json.loads(record)
+    if not isinstance(value, dict):
+        raise ConfigurationError("Invalid internal checkpoint request")
+    return value
+
+
+def _document_size(document: Any) -> None:
+    try:
+        encoded = json.dumps(document, allow_nan=False, separators=(",", ":")).encode()
+    except (TypeError, ValueError):
+        raise ConfigurationError("Checkpoint must be strict JSON") from None
+    if len(encoded) > _MAX_DOCUMENT_BYTES:
+        raise ConfigurationError("Checkpoint exceeds the maximum document size")
+
+
+def _exact_fields(value: Mapping[str, Any], expected: frozenset[str], name: str) -> None:
+    if set(value) != expected:
+        raise ConfigurationError(f"{name} fields do not match the checkpoint schema")