mostlyright-data 0.9.0__py3-none-any.whl

mostlyright/data_harness/thin/acquire_cancel.py

@@ -0,0 +1,579 @@
+"""Cancel one still-active hosted public acquisition, once and durably.
+
+The acquisition receipt names the workspace and crawler-session coordinate.  Cancellation is a
+separate, authority-bearing act: this module owns its one V3 mutation rather than teaching the
+read-only run client a cancel route.  A small local journal is written *before* the request so a
+lost response is retried with the same idempotency key.  It deliberately stores only public
+coordinates, the idempotency key, and Studio's cancellation receipt -- never a device key, bearer
+token, Cloud URL, or environment variable name/value.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import re
+import secrets
+import stat
+from collections.abc import Mapping
+from datetime import datetime
+from pathlib import Path
+from typing import Any
+from uuid import UUID
+
+from mostlyright.data_harness.canonical import CanonicalJSONError, canonical_sha256, parse_json
+from mostlyright.data_harness.thin.runs import StudioApiClient
+from mostlyright.data_harness.thin.session import StudioSession, open_studio_session
+from mostlyright.data_harness.thin.transport import ThinLaneError
+from mostlyright.data_harness.ux.credentials import (
+    default_credentials_dir,
+    resolve_cloud_credentials,
+)
+
+CANCEL_PUBLIC_ACQUISITION_PATH = (
+    "/v3/workspaces/{workspace_id}/public-acquisitions/{crawler_session_id}:cancel"
+)
+CANCELLATION_RECEIPT_VERSION = "mostlyright-public-acquisition-cancellation-receipt.v1"
+CANCELLATION_JOURNAL_VERSION = "mostlyright-thin-client-public-acquisition-cancellation-journal.v1"
+MAX_JOURNAL_BYTES = 64 * 1024
+MAX_RETRY_AFTER_SECONDS = 900
+_JOURNAL_STATE_DIRECTORY = "mr-data/public-acquisition-cancellations"
+
+RECEIPT_FIELDS = frozenset(
+    {
+        "schema_version",
+        "crawler_session_id",
+        "workspace_id",
+        "requester_principal_id",
+        "status",
+        "cancelled_at",
+    }
+)
+JOURNAL_FIELDS = frozenset(
+    {
+        "schema_version",
+        "workspace_id",
+        "crawler_session_id",
+        "idempotency_key",
+        "state",
+        "receipt",
+    }
+)
+
+_UUID = re.compile(r"^[0-9a-f]{8}-[0-9a-f]{4}-[1-8][0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$")
+_IDEMPOTENCY = re.compile(r"^mr-data\.acquire-cancel\.v1:[0-9a-f]{64}$")
+_SECONDS_Z = re.compile(r"^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}Z$")
+
+
+class CancellationRefusal(ThinLaneError):
+    """A cancellation refusal whose retry facts the caller must be able to retain."""
+
+    def __init__(
+        self,
+        code: str,
+        detail: str,
+        *,
+        workspace_id: str | None = None,
+        crawler_session_id: str | None = None,
+        retry_after_seconds: int | None = None,
+    ) -> None:
+        super().__init__(code, detail)
+        self.workspace_id = workspace_id
+        self.crawler_session_id = crawler_session_id
+        self.retry_after_seconds = retry_after_seconds
+
+    def to_payload(self) -> dict[str, Any]:
+        payload: dict[str, Any] = {
+            "schema_version": "mr-data-error.v1",
+            "code": self.code,
+            "error": self.detail,
+            "headline": self.detail,
+        }
+        if self.retry_after_seconds is not None:
+            payload["retry_after_seconds"] = self.retry_after_seconds
+        if self.workspace_id is not None and self.crawler_session_id is not None:
+            payload["active_public_acquisition"] = {
+                "workspace_id": self.workspace_id,
+                "crawler_session_id": self.crawler_session_id,
+            }
+        return payload
+
+
+class StudioAcquisitionCancellationClient(StudioApiClient):
+    """The one public-acquisition cancellation route and no hidden follow-up mutation."""
+
+    def cancel(
+        self,
+        *,
+        workspace_id: str,
+        crawler_session_id: str,
+        idempotency_key: str,
+    ) -> tuple[dict[str, Any], bool]:
+        headers: dict[str, str] = {}
+        try:
+            receipt = self._call(
+                "POST",
+                CANCEL_PUBLIC_ACQUISITION_PATH.format(
+                    workspace_id=workspace_id,
+                    crawler_session_id=crawler_session_id,
+                ),
+                body={
+                    "schema_version": "3.0.0",
+                    "workspace_id": workspace_id,
+                    "crawler_session_id": crawler_session_id,
+                },
+                extra_headers={"Idempotency-Key": idempotency_key},
+                expected=(200,),
+                response_headers=headers,
+            )
+        except ThinLaneError as refusal:
+            status = getattr(refusal, "http_status", None)
+            if status == 429:
+                window = _retry_after(headers)
+                if window is None:
+                    raise ThinLaneError(
+                        "THIN_CANCELLATION_RETRY_AFTER_INVALID",
+                        "Studio rate-limited cancellation without a retry window from 1 to 900 "
+                        "seconds",
+                    ) from refusal
+                raise CancellationRefusal(
+                    refusal.code,
+                    f"{refusal.detail}; nothing was cancelled and Studio asked for a retry after "
+                    f"{window} seconds",
+                    workspace_id=workspace_id,
+                    crawler_session_id=crawler_session_id,
+                    retry_after_seconds=window,
+                ) from refusal
+            if status == 409:
+                raise CancellationRefusal(
+                    "THIN_ACQUISITION_CANCEL_TERMINAL",
+                    "Studio cannot cancel this public acquisition because it already completed "
+                    "or failed",
+                ) from refusal
+            raise
+        replay = headers.get("idempotent-replay")
+        if replay not in {"true", "false"}:
+            raise ThinLaneError(
+                "THIN_CANCELLATION_REPLAY_INVALID",
+                "Studio answered cancellation without Idempotent-Replay true or false",
+            )
+        return receipt, replay == "true"
+
+
+class CancellationJournal:
+    """A narrowly shaped, atomically written receipt journal for one coordinate."""
+
+    def __init__(self, path: Path) -> None:
+        self.path = path
+        self.state: dict[str, Any] = {}
+
+    def load_or_start(self, *, workspace_id: str, crawler_session_id: str) -> None:
+        _journal_location(self.path)
+        if self.path.exists():
+            self.state = _read_journal(self.path)
+            if (
+                self.state["workspace_id"] != workspace_id
+                or self.state["crawler_session_id"] != crawler_session_id
+            ):
+                raise ThinLaneError(
+                    "THIN_CANCELLATION_JOURNAL_SCOPE",
+                    "the cancellation journal belongs to a different public acquisition coordinate",
+                )
+            expected = _idempotency_key(workspace_id, crawler_session_id)
+            if self.state["idempotency_key"] != expected:
+                raise ThinLaneError(
+                    "THIN_CANCELLATION_JOURNAL_INVALID",
+                    "the cancellation journal has an idempotency key this client would not use",
+                )
+            return
+        self.state = {
+            "schema_version": CANCELLATION_JOURNAL_VERSION,
+            "workspace_id": workspace_id,
+            "crawler_session_id": crawler_session_id,
+            "idempotency_key": _idempotency_key(workspace_id, crawler_session_id),
+            "state": "prepared",
+            "receipt": None,
+        }
+        self.save()
+
+    def save(self) -> None:
+        rendered = json.dumps(self.state, ensure_ascii=False, sort_keys=True).encode("utf-8")
+        if len(rendered) > MAX_JOURNAL_BYTES:
+            raise ThinLaneError(
+                "THIN_CANCELLATION_JOURNAL_INVALID",
+                "the cancellation journal exceeds its byte limit",
+            )
+        parent = _open_private_parent(self.path.parent)
+        staged = f".{self.path.name}.{secrets.token_hex(16)}.partial"
+        descriptor: int | None = None
+        try:
+            try:
+                flags = (
+                    os.O_WRONLY
+                    | os.O_CREAT
+                    | os.O_EXCL
+                    | getattr(os, "O_NOFOLLOW", 0)
+                    | getattr(os, "O_CLOEXEC", 0)
+                    | getattr(os, "O_BINARY", 0)
+                )
+                descriptor = os.open(staged, flags, 0o600, dir_fd=parent)
+                _validate_private_file(os.fstat(descriptor))
+                offset = 0
+                while offset < len(rendered):
+                    written = os.write(descriptor, rendered[offset:])
+                    if written <= 0:
+                        raise OSError("could not finish writing cancellation journal")
+                    offset += written
+                os.fsync(descriptor)
+            finally:
+                if descriptor is not None:
+                    os.close(descriptor)
+            os.replace(staged, self.path.name, src_dir_fd=parent, dst_dir_fd=parent)
+            os.fsync(parent)
+        except BaseException:
+            _remove_stage(parent, staged)
+            raise
+        finally:
+            os.close(parent)
+
+    def cancelled(self, receipt: Mapping[str, Any]) -> None:
+        self.state["state"] = "cancelled"
+        self.state["receipt"] = dict(receipt)
+        self.save()
+
+
+def declare_arguments(parser: argparse.ArgumentParser) -> None:
+    parser.add_argument(
+        "--workspace-id",
+        required=True,
+        help="the workspace_id printed by mr-data acquire for this public acquisition",
+    )
+    parser.add_argument(
+        "--crawler-session-id",
+        required=True,
+        help="the crawler_session_id printed by mr-data acquire for this public acquisition",
+    )
+    parser.add_argument(
+        "--journal",
+        metavar="FILE",
+        help=(
+            "override the protected Receipt journal for this address; retries reuse its "
+            "idempotency key and it stores no credential"
+        ),
+    )
+
+
+def acquire_cancel(
+    args: argparse.Namespace,
+    *,
+    client: StudioAcquisitionCancellationClient | None = None,
+) -> dict[str, Any]:
+    """Cancel an active acquisition, or return its locally confirmed terminal receipt."""
+
+    workspace_id = _coordinate(getattr(args, "workspace_id", None), "--workspace-id")
+    crawler_session_id = _coordinate(
+        getattr(args, "crawler_session_id", None), "--crawler-session-id"
+    )
+    supplied_journal = getattr(args, "journal", None)
+    journal = CancellationJournal(
+        Path(str(supplied_journal)).expanduser()
+        if supplied_journal
+        else default_cancellation_journal_path(workspace_id, crawler_session_id)
+    )
+    # A bad coordinate or journal path is refused before Cloud is contacted; creating the durable
+    # record itself waits until the effective hosted credential exists, so an unauthenticated typo
+    # does not leave a file behind.
+    _journal_location(journal.path, create_parent=supplied_journal is None)
+    selected = client or StudioAcquisitionCancellationClient(_session(args))
+    journal.load_or_start(workspace_id=workspace_id, crawler_session_id=crawler_session_id)
+    receipt = journal.state.get("receipt")
+    if journal.state.get("state") == "cancelled":
+        if not isinstance(receipt, Mapping):
+            raise ThinLaneError(
+                "THIN_CANCELLATION_JOURNAL_INVALID",
+                "the cancellation journal says cancelled but holds no receipt",
+            )
+        _validate_receipt(receipt, workspace_id=workspace_id, crawler_session_id=crawler_session_id)
+        return dict(receipt)
+    returned, _replayed = selected.cancel(
+        workspace_id=workspace_id,
+        crawler_session_id=crawler_session_id,
+        idempotency_key=str(journal.state["idempotency_key"]),
+    )
+    _validate_receipt(returned, workspace_id=workspace_id, crawler_session_id=crawler_session_id)
+    journal.cancelled(returned)
+    return dict(returned)
+
+
+def _session(args: argparse.Namespace) -> StudioSession:
+    return open_studio_session(resolve_cloud_credentials())
+
+
+def _coordinate(value: Any, flag: str) -> str:
+    if not isinstance(value, str) or _UUID.fullmatch(value) is None:
+        raise ThinLaneError(
+            "THIN_CANCELLATION_COORDINATE_INVALID",
+            f"{flag} must be one lowercase RFC 4122 UUID",
+        )
+    try:
+        return str(UUID(value))
+    except ValueError as error:  # pragma: no cover - the expression excludes malformed UUIDs
+        raise ThinLaneError(
+            "THIN_CANCELLATION_COORDINATE_INVALID", f"{flag} is not a UUID"
+        ) from error
+
+
+def _idempotency_key(workspace_id: str, crawler_session_id: str) -> str:
+    return "mr-data.acquire-cancel.v1:" + canonical_sha256(
+        {
+            "schema_version": "hosted-public-acquisition-cancellation-idempotency.v1",
+            "workspace_id": workspace_id,
+            "crawler_session_id": crawler_session_id,
+        }
+    )
+
+
+def _retry_after(headers: Mapping[str, str]) -> int | None:
+    value = headers.get("retry-after", "")
+    if not value.isdecimal():
+        return None
+    seconds = int(value)
+    return seconds if 1 <= seconds <= MAX_RETRY_AFTER_SECONDS else None
+
+
+def _validate_receipt(
+    receipt: Mapping[str, Any], *, workspace_id: str, crawler_session_id: str
+) -> None:
+    if set(receipt) != RECEIPT_FIELDS:
+        raise ThinLaneError(
+            "THIN_CANCELLATION_RECEIPT_INVALID",
+            "Studio's cancellation receipt has missing, malformed, or extra fields",
+        )
+    cancelled_at = receipt.get("cancelled_at")
+    if (
+        receipt.get("schema_version") != CANCELLATION_RECEIPT_VERSION
+        or receipt.get("workspace_id") != workspace_id
+        or receipt.get("crawler_session_id") != crawler_session_id
+        or receipt.get("status") != "cancelled"
+        or not isinstance(receipt.get("requester_principal_id"), str)
+        or _UUID.fullmatch(str(receipt.get("requester_principal_id"))) is None
+        or not isinstance(cancelled_at, str)
+        or _SECONDS_Z.fullmatch(cancelled_at) is None
+    ):
+        raise ThinLaneError(
+            "THIN_CANCELLATION_RECEIPT_INVALID",
+            "Studio's cancellation receipt is not for this coordinate or is not terminally "
+            "cancelled",
+        )
+    try:
+        datetime.strptime(cancelled_at, "%Y-%m-%dT%H:%M:%SZ")
+    except ValueError as error:
+        raise ThinLaneError(
+            "THIN_CANCELLATION_RECEIPT_INVALID",
+            "Studio's cancellation receipt has no RFC3339 seconds timestamp",
+        ) from error
+
+
+def default_cancellation_journal_path(workspace_id: str, crawler_session_id: str) -> Path:
+    """The private, deterministic receipt location for one acquisition coordinate."""
+
+    return (
+        default_credentials_dir()
+        / _JOURNAL_STATE_DIRECTORY
+        / (f"{workspace_id}-{crawler_session_id}.json")
+    )
+
+
+def _journal_location(path: Path, *, create_parent: bool = False) -> None:
+    _require_secure_journal_primitives()
+    if path.name in {"", ".", ".."}:
+        raise ThinLaneError(
+            "THIN_CANCELLATION_JOURNAL_INVALID",
+            "--journal must name one ordinary file",
+        )
+    parent = path.parent
+    if create_parent:
+        _create_private_state_parent(parent)
+    try:
+        descriptor = _open_private_parent(parent)
+    except OSError as error:
+        raise ThinLaneError(
+            "THIN_CANCELLATION_JOURNAL_INVALID",
+            "--journal's parent must be one private directory owned by this user",
+        ) from error
+    else:
+        os.close(descriptor)
+
+
+def _create_private_state_parent(parent: Path) -> None:
+    """Create only our two state directories, then verify them through no-follow opens."""
+
+    root = default_credentials_dir()
+    try:
+        root.mkdir(mode=0o700, parents=True, exist_ok=True)
+        _assert_private_directory(root)
+        current = root
+        for part in Path(_JOURNAL_STATE_DIRECTORY).parts:
+            current = current / part
+            current.mkdir(mode=0o700, exist_ok=True)
+            _assert_private_directory(current)
+    except OSError as error:
+        raise ThinLaneError(
+            "THIN_CANCELLATION_JOURNAL_INVALID",
+            "the default cancellation journal directory could not be created securely",
+        ) from error
+    if parent != current:
+        raise ThinLaneError(
+            "THIN_CANCELLATION_JOURNAL_INVALID", "the default cancellation journal path is invalid"
+        )
+
+
+def _assert_private_directory(path: Path) -> None:
+    descriptor = _open_private_parent(path)
+    os.close(descriptor)
+
+
+def _open_private_parent(path: Path) -> int:
+    _require_secure_journal_primitives()
+    flags = (
+        os.O_RDONLY
+        | getattr(os, "O_DIRECTORY", 0)
+        | getattr(os, "O_NOFOLLOW", 0)
+        | getattr(os, "O_CLOEXEC", 0)
+    )
+    descriptor = os.open(path, flags)
+    try:
+        _validate_private_directory(os.fstat(descriptor))
+    except BaseException:
+        os.close(descriptor)
+        raise
+    return descriptor
+
+
+def _require_secure_journal_primitives() -> None:
+    """Refuse platforms without the descriptor operations that make journal writes safe.
+
+    Windows does not expose the POSIX ``dir_fd``/``O_NOFOLLOW`` contract needed to protect this
+    state from reparse-point replacement.  A path-based fallback would make a cancellation retry
+    less safe than the first cancellation, so portability here means an explicit, pre-network
+    refusal until an equally strong native primitive is available.
+    """
+
+    supported = getattr(os, "supports_dir_fd", frozenset())
+    if (
+        os.name != "posix"
+        or not all(operation in supported for operation in (os.open, os.rename, os.unlink))
+        or not hasattr(os, "O_NOFOLLOW")
+        or not hasattr(os, "O_DIRECTORY")
+    ):
+        raise ThinLaneError(
+            "THIN_CANCELLATION_JOURNAL_UNSUPPORTED",
+            "this platform lacks the secure filesystem primitives required for cancellation "
+            "receipt journaling",
+        )
+
+
+def _validate_private_directory(info: os.stat_result) -> None:
+    if not stat.S_ISDIR(info.st_mode):
+        raise OSError("journal parent is not a directory")
+    if os.name == "posix" and (info.st_uid != os.getuid() or stat.S_IMODE(info.st_mode) & 0o022):
+        raise OSError("journal parent is not private to this user")
+
+
+def _validate_private_file(info: os.stat_result) -> None:
+    if not stat.S_ISREG(info.st_mode) or info.st_nlink != 1:
+        raise OSError("journal is not a single-link regular file")
+    if os.name == "posix" and (info.st_uid != os.getuid() or stat.S_IMODE(info.st_mode) != 0o600):
+        raise OSError("journal is not a private file")
+
+
+def _remove_stage(parent: int, staged: str) -> None:
+    """Best-effort cleanup of our unguessable staging name; unlink never follows links."""
+
+    try:
+        os.unlink(staged, dir_fd=parent)
+    except FileNotFoundError:
+        pass
+
+
+def _read_journal(path: Path) -> dict[str, Any]:
+    parent: int | None = None
+    descriptor: int | None = None
+    try:
+        parent = _open_private_parent(path.parent)
+        flags = (
+            os.O_RDONLY
+            | os.O_NONBLOCK
+            | getattr(os, "O_NOFOLLOW", 0)
+            | getattr(os, "O_CLOEXEC", 0)
+            | getattr(os, "O_BINARY", 0)
+        )
+        descriptor = os.open(path.name, flags, dir_fd=parent)
+        info = os.fstat(descriptor)
+        _validate_private_file(info)
+        if info.st_size > MAX_JOURNAL_BYTES:
+            raise OSError("journal exceeds its byte limit")
+        pieces: list[bytes] = []
+        remaining = MAX_JOURNAL_BYTES + 1
+        while remaining:
+            piece = os.read(descriptor, min(remaining, 1 << 16))
+            if not piece:
+                break
+            pieces.append(piece)
+            remaining -= len(piece)
+        raw = b"".join(pieces)
+        if len(raw) > MAX_JOURNAL_BYTES:
+            raise OSError("journal exceeds its byte limit")
+        value = parse_json(raw)
+    except (OSError, CanonicalJSONError) as error:
+        raise ThinLaneError(
+            "THIN_CANCELLATION_JOURNAL_INVALID",
+            "the cancellation journal could not be read as one canonical journal this client wrote",
+        ) from error
+    finally:
+        if descriptor is not None:
+            os.close(descriptor)
+        if parent is not None:
+            os.close(parent)
+    if not isinstance(value, dict) or set(value) != JOURNAL_FIELDS:
+        raise ThinLaneError(
+            "THIN_CANCELLATION_JOURNAL_INVALID",
+            "the cancellation journal has missing, malformed, or extra fields",
+        )
+    if (
+        value.get("schema_version") != CANCELLATION_JOURNAL_VERSION
+        or value.get("state") not in {"prepared", "cancelled"}
+        or not isinstance(value.get("workspace_id"), str)
+        or not isinstance(value.get("crawler_session_id"), str)
+        or not isinstance(value.get("idempotency_key"), str)
+        or _IDEMPOTENCY.fullmatch(str(value.get("idempotency_key"))) is None
+    ):
+        raise ThinLaneError(
+            "THIN_CANCELLATION_JOURNAL_INVALID", "the cancellation journal is invalid"
+        )
+    if value["state"] == "prepared" and value.get("receipt") is not None:
+        raise ThinLaneError(
+            "THIN_CANCELLATION_JOURNAL_INVALID",
+            "the prepared cancellation journal already holds a receipt",
+        )
+    if value["state"] == "cancelled" and not isinstance(value.get("receipt"), dict):
+        raise ThinLaneError(
+            "THIN_CANCELLATION_JOURNAL_INVALID",
+            "the cancelled cancellation journal holds no receipt",
+        )
+    return value
+
+
+__all__ = [
+    "CANCELLATION_JOURNAL_VERSION",
+    "CANCELLATION_RECEIPT_VERSION",
+    "CANCEL_PUBLIC_ACQUISITION_PATH",
+    "JOURNAL_FIELDS",
+    "CancellationJournal",
+    "CancellationRefusal",
+    "StudioAcquisitionCancellationClient",
+    "acquire_cancel",
+    "declare_arguments",
+    "default_cancellation_journal_path",
+]