mostlyright-data 0.9.0__py3-none-any.whl

mostlyright/data_harness/thin/runs.py

@@ -0,0 +1,519 @@
+"""The Studio V3 run calls the thin lane makes: start one, read one, and fetch what it produced.
+
+Written against the contract rather than against the generated client, and that is a decision
+rather than an omission. The generated client is an optional extra (``mostlyright-studio-client-v3``
+arrives with ``[hosted]``), it pulls ``httpx`` and its whole closure, and its own
+``stream_run_events`` buffers the entire Server-Sent-Events body into one ``str`` before returning
+it -- unusable for a live stream by construction. A thin client that must run with no third-party
+package at all therefore speaks the wire directly, and :data:`PINNED_CONTRACT_VERSION` plus the
+route table below are what a reviewer diffs against ``contracts/openapi/studio-v3.yaml``.
+
+Every route here is a READ or a start. Nothing in this module cancels, retries, approves or
+releases: those are authority-bearing mutations whose refusals belong with the operator surfaces
+that already own them. The three lanes that DO write --
+:mod:`mostlyright.data_harness.thin.acquire`, which opens a hosted acquisition session,
+:mod:`mostlyright.data_harness.thin.approvals`, which records an ask a human must still settle,
+and :mod:`mostlyright.data_harness.thin.narrative`, which appends one unsealed cell of prose --
+keep their own routes in their own modules and share only :class:`StudioApiClient`, so this
+module's promise stays checkable by reading it.
+"""
+
+from __future__ import annotations
+
+import re
+import secrets
+from collections.abc import Mapping
+from typing import Any
+from urllib.parse import urlencode
+from uuid import UUID
+
+from mostlyright.data_harness.canonical import (
+    CanonicalJSONError,
+    canonical_json_bytes,
+    parse_foreign_json,
+)
+from mostlyright.data_harness.thin.session import StudioSession
+from mostlyright.data_harness.thin.transport import (
+    MAX_JSON_RESPONSE_BYTES,
+    ThinLaneError,
+    Transport,
+    UrlLibTransport,
+    signed_transfer_url,
+)
+
+#: The contract version every payload below is written against. ``studio_boundary`` pins the
+#: generated client to the same 3.0.0 and to one OpenAPI source digest; this constant is the thin
+#: lane's half of that pin, and ``tests/test_thin_client.py`` holds the two equal.
+PINNED_CONTRACT_VERSION = "3.0.0"
+
+#: The routes. Kept as one table because it is the thing a reviewer checks against the contract.
+START_RUN_PATH = "/v3/runs"
+#: The same collection, read rather than added to. Written out rather than aliased to
+#: :data:`START_RUN_PATH`, because the cross-repo gate checks each name against its own generated
+#: operation and an alias would check one of them twice.
+LIST_RUNS_PATH = "/v3/runs"
+GET_RUN_PATH = "/v3/runs/{run_id}"
+RUN_EVENTS_STREAM_PATH = "/v3/runs/{run_id}/events:stream"
+RUN_EXECUTION_SUMMARY_PATH = "/v3/runs/{run_id}/execution-summary"
+RUN_CANDIDATE_EVIDENCE_PATH = "/v3/runs/{run_id}/candidate-evidence"
+RUN_AGENT_EVIDENCE_PATH = "/v3/runs/{run_id}/agent-evidence"
+RUN_FINDINGS_PATH = "/v3/runs/{run_id}/findings"
+#: Studio renders this itself, from the run's event log and from nothing else, once the run is
+#: terminal. It is the whole reason the vendored client moved to 46e64f1; faa7dbb after it is what
+#: made `acquire`, `approve` and `recipe-approve` writable, and fd14830 after that is what made
+#: `note` able to put a sentence into the log this document is rendered from.
+RUN_NOTEBOOK_PATH = "/v3/runs/{run_id}/notebook"
+TABLE_VERSIONS_PATH = "/v3/tables/{table_id}/versions"
+TABLE_VERSION_PATH = "/v3/tables/{table_id}/versions/{table_version_id}"
+TABLE_VERSION_PREVIEW_PATH = "/v3/tables/{table_id}/versions/{table_version_id}/preview"
+TABLE_VERSION_DOWNLOAD_PATH = "/v3/tables/{table_id}/versions/{table_version_id}/download"
+ARTIFACT_DOWNLOAD_SESSION_PATH = "/v3/artifacts/{artifact_id}/download-session"
+#: Two reads every lane that touches a recipe proposal shares. The proposal itself, for its status
+#: and the ETag its approval is pinned to; and the Build an approval STARTED, which Studio #124
+#: made possible: the dashboard's confirmation can start the first hosted Build in the same
+#: transaction, after which the proposal is `activation_started`, `POST /v3/runs` refuses it, and
+#: this projection is the only way to recover the run that already exists. It answers 404 while
+#: the approval is undecided or unstarted, and that 404 is an ordinary answer here.
+GET_RECIPE_PROPOSAL_PATH = "/v3/recipe-proposals/{recipe_proposal_id}"
+TABLE_RECIPE_APPROVAL_START_PATH = "/v3/table-recipe-approvals/{approval_request_id}/start"
+
+#: The Run statuses this contract admits (``run.schema.json#/properties/status``). Held equal to
+#: the schema's own enum by a cross-repo gate, so a filter can never name a status no Run holds.
+RUN_STATUSES = (
+    "queued",
+    "producer_retry_pending",
+    "running",
+    "awaiting_candidate_selection",
+    "verifying",
+    "verifier_retry_pending",
+    "releasable",
+    "released",
+    "repair_required",
+    "failed",
+    "cancelled",
+    "abandoned",
+    "human_direction_required",
+)
+
+#: Studio's own ceiling on one page of any V3 list (``components.parameters.Limit``). Asking for
+#: more is a 422, so it is bounded here and the refusal is a sentence rather than a round trip.
+MAX_LIST_LIMIT = 200
+
+#: ``kind`` values a start-run request may name. The Run resource itself admits more (a successor
+#: or a repair successor is created by Studio, never asked for from out here).
+START_RUN_KINDS = ("initial_build", "incremental_refresh", "bounded_backfill")
+
+#: The closed capacity vocabulary a start-run may name
+#: (``common.schema.json#/$defs/resource_class``).
+RESOURCE_CLASSES = ("standard", "large")
+
+#: ``expected_recipe_digest`` is the one digest in this contract carried WITHOUT the ``sha256:``
+#: prefix. Refusing a prefixed value here turns a silent 422 from Studio into a local sentence.
+_RECIPE_DIGEST = re.compile(r"^[0-9a-f]{64}$")
+_PREFIXED_DIGEST = re.compile(r"^sha256:[0-9a-f]{64}$")
+
+#: ``common.schema.json#/$defs/idempotency_key``.
+_IDEMPOTENCY_KEY = re.compile(r"^[A-Za-z0-9._:-]{16,128}$")
+
+
+def new_idempotency_key(prefix: str = "mr-data") -> str:
+    """A fresh idempotency key for one submission.
+
+    Random rather than derived from the request body: two deliberate submissions of the same
+    recipe are two runs, and a content-derived key would silently return the first one's Run and
+    look like a success.
+    """
+
+    key = f"{prefix}-{secrets.token_hex(16)}"
+    if _IDEMPOTENCY_KEY.fullmatch(key) is None:  # pragma: no cover - the format is fixed above
+        raise ThinLaneError("THIN_REQUEST_INVALID", "the generated idempotency key is malformed")
+    return key
+
+
+def start_run_body(
+    *,
+    workspace_id: UUID,
+    expected_recipe_digest: str,
+    recipe_proposal_id: UUID,
+    resource_class: str | None = None,
+) -> dict[str, Any]:
+    """The ``initial_build`` start-run body, with the contract's conditional applied here.
+
+    ⚠ ``initial_build`` is exclusive. The contract requires ``recipe_proposal_id`` and FORBIDS
+    ``dataset_id``, ``table_id`` and ``requested_range``; the other two kinds invert that exactly.
+    This builder covers the initial build only, so the forbidden trio can never be added by
+    accident, and a caller that needs a refresh or a backfill gets a builder of its own rather
+    than a flag on this one.
+    """
+
+    if _PREFIXED_DIGEST.fullmatch(expected_recipe_digest):
+        raise ThinLaneError(
+            "THIN_REQUEST_INVALID",
+            "expected_recipe_digest is bare lowercase hex in this contract, without a sha256: "
+            "prefix",
+        )
+    if _RECIPE_DIGEST.fullmatch(expected_recipe_digest) is None:
+        raise ThinLaneError(
+            "THIN_REQUEST_INVALID", "expected_recipe_digest is not a 64-character hex digest"
+        )
+    if resource_class is not None and resource_class not in RESOURCE_CLASSES:
+        allowed = ", ".join(RESOURCE_CLASSES)
+        raise ThinLaneError("THIN_REQUEST_INVALID", f"resource_class must be one of: {allowed}")
+    body: dict[str, Any] = {
+        "workspace_id": str(workspace_id),
+        "expected_recipe_digest": expected_recipe_digest,
+        "kind": "initial_build",
+        "recipe_proposal_id": str(recipe_proposal_id),
+    }
+    if resource_class is not None:
+        body["resource_class"] = resource_class
+    return body
+
+
+class StudioApiClient:
+    """One bearer-authenticated V3 call, and the refusal it becomes when it does not answer.
+
+    Extracted from :class:`StudioRunClient` rather than copied, because the two lanes that came
+    after it -- hosted acquisition and the approval ask -- speak the same wire under the same
+    hardening, and a second ``_call`` would be a second place for the byte ceiling, the JSON
+    hardening and the Studio-code translation to drift apart. Subclasses add routes; none of them
+    adds a transport.
+    """
+
+    def __init__(self, session: StudioSession, *, transport: Transport | None = None) -> None:
+        self._session = session
+        self._transport = transport or UrlLibTransport()
+
+    @property
+    def session(self) -> StudioSession:
+        return self._session
+
+    @property
+    def transport(self) -> Transport:
+        """The bounded transport this client speaks through.
+
+        Exposed so a signed download made ON BEHALF of a call this client made goes out over the
+        same substituted seam -- otherwise a test that fakes Studio would still reach the network
+        for the bytes Studio pointed it at.
+        """
+
+        return self._transport
+
+    def _call(
+        self,
+        method: str,
+        path: str,
+        *,
+        body: Mapping[str, Any] | None = None,
+        extra_headers: Mapping[str, str] | None = None,
+        query: Mapping[str, Any] | None = None,
+        expected: tuple[int, ...],
+        shape: type = dict,
+        response_headers: dict[str, str] | None = None,
+    ) -> Any:
+        headers = {"Accept": "application/json", **self._session.authorization()}
+        if extra_headers:
+            headers.update(extra_headers)
+        raw_body: bytes | None = None
+        if body is not None:
+            headers["Content-Type"] = "application/json"
+            raw_body = canonical_json_bytes(body)
+        url = f"{self._session.studio_base_url}{path}{_query_string(query)}"
+        try:
+            status, raw, answered_headers = self._transport.request(
+                method, url, headers, raw_body, MAX_JSON_RESPONSE_BYTES
+            )
+            if response_headers is not None:
+                response_headers.update(answered_headers)
+        except ThinLaneError:
+            raise
+        except Exception as error:
+            raise ThinLaneError(
+                "THIN_STUDIO_UNREACHABLE", f"Studio did not answer {path}"
+            ) from error
+        try:
+            # ⚠ `parse_foreign_json`, NOT `parse_json`. Every hardening stays -- duplicate keys,
+            # byte-order marks, non-strict UTF-8, NaN, out-of-range integers, depth and member
+            # ceilings are all refused identically -- and the single difference is that a finite
+            # fractional number is admissible. It has to be: a table preview's rows carry the
+            # table's own values, and a temperature column is not integers. `parse_json` refused
+            # the whole response for one `20.5`, which turned a perfectly good preview into
+            # "Studio did not answer with an object". Nothing read here can be sealed either way:
+            # `canonical_json_bytes` still refuses a float, so a value parsed here can be shown
+            # and can never reach a digest.
+            parsed = parse_foreign_json(raw)
+        except CanonicalJSONError:
+            parsed = None
+        if status not in expected:
+            raise _studio_refusal(status, parsed, path)
+        if not isinstance(parsed, shape):
+            wanted = "a list" if shape is list else "an object"
+            raise ThinLaneError("THIN_RESPONSE_INVALID", f"{path} did not answer with {wanted}")
+        return parsed
+
+    def _call_list(
+        self, path: str, *, query: Mapping[str, Any] | None = None
+    ) -> list[dict[str, Any]]:
+        """One V3 list route, which answers with a JSON array rather than with an object.
+
+        Separate from :meth:`_call` rather than folded into it with a flag: every list route in
+        this contract answers with a bare array and every single-resource route answers with an
+        object, so a reader of either method knows which shape it is looking at without checking a
+        parameter. Non-object members are dropped rather than refused -- a list is a page of
+        resources, and one malformed member is not a reason to refuse the page -- but the array
+        itself being absent is.
+        """
+
+        parsed = self._call("GET", path, query=query, expected=(200,), shape=list)
+        return [item for item in parsed if isinstance(item, dict)]
+
+    def recipe_proposal(
+        self, recipe_proposal_id: str, *, response_headers: dict[str, str] | None = None
+    ) -> dict[str, Any]:
+        """One recipe proposal as Studio holds it, with the ETag it answered under."""
+
+        return self._call(
+            "GET",
+            GET_RECIPE_PROPOSAL_PATH.format(recipe_proposal_id=recipe_proposal_id),
+            expected=(200,),
+            response_headers=response_headers,
+        )
+
+    def approval_start(self, approval_request_id: str) -> dict[str, Any] | None:
+        """The Build a Table recipe approval started, or ``None`` while it has started none.
+
+        404 is the ordinary answer for an approval that is undecided, unstarted, or local-build,
+        and it is returned as ``None`` rather than raised: every caller asks this question to
+        decide between "print the build command" and "print the run that already exists".
+        """
+
+        try:
+            return self._call(
+                "GET",
+                TABLE_RECIPE_APPROVAL_START_PATH.format(approval_request_id=approval_request_id),
+                expected=(200,),
+            )
+        except ThinLaneError as refusal:
+            if refusal.code == "THIN_NOT_FOUND":
+                return None
+            raise
+
+
+class StudioRunClient(StudioApiClient):
+    """Bearer-authenticated V3 reads and one start, over the thin transport."""
+
+    def start_initial_build(
+        self,
+        *,
+        expected_recipe_digest: str,
+        recipe_proposal_id: UUID,
+        resource_class: str | None = None,
+        idempotency_key: str | None = None,
+    ) -> dict[str, Any]:
+        """Queue one hosted initial build and return Studio's Run."""
+
+        return self._call(
+            "POST",
+            START_RUN_PATH,
+            body=start_run_body(
+                workspace_id=self._session.workspace_id,
+                expected_recipe_digest=expected_recipe_digest,
+                recipe_proposal_id=recipe_proposal_id,
+                resource_class=resource_class,
+            ),
+            extra_headers={"Idempotency-Key": idempotency_key or new_idempotency_key()},
+            expected=(202,),
+        )
+
+    def get_run(self, run_id: str) -> dict[str, Any]:
+        return self._call("GET", GET_RUN_PATH.format(run_id=run_id), expected=(200,))
+
+    def list_runs(
+        self,
+        *,
+        status: str | None = None,
+        limit: int | None = None,
+        dataset_id: str | None = None,
+        question_id: str | None = None,
+        created_after: str | None = None,
+        cursor: str | None = None,
+    ) -> list[dict[str, Any]]:
+        """One page of this workspace's runs, newest filters applied by Studio rather than here.
+
+        The workspace is the bearer token's, never a parameter: this contract scopes every list to
+        the authority that asked, and passing a workspace would be asking for one we might not
+        hold.
+        """
+
+        if status is not None and status not in RUN_STATUSES:
+            raise ThinLaneError(
+                "THIN_REQUEST_INVALID",
+                f"{status!r} is not a run state; this contract admits " + ", ".join(RUN_STATUSES),
+            )
+        if limit is not None and not 1 <= limit <= MAX_LIST_LIMIT:
+            raise ThinLaneError(
+                "THIN_REQUEST_INVALID", f"how many runs to list must be from 1 to {MAX_LIST_LIMIT}"
+            )
+        return self._call_list(
+            LIST_RUNS_PATH,
+            query={
+                "status": status,
+                "limit": limit,
+                "dataset_id": dataset_id,
+                "question_id": question_id,
+                "created_after": created_after,
+                "cursor": cursor,
+            },
+        )
+
+    def run_findings(self, run_id: str) -> list[dict[str, Any]]:
+        return self._call_list(RUN_FINDINGS_PATH.format(run_id=run_id))
+
+    def run_agent_evidence(self, run_id: str) -> list[dict[str, Any]]:
+        return self._call_list(RUN_AGENT_EVIDENCE_PATH.format(run_id=run_id))
+
+    def run_notebook(self, run_id: str) -> dict[str, Any]:
+        """The notebook Studio rendered for one terminal run, and a signed download per document.
+
+        404 here is an ordinary answer rather than a fault: a run that has not finished has no
+        notebook, because the live view is the stream. The caller turns that into the sentence.
+        """
+
+        return self._call("GET", RUN_NOTEBOOK_PATH.format(run_id=run_id), expected=(200,))
+
+    def table_versions(self, table_id: str) -> list[dict[str, Any]]:
+        return self._call_list(TABLE_VERSIONS_PATH.format(table_id=table_id))
+
+    def table_version(self, table_id: str, table_version_id: str) -> dict[str, Any]:
+        return self._call(
+            "GET",
+            TABLE_VERSION_PATH.format(table_id=table_id, table_version_id=table_version_id),
+            expected=(200,),
+        )
+
+    def table_version_preview(self, table_id: str, table_version_id: str) -> dict[str, Any]:
+        """The bounded preview Studio publishes beside one released version.
+
+        Bounded by Studio, not by this client: the preview is an artifact the release wrote, so a
+        caller can narrow what it shows and can never widen it.
+        """
+
+        return self._call(
+            "GET",
+            TABLE_VERSION_PREVIEW_PATH.format(table_id=table_id, table_version_id=table_version_id),
+            expected=(200,),
+        )
+
+    def execution_summary(self, run_id: str) -> dict[str, Any]:
+        return self._call("GET", RUN_EXECUTION_SUMMARY_PATH.format(run_id=run_id), expected=(200,))
+
+    def candidate_evidence(self, run_id: str) -> dict[str, Any]:
+        return self._call("GET", RUN_CANDIDATE_EVIDENCE_PATH.format(run_id=run_id), expected=(200,))
+
+    def table_version_download(self, table_id: str, table_version_id: str) -> dict[str, Any]:
+        return self._call(
+            "GET",
+            TABLE_VERSION_DOWNLOAD_PATH.format(
+                table_id=table_id, table_version_id=table_version_id
+            ),
+            expected=(200,),
+        )
+
+    def artifact_download_session(self, artifact_id: str) -> dict[str, Any]:
+        """Ask Studio to sign one short-lived download of one artifact.
+
+        The session is created, never chosen: kind, purpose and classification come from sealed
+        server metadata, so this call carries no body at all beyond its idempotency key.
+        """
+
+        session = self._call(
+            "POST",
+            ARTIFACT_DOWNLOAD_SESSION_PATH.format(artifact_id=artifact_id),
+            body={},
+            extra_headers={"Idempotency-Key": new_idempotency_key("mr-data-dl")},
+            expected=(201,),
+        )
+        if session.get("direction") != "download" or session.get("method") != "GET":
+            raise ThinLaneError(
+                "THIN_RESPONSE_INVALID", "Studio signed a session that is not a download"
+            )
+        signed_transfer_url(session.get("signed_url", ""), "the signed download URL")
+        return session
+
+
+def _query_string(query: Mapping[str, Any] | None) -> str:
+    """Render the filters one list route was given, and nothing for the ones it was not.
+
+    ``None`` means "not asked for" and is dropped rather than sent empty: this contract's list
+    parameters are all optional, and `?status=` is a different request from no `status` at all.
+    Sorted, so two identical requests produce one identical URL -- which is what makes a recorded
+    request in a test an assertion about the request rather than about dict ordering.
+    """
+
+    if not query:
+        return ""
+    selected = {key: value for key, value in sorted(query.items()) if value is not None}
+    if not selected:
+        return ""
+    return "?" + urlencode({key: str(value) for key, value in selected.items()})
+
+
+def _studio_refusal(status: int, parsed: Any, path: str) -> ThinLaneError:
+    """Turn one ``application/problem+json`` body into a typed refusal.
+
+    Studio's error code is a closed vocabulary and is far more useful than the status, so it is
+    carried through verbatim as the refusal detail's first token. The message is Studio's own and
+    is bounded by its contract at 2000 characters; it is not reformatted here.
+    """
+
+    studio_code = parsed.get("code") if isinstance(parsed, dict) else None
+    if not isinstance(studio_code, str):
+        studio_code = None
+    if status in {401, 403}:
+        refusal = ThinLaneError(
+            "THIN_AUTHORIZATION_DENIED",
+            "the Studio token does not carry the authority this call needs",
+        )
+    elif status == 404:
+        refusal = ThinLaneError("THIN_NOT_FOUND", f"Studio has no such resource for {path}")
+    elif studio_code is not None:
+        message = parsed.get("message") if isinstance(parsed, dict) else None
+        detail = message if isinstance(message, str) and message else "Studio refused the call"
+        refusal = ThinLaneError(f"THIN_STUDIO_{studio_code}", detail)
+    else:
+        refusal = ThinLaneError("THIN_STUDIO_REFUSED", f"Studio answered {path} with HTTP {status}")
+    # The status and Studio's own code ride on the refusal as facts about the answer, beside the
+    # translated code a script branches on: a write lane deciding whether a refused request was
+    # applied needs to tell a validation 422 from a gateway 504 carrying the same translated code.
+    refusal.http_status = status  # type: ignore[attr-defined]
+    refusal.studio_code = studio_code  # type: ignore[attr-defined]
+    return refusal
+
+
+__all__ = [
+    "ARTIFACT_DOWNLOAD_SESSION_PATH",
+    "GET_RECIPE_PROPOSAL_PATH",
+    "GET_RUN_PATH",
+    "LIST_RUNS_PATH",
+    "MAX_LIST_LIMIT",
+    "PINNED_CONTRACT_VERSION",
+    "RESOURCE_CLASSES",
+    "RUN_AGENT_EVIDENCE_PATH",
+    "RUN_CANDIDATE_EVIDENCE_PATH",
+    "RUN_EVENTS_STREAM_PATH",
+    "RUN_EXECUTION_SUMMARY_PATH",
+    "RUN_FINDINGS_PATH",
+    "RUN_NOTEBOOK_PATH",
+    "RUN_STATUSES",
+    "START_RUN_KINDS",
+    "START_RUN_PATH",
+    "TABLE_RECIPE_APPROVAL_START_PATH",
+    "TABLE_VERSIONS_PATH",
+    "TABLE_VERSION_DOWNLOAD_PATH",
+    "TABLE_VERSION_PATH",
+    "TABLE_VERSION_PREVIEW_PATH",
+    "StudioApiClient",
+    "StudioRunClient",
+    "new_idempotency_key",
+    "start_run_body",
+]