mostlyright-data 0.9.0__py3-none-any.whl

mostlyright/data_harness/thin/commands.py

@@ -0,0 +1,406 @@
+"""The hosted lane behind the command names people already type.
+
+Each function here returns one payload mapping and prints nothing. That is the convention
+``ux.commands.CommandSpec`` already establishes for every auto-discovered subcommand, and it is
+what lets :mod:`mostlyright.data_harness.thin.router` decide plain lines or ``--json`` in one
+place for the whole lane.
+
+WHAT THE PAYLOADS PROMISE. Every fact in them comes out of one Studio response or out of the
+submission this command just made. Nothing is derived from a clock, no percentage is manufactured,
+and no progress record is presented as evidence -- the unsealed stream says what it measured and
+the sealed candidate evidence says what was proved, and the two are reported under separate keys
+so a reader cannot mistake one for the other.
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from collections.abc import Mapping
+from pathlib import Path
+from typing import Any
+from uuid import UUID
+
+from mostlyright.data_harness.formats import FORMAT_MEDIA_TYPES, FORMAT_SUFFIXES
+from mostlyright.data_harness.thin import THIN_SCHEMA_PREFIX
+from mostlyright.data_harness.thin.download import download_signed_artifact
+from mostlyright.data_harness.thin.runs import RESOURCE_CLASSES, StudioRunClient
+from mostlyright.data_harness.thin.session import (
+    StudioSession,
+    open_studio_session,
+    run_dashboard_url,
+)
+from mostlyright.data_harness.thin.stream import TERMINAL_RUN_STATUSES, RunEventStream
+from mostlyright.data_harness.thin.transport import ThinLaneError
+from mostlyright.data_harness.ux.credentials import resolve_cloud_credentials
+
+SUBMIT_SCHEMA = f"{THIN_SCHEMA_PREFIX}-submission.v1"
+WATCH_SCHEMA = f"{THIN_SCHEMA_PREFIX}-watch.v1"
+STATUS_SCHEMA = f"{THIN_SCHEMA_PREFIX}-run-status.v1"
+DOWNLOAD_SCHEMA = f"{THIN_SCHEMA_PREFIX}-download.v1"
+
+#: The artifact kinds ``deploy-dataset`` brings back by default: the dataset itself, and nothing
+#: else. A run also produces evidence, previews, profiles and lineage; those are fetched only when
+#: a caller names them, because a folder that quietly filled with nine files is not a dataset.
+DEFAULT_DOWNLOAD_KINDS = ("table_package", "candidate_parquet")
+
+
+def _session(args: argparse.Namespace) -> StudioSession:
+    """Resolve this machine's credential and open a short-lived Studio session."""
+
+    credentials = resolve_cloud_credentials()
+    return open_studio_session(credentials)
+
+
+def _client(args: argparse.Namespace) -> StudioRunClient:
+    return StudioRunClient(_session(args))
+
+
+def _run_id(value: str) -> str:
+    """Refuse anything that is not a run identifier before it reaches a URL path."""
+
+    try:
+        return str(UUID(value))
+    except ValueError as error:
+        raise ThinLaneError("THIN_REQUEST_INVALID", f"{value!r} is not a run identifier") from error
+
+
+def run_receipt(run: Mapping[str, Any]) -> dict[str, Any]:
+    """The Run fields a person reads, exactly as Studio reported them."""
+
+    return {
+        key: run[key]
+        for key in (
+            "run_id",
+            "status",
+            "kind",
+            "resource_class",
+            "table_id",
+            "dataset_id",
+            "recipe_digest",
+            "recipe_version",
+            "created_at",
+            "started_at",
+            "completed_at",
+            "failure_code",
+        )
+        if key in run
+    }
+
+
+def submit(args: argparse.Namespace, *, client: StudioRunClient | None = None) -> dict[str, Any]:
+    """Queue one hosted initial build and say where to watch it.
+
+    Backs ``mr-data build`` and ``mr-data run`` in the hosted lane. The two names stay distinct
+    everywhere else in this product and mean the same thing here, because a hosted submission is
+    one act: there is no local workbench to build in and no local plan to run.
+    """
+
+    selected = client or _client(args)
+    try:
+        run = selected.start_initial_build(
+            expected_recipe_digest=args.recipe_digest,
+            recipe_proposal_id=UUID(args.recipe_proposal),
+            resource_class=getattr(args, "resource_class", None),
+        )
+    except ThinLaneError as refusal:
+        started = _already_started(selected, args, refusal)
+        if started is None:
+            raise
+        return started
+    run_id = run.get("run_id")
+    if not isinstance(run_id, str):
+        raise ThinLaneError("THIN_RESPONSE_INVALID", "Studio queued a run without an identifier")
+    return {
+        "schema_version": SUBMIT_SCHEMA,
+        "status": "run_queued",
+        "lane": "hosted",
+        "dashboard_url": run_dashboard_url(selected.session.cloud_url, run_id),
+        "watch_command": f"mr-data watch {run_id} --hosted",
+        "run": run_receipt(run),
+        "studio": selected.session.to_receipt(),
+    }
+
+
+def _already_started(
+    client: StudioRunClient, args: argparse.Namespace, refusal: ThinLaneError
+) -> dict[str, Any] | None:
+    """The run an approval already started, when that is why the submission was refused.
+
+    Since Studio #124 a person confirming a Table recipe in the dashboard can start its first
+    hosted Build in the same transaction. The proposal is then ``activation_started``, and
+    ``POST /v3/runs`` refuses it with ``RELEASE_GATE_BLOCKED`` -- before it reaches the activation
+    receipt that would otherwise replay the run. So a build command printed before the person
+    decided is not wrong, it is late: the run it asked for exists. This reads the proposal, and
+    if the approval started the Build, reads that Build back and reports it as already queued. Any
+    other refusal, or a proposal in any other state, is the refusal it was.
+    """
+
+    if refusal.code != "THIN_STUDIO_RELEASE_GATE_BLOCKED":
+        return None
+    proposal = client.recipe_proposal(args.recipe_proposal)
+    approval_request_id = proposal.get("approval_request_id")
+    if proposal.get("status") != "activation_started" or not isinstance(approval_request_id, str):
+        return None
+    if proposal.get("recipe_digest") != args.recipe_digest:
+        # The run that exists was started from a different recipe than the one asked for, which
+        # is the refusal Studio gave, not a run to hand back.
+        return None
+    started = client.approval_start(approval_request_id)
+    if started is None:
+        return None
+    run = started.get("run")
+    run_id = run.get("run_id") if isinstance(run, dict) else None
+    if not isinstance(run_id, str):
+        raise ThinLaneError("THIN_RESPONSE_INVALID", "Studio named a started Build without a run")
+    return {
+        "schema_version": SUBMIT_SCHEMA,
+        "status": "run_already_queued",
+        "lane": "hosted",
+        "approval_request_id": approval_request_id,
+        "dashboard_url": run_dashboard_url(client.session.cloud_url, run_id),
+        "watch_command": f"mr-data watch {run_id} --hosted",
+        "run": run_receipt(run),
+        "note": (
+            "The approval that confirmed this recipe started its first hosted Build in the same "
+            "act, so there was nothing left for this command to queue. This is that Build."
+        ),
+        "studio": client.session.to_receipt(),
+    }
+
+
+def watch(args: argparse.Namespace, *, client: StudioRunClient | None = None) -> dict[str, Any]:
+    """Stream one run's event log to the terminal, resuming across stream cuts.
+
+    Lines are written as they arrive rather than collected and printed at the end -- the point of
+    a live view is that it is live -- and the returned payload is the summary of what was seen, so
+    ``--json`` gets one object and a person gets a timeline followed by that object.
+    """
+
+    # The identifier is checked BEFORE the credential is resolved and a token is minted. A typo in
+    # a run id should cost a sentence, not a network round trip and a Studio 404.
+    run_id = _run_id(args.run_id)
+    selected = client or _client(args)
+    stream = RunEventStream(selected.session, run_id)
+    announce = getattr(args, "json", False) is False
+
+    def _write(line: str) -> None:
+        if announce:
+            print(line, flush=True)
+
+    if announce:
+        print(f"watching {run_id}", flush=True)
+        print(run_dashboard_url(selected.session.cloud_url, run_id), flush=True)
+    progress = stream.follow(
+        from_seq=int(getattr(args, "from_seq", 0) or 0),
+        last_event_id=getattr(args, "last_event_id", None),
+        on_line=_write,
+    )
+    run = selected.get_run(run_id)
+    return {
+        "schema_version": WATCH_SCHEMA,
+        "status": "run_terminal" if run.get("status") in TERMINAL_RUN_STATUSES else "watch_ended",
+        "lane": "hosted",
+        "dashboard_url": run_dashboard_url(selected.session.cloud_url, run_id),
+        "run": run_receipt(run),
+        "stream": progress.to_receipt(),
+    }
+
+
+def status(args: argparse.Namespace, *, client: StudioRunClient | None = None) -> dict[str, Any]:
+    """Report Studio's own reading of one run, with its receipts when it has any.
+
+    The execution summary and the candidate evidence are read only once the run has something to
+    say: asking for candidate evidence on a queued run is a 404 by design, and reporting that 404
+    as a failure would be reporting the ordinary state of every new run as an error.
+    """
+
+    # The identifier is checked BEFORE the credential is resolved and a token is minted. A typo in
+    # a run id should cost a sentence, not a network round trip and a Studio 404.
+    run_id = _run_id(args.run_id)
+    selected = client or _client(args)
+    run = selected.get_run(run_id)
+    payload: dict[str, Any] = {
+        "schema_version": STATUS_SCHEMA,
+        "status": "run_reported",
+        "lane": "hosted",
+        "dashboard_url": run_dashboard_url(selected.session.cloud_url, run_id),
+        "run": run_receipt(run),
+    }
+    if not getattr(args, "receipts", False):
+        return payload
+    unavailable: dict[str, str] = {}
+    payload["execution_summary"] = optional_receipt(
+        "execution_summary", lambda: selected.execution_summary(run_id), unavailable
+    )
+    evidence = optional_receipt(
+        "candidate_evidence", lambda: selected.candidate_evidence(run_id), unavailable
+    )
+    if isinstance(evidence, Mapping):
+        payload["candidate_evidence"] = {
+            key: evidence[key]
+            for key in ("execution_scope_digest", "release_policy_digest", "admission")
+            if key in evidence
+        }
+        payload["artifacts"] = [
+            {
+                key: artifact[key]
+                for key in (
+                    "artifact_id",
+                    "kind",
+                    "purpose",
+                    "media_type",
+                    "size_bytes",
+                    "content_digest",
+                    "classification",
+                )
+                if key in artifact
+            }
+            for artifact in artifact_records(evidence)
+        ]
+    else:
+        payload["candidate_evidence"] = None
+        payload["artifacts"] = []
+    payload["receipts_unavailable"] = unavailable
+    return payload
+
+
+def download(args: argparse.Namespace, *, client: StudioRunClient | None = None) -> dict[str, Any]:
+    """Bring one released run's dataset back to this machine, verified.
+
+    Backs ``mr-data deploy-dataset`` in the hosted lane, which is the command that already means
+    "bring the Build a deployment released back here". Under the thin profile it brings back the
+    bytes and their proof; it does not write a dataset notebook, because rendering one is
+    server-side work under ADR 0021, and a locally rendered notebook beside a hosted build is
+    exactly the kind of second answer this architecture exists to remove.
+    """
+
+    # The identifier is checked BEFORE the credential is resolved and a token is minted. A typo in
+    # a run id should cost a sentence, not a network round trip and a Studio 404.
+    run_id = _run_id(args.run_id)
+    selected = client or _client(args)
+    evidence = selected.candidate_evidence(run_id)
+    wanted = tuple(getattr(args, "kind", None) or DEFAULT_DOWNLOAD_KINDS)
+    output = Path(getattr(args, "output", None) or "hosted-dataset").expanduser()
+    selectedartifact_records = [
+        artifact for artifact in artifact_records(evidence) if artifact.get("kind") in wanted
+    ]
+    if not selectedartifact_records:
+        available = sorted({str(item.get("kind")) for item in artifact_records(evidence)})
+        raise ThinLaneError(
+            "THIN_NO_MATCHING_ARTIFACT",
+            f"this run published no artifact of kind {', '.join(wanted)}"
+            + (f"; it published {', '.join(available)}" if available else ""),
+        )
+    downloaded = []
+    for artifact in selectedartifact_records:
+        artifact_id = artifact.get("artifact_id")
+        if not isinstance(artifact_id, str):
+            raise ThinLaneError("THIN_RESPONSE_INVALID", "an artifact carried no identifier")
+        session = selected.artifact_download_session(artifact_id)
+        destination = output / f"{artifact_id}{suffix_for(artifact.get('media_type'))}"
+        downloaded.append(download_signed_artifact(session, destination).to_receipt())
+    return {
+        "schema_version": DOWNLOAD_SCHEMA,
+        "status": "artifacts_downloaded",
+        "lane": "hosted",
+        "dashboard_url": run_dashboard_url(selected.session.cloud_url, run_id),
+        "run_id": run_id,
+        "output": str(output),
+        "artifacts": downloaded,
+    }
+
+
+def artifact_records(evidence: Mapping[str, Any]) -> list[Mapping[str, Any]]:
+    artifacts = evidence.get("artifacts")
+    if not isinstance(artifacts, list):
+        return []
+    return [item for item in artifacts if isinstance(item, Mapping)]
+
+
+#: Media type to file extension, derived from the one format table rather than restated. Both
+#: halves come from :mod:`mostlyright.data_harness.formats`, so a format added there reaches a
+#: downloaded file's name without anybody remembering this line exists.
+SUFFIX_BY_MEDIA_TYPE = {
+    media_type: min(FORMAT_SUFFIXES[data_format])
+    for data_format, media_types in FORMAT_MEDIA_TYPES.items()
+    if data_format in FORMAT_SUFFIXES
+    for media_type in media_types
+}
+
+
+def media_type_essence(media_type: Any) -> str:
+    """The type/subtype half of a media type, with its parameters and its casing removed.
+
+    ⚠ A MEDIA TYPE IS NOT A STRING TO COMPARE. RFC 9110 §8.3 makes everything after the first
+    ``;`` a parameter list, and Studio sends ``text/html; charset=utf-8`` where this repository
+    had written ``text/html``. An equality test against the whole header therefore missed, the
+    rendered notebook was named ``.bin``, and nothing failed loudly -- which is exactly the shape
+    of defect a comparison against a parsed value cannot have. Type and subtype are
+    case-insensitive in that same section, so they are lowered here rather than at each caller.
+
+    Only the essence is returned. No caller in this lane acts on a parameter, and returning them
+    would invite the next comparison to be made against the whole string again.
+    """
+
+    if not isinstance(media_type, str):
+        return ""
+    return media_type.partition(";")[0].strip().lower()
+
+
+def suffix_for(media_type: Any) -> str:
+    """A file extension chosen from the media type, and ``.bin`` when it is not one we name.
+
+    Deliberately the project's own table rather than ``mimetypes``: the stdlib mapping differs
+    between hosts, and a dataset whose extension depends on the machine that downloaded it is a
+    support question waiting to happen.
+
+    The lookup is on the parsed essence, so ``text/csv`` and ``text/csv; charset=utf-8`` name the
+    same file. The table's own keys carry no parameters, so this widens what is recognised and
+    can never change what an already-recognised type resolves to.
+    """
+
+    return SUFFIX_BY_MEDIA_TYPE.get(media_type_essence(media_type), ".bin")
+
+
+#: The two refusals a receipt read may legitimately answer with, and what each one means.
+#: ``THIN_NOT_FOUND`` is the ordinary state of every run that has not produced that record yet;
+#: ``THIN_AUTHORIZATION_DENIED`` is a real answer too -- candidate evidence needs `review:read`,
+#: which a viewer may not hold. Neither is reported as absence, because "there is nothing" and
+#: "you may not see it" are different sentences and a reader must be told which one they got.
+RECEIPT_UNAVAILABLE_CODES = ("THIN_AUTHORIZATION_DENIED", "THIN_NOT_FOUND")
+
+
+def optional_receipt(name: str, call: Any, unavailable: dict[str, str]) -> Any:
+    """Return what ``call`` answers, recording WHY it could not be read when it could not."""
+
+    try:
+        return call()
+    except ThinLaneError as refusal:
+        if refusal.code in RECEIPT_UNAVAILABLE_CODES:
+            unavailable[name] = refusal.code
+            return None
+        raise
+
+
+def print_stderr(text: str) -> None:
+    print(text, file=sys.stderr, flush=True)
+
+
+__all__ = [
+    "DEFAULT_DOWNLOAD_KINDS",
+    "DOWNLOAD_SCHEMA",
+    "RECEIPT_UNAVAILABLE_CODES",
+    "RESOURCE_CLASSES",
+    "STATUS_SCHEMA",
+    "SUBMIT_SCHEMA",
+    "WATCH_SCHEMA",
+    "artifact_records",
+    "download",
+    "optional_receipt",
+    "run_receipt",
+    "status",
+    "submit",
+    "suffix_for",
+    "watch",
+]

mostlyright/data_harness/thin/download.py

@@ -0,0 +1,194 @@
+"""Fetch one artifact through a signed URL, and refuse the bytes unless they are the bytes.
+
+The signed URL is a short-lived bearer capability that points straight at object storage, so two
+things are true at once: nothing here needs a Studio credential, and nothing here may trust what
+comes back. Studio hands over ``expected_content_digest`` and ``expected_size_bytes`` alongside the
+URL precisely so the client can settle that itself, and this module treats both as mandatory
+rather than advisory -- a download whose digest does not match is deleted, not written with a
+warning.
+
+The file is written with ``O_EXCL``: a hosted download never overwrites something already on this
+machine, because the one thing worse than not having the dataset is silently replacing a different
+one with it.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import os
+import re
+from collections.abc import Mapping
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+from mostlyright.data_harness.thin.transport import (
+    REQUEST_TIMEOUT_SECONDS,
+    ThinLaneError,
+    Transport,
+    UrlLibTransport,
+    signed_transfer_url,
+)
+
+#: Read granularity. The same 64 KiB ``acquisition.result_download`` streams a crawl result at.
+DOWNLOAD_CHUNK_BYTES = 65_536
+
+#: A ceiling on any single artifact this lane will write, so a wrong ``expected_size_bytes`` cannot
+#: fill the volume before the digest check gets a chance to refuse. Sixteen gibibytes is far above
+#: any table this product builds and far below a disk.
+MAX_ARTIFACT_BYTES = 16 * 1024 * 1024 * 1024
+
+_PREFIXED_DIGEST = re.compile(r"^sha256:[0-9a-f]{64}$")
+
+
+@dataclass(frozen=True)
+class DownloadedArtifact:
+    """What one verified download produced."""
+
+    path: Path
+    artifact_id: str
+    media_type: str
+    size_bytes: int
+    content_digest: str
+
+    def to_receipt(self) -> dict[str, Any]:
+        return {
+            "path": str(self.path),
+            "artifact_id": self.artifact_id,
+            "media_type": self.media_type,
+            "size_bytes": self.size_bytes,
+            "content_digest": self.content_digest,
+        }
+
+
+def _required_headers(session: Mapping[str, Any]) -> dict[str, str]:
+    """The headers the signature covers. Sending fewer, or more, invalidates it."""
+
+    declared = session.get("required_headers")
+    if declared is None:
+        return {}
+    if not isinstance(declared, list):
+        raise ThinLaneError("THIN_RESPONSE_INVALID", "required_headers is not a list")
+    headers: dict[str, str] = {}
+    for item in declared:
+        if not isinstance(item, Mapping):
+            raise ThinLaneError("THIN_RESPONSE_INVALID", "a required header is not an object")
+        name = item.get("name")
+        value = item.get("value")
+        if not isinstance(name, str) or not isinstance(value, str) or not name:
+            raise ThinLaneError(
+                "THIN_RESPONSE_INVALID", "a required header is not a name/value pair"
+            )
+        headers[name] = value
+    return headers
+
+
+def download_signed_artifact(
+    session: Mapping[str, Any],
+    destination: Path,
+    *,
+    transport: Transport | None = None,
+) -> DownloadedArtifact:
+    """Fetch the bytes one signed session points at, verify them, and write them once."""
+
+    url = signed_transfer_url(session.get("signed_url", ""), "the signed download URL")
+    expected_digest = session.get("expected_content_digest")
+    expected_size = session.get("expected_size_bytes")
+    if not isinstance(expected_digest, str) or _PREFIXED_DIGEST.fullmatch(expected_digest) is None:
+        raise ThinLaneError(
+            "THIN_RESPONSE_INVALID", "the signed session carried no sha256 content digest"
+        )
+    if not isinstance(expected_size, int) or isinstance(expected_size, bool) or expected_size < 0:
+        raise ThinLaneError("THIN_RESPONSE_INVALID", "the signed session carried no byte size")
+    if expected_size > MAX_ARTIFACT_BYTES:
+        raise ThinLaneError(
+            "THIN_ARTIFACT_TOO_LARGE",
+            f"the artifact declares {expected_size} bytes, above this lane's ceiling",
+        )
+
+    selected = transport or UrlLibTransport()
+    destination.parent.mkdir(parents=True, exist_ok=True)
+    flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL
+    if hasattr(os, "O_NOFOLLOW"):
+        flags |= os.O_NOFOLLOW
+    # Windows only. The artifact this writes is Parquet, and its digest is verified against the
+    # one the signed session declared -- so a CRLF translation on the way to disk would not be a
+    # cosmetic difference, it would be a download that always fails its own digest check.
+    flags |= getattr(os, "O_BINARY", 0)
+    try:
+        descriptor = os.open(destination, flags, 0o600)
+    except FileExistsError as error:
+        raise ThinLaneError(
+            "THIN_DESTINATION_EXISTS",
+            f"{destination} already exists; a hosted download never overwrites a local file",
+        ) from error
+
+    try:
+        handle = os.fdopen(descriptor, "wb", closefd=True)
+    except BaseException:
+        # `fdopen` owns the descriptor only once it succeeds, so this is the one window where
+        # closing it is still ours to do.
+        os.close(descriptor)
+        destination.unlink(missing_ok=True)
+        raise
+
+    digest = hashlib.sha256()
+    written = 0
+    try:
+        # The file handle enters FIRST so that a stream that fails to open still closes it.
+        with (
+            handle,
+            selected.stream(url, _required_headers(session), REQUEST_TIMEOUT_SECONDS) as response,
+        ):
+            status = int(getattr(response, "status", 200))
+            if status != 200:
+                raise ThinLaneError(
+                    "THIN_DOWNLOAD_REFUSED",
+                    f"the signed download answered HTTP {status}; the session may have expired",
+                )
+            while True:
+                chunk = response.read(DOWNLOAD_CHUNK_BYTES)
+                if not chunk:
+                    break
+                written += len(chunk)
+                if written > expected_size:
+                    raise ThinLaneError(
+                        "THIN_DOWNLOAD_OVERRUN",
+                        "the signed download returned more bytes than it declared",
+                    )
+                digest.update(chunk)
+                handle.write(chunk)
+        if written != expected_size:
+            raise ThinLaneError(
+                "THIN_DOWNLOAD_TRUNCATED",
+                f"the signed download returned {written} of {expected_size} declared bytes",
+            )
+        observed = "sha256:" + digest.hexdigest()
+        if observed != expected_digest:
+            raise ThinLaneError(
+                "THIN_DOWNLOAD_DIGEST_MISMATCH",
+                "the downloaded bytes do not hash to the digest Studio declared",
+            )
+    except BaseException:
+        # Nothing partial and nothing unverified is left behind for a later command to mistake
+        # for a dataset.
+        destination.unlink(missing_ok=True)
+        raise
+
+    media_type = session.get("media_type")
+    artifact_id = session.get("artifact_id")
+    return DownloadedArtifact(
+        path=destination,
+        artifact_id=artifact_id if isinstance(artifact_id, str) else "",
+        media_type=media_type if isinstance(media_type, str) else "application/octet-stream",
+        size_bytes=written,
+        content_digest=expected_digest,
+    )
+
+
+__all__ = [
+    "DOWNLOAD_CHUNK_BYTES",
+    "MAX_ARTIFACT_BYTES",
+    "DownloadedArtifact",
+    "download_signed_artifact",
+]