snowflake-sandbox-python 0.2.1a1__py3-none-any.whl

snowflake/sandbox/exec_stream.py

@@ -0,0 +1,264 @@
+"""Public ``ExecStream`` async iterator over SSE-framed exec output.
+
+Returned by `AsyncSandbox.exec_stream()`. Wraps the raw SSE event stream
+from ``_transport`` into ``StreamLine(stream, data)`` tuples and captures
+the terminal ``exit`` frame so the caller can read ``exit_code`` and
+``elapsed_ms`` after the iterator drains.
+
+The public stream API does NOT raise on non-zero exit. Users check
+``stream.exit_code`` after iterating. ``AsyncSandbox.exec`` (the
+drain-to-completion convenience) DOES raise ``SandboxExecError`` on
+non-zero exit.
+
+It does raise on an ``error`` frame, which is a different thing from a
+non-zero exit: it means the exec never produced a result at all. It also raises
+when the stream ends with **no** terminal frame, which means the output was
+truncated and the exit status is unknown -- silently ending there made a dropped
+relay indistinguishable from a clean finish.
+"""
+
+from __future__ import annotations
+
+from collections.abc import AsyncIterator
+from typing import TYPE_CHECKING, Self
+
+from snowflake.sandbox.exceptions import (
+    SandboxContractWarning,
+    SandboxExecError,
+)
+from snowflake.sandbox.types import StreamLine, StreamName
+
+if TYPE_CHECKING:
+    from snowflake.sandbox._transport import SSEEvent
+
+__all__ = ["ExecStream"]
+
+
+class ExecStream(AsyncIterator[StreamLine]):
+    """Async iterator over the lines of a streaming exec.
+
+    Properties (populated after the iterator drains, or on early
+    `aclose()`):
+
+    * ``exit_code`` -- process exit code, or ``None`` if cancelled before
+      the exit frame arrived.
+    * ``elapsed_ms`` -- server-reported wall-clock duration, or ``None``.
+    * ``error`` -- the server's message from an ``error`` frame, if one
+      arrived. Set just before ``SandboxExecError`` is raised.
+    * ``gaps`` -- how many times output was lost to backlog eviction while the
+      stream was resuming across a transport cut. ``0`` for a stream that was
+      never cut, or one whose every reconnect landed inside the retained
+      backlog. Non-zero means some output between two reconnect points is gone.
+    """
+
+    def __init__(self, source: AsyncIterator[SSEEvent]) -> None:
+        self._source = source
+        self.exit_code: int | None = None
+        self.elapsed_ms: int | None = None
+        self.error: str | None = None
+        self.gaps = 0
+        self._warned_gaps = False
+        self._closed = False
+        # Whether a terminal frame (`exit` or `error`) was seen. A stream that
+        # ends without one was truncated, which must not read as completion.
+        self._terminated = False
+
+    def __aiter__(self) -> Self:
+        return self
+
+    async def __anext__(self) -> StreamLine:
+        if self._closed:
+            raise StopAsyncIteration
+        while True:
+            try:
+                evt = await self._source.__anext__()
+            except StopAsyncIteration:
+                self._closed = True
+                if not self._terminated:
+                    # The contract is that every exec ends with a terminal `exit`
+                    # frame carrying its code. A clean close without one means the
+                    # relay, the container, or a proxy dropped the stream partway
+                    # -- and silently raising StopAsyncIteration with exit_code
+                    # still None made a truncated exec indistinguishable from one
+                    # that completed or was cancelled.
+                    raise SandboxExecError(
+                        "exec stream ended without a terminal exit frame: the output "
+                        "may be incomplete and the exit status is unknown",
+                        exit_code=-1,
+                    ) from None
+                raise
+            event_name = evt.event
+            if event_name == "exit":
+                self._apply_exit_frame(evt)
+                # exit terminates the stream.
+                self._terminated = True
+                self._closed = True
+                await self._maybe_aclose_source()
+                self._warn_if_gaps()
+                raise StopAsyncIteration
+            if event_name in ("stdout", "stderr"):
+                stream: StreamName = "stderr" if event_name == "stderr" else "stdout"
+                data = _extract_data(evt)
+                # Strip a single trailing newline; preserve internal ones.
+                if data.endswith("\n"):
+                    data = data[:-1]
+                return StreamLine(stream=stream, data=data)
+            if event_name == "error":
+                message = _error_frame_message(evt)
+                self.error = message
+                self._terminated = True
+                self._closed = True
+                await self._maybe_aclose_source()
+                raise SandboxExecError(f"exec stream failed server-side: {message}", exit_code=-1)
+            if event_name == "gap":
+                # The stream resumed past output the server's backlog had already
+                # evicted (see _exec_stream_frames' reconnect). Count it so the
+                # caller can tell a clean run from one with a hole, rather than
+                # letting the loss pass as an unknown event.
+                self.gaps += 1
+                continue
+            # Unknown event type -- skip.
+            continue
+
+    def _apply_exit_frame(self, evt: SSEEvent) -> None:
+        """Read ``exit_code`` and ``elapsed_ms`` out of an ``exit`` frame.
+
+        Best-effort: a frame whose body is not the expected JSON object leaves
+        both fields untouched rather than raising, since the frame's arrival is
+        itself the signal that the stream terminated.
+        """
+        try:
+            payload = evt.json() or {}
+        except Exception:
+            payload = {}
+        if isinstance(payload, dict):
+            self.exit_code = _coerce_exit_code(payload.get("code"))
+            elapsed = payload.get("elapsed_ms")
+            if isinstance(elapsed, (int, float)) and not isinstance(elapsed, bool):
+                self.elapsed_ms = int(elapsed)
+
+    async def aclose(self) -> None:
+        """Close the local SSE stream and stop iterating.
+
+        This tears down the client-side stream only. It does **not** kill the
+        command running in the container: there is no server-side cancel/DELETE
+        route, so the process outlives the closed stream and keeps consuming
+        CPU/cost until it exits on its own (or the sandbox is destroyed). If you
+        need the work to stop, run it under a ``timeout`` or destroy the sandbox.
+        """
+        if self._closed:
+            return
+        # An explicit cancel is a deliberate early exit, not a truncation, so it
+        # does not raise on the next __anext__.
+        self._terminated = True
+        self._closed = True
+        await self._maybe_aclose_source()
+        self._warn_if_gaps()
+
+    def _warn_if_gaps(self) -> None:
+        """Warn once, when the stream finishes, if output was lost.
+
+        `gaps` was counted and never surfaced. The reading that matters is a run that
+        finishes ``exit_code 0`` with a hole in the middle of its output and says
+        nothing -- the case a caller is least likely to look for. Programmatic callers
+        still read ``.gaps``.
+
+        Called only from the two paths that are NOT already raising: the terminal
+        ``exit`` frame and an explicit ``aclose()``. On the truncation and server-error
+        paths an exception already carries the diagnosis, and warning there would fire
+        *while* one is being handled -- which, with warnings-as-errors, breaks an async
+        generator's teardown ("generator didn't stop after athrow()").
+        """
+        if self._warned_gaps or not self.gaps:
+            return
+        self._warned_gaps = True
+        import warnings
+
+        warnings.warn(
+            f"exec stream lost output in {self.gaps} place(s): it resumed past data "
+            f"the server's backlog had already evicted, so output between two "
+            f"reconnect points is missing. Read the stream's .gaps attribute to detect this "
+            f"programmatically.",
+            SandboxContractWarning,
+            stacklevel=3,
+        )
+
+    async def _maybe_aclose_source(self) -> None:
+        aclose = getattr(self._source, "aclose", None)
+        if aclose is not None:
+            try:
+                await aclose()
+            except Exception:
+                pass
+
+
+def _error_frame_message(evt: SSEEvent) -> str:
+    """Extract the human-readable message from an ``error`` frame.
+
+    The gateway emits ``event: error`` with ``{"message": ...}`` when its relay
+    fails *after* the SSE headers are committed, so a JSON error response is no
+    longer possible. Skipping it as an unknown event -- which is what we used to
+    do -- ended the stream with ``exit_code`` still ``None`` and threw the
+    server's only explanation away. A failed exec must not be silently
+    indistinguishable from one that printed nothing. Falls back to the raw frame
+    ``data`` (then a placeholder) when no JSON ``message`` is present.
+    """
+    try:
+        payload = evt.json() or {}
+    except Exception:
+        payload = {}
+    if isinstance(payload, dict):
+        raw = payload.get("message")
+        if isinstance(raw, str) and raw:
+            return raw
+    return evt.data or "no detail"
+
+
+def _coerce_exit_code(raw: object) -> int | None:
+    """Read an exit code that may arrive as an int, a float, or a string.
+
+    The frame's ``code`` is an int on the wire today, but a stricter
+    ``isinstance(code, int)`` test silently dropped anything else -- leaving
+    ``exit_code`` at ``None``, i.e. reading a real exit status as "never
+    finished". A value we cannot read at all stays ``None``.
+    """
+    if isinstance(raw, bool):
+        return int(raw)
+    if isinstance(raw, int):
+        return raw
+    if isinstance(raw, float):
+        return int(raw)
+    if isinstance(raw, str):
+        try:
+            return int(raw.strip(), 10)
+        except ValueError:
+            try:
+                return int(float(raw.strip()))
+            except ValueError:
+                return None
+    return None
+
+
+def _extract_data(evt: SSEEvent) -> str:
+    """Pull the ``data`` field from a stdout/stderr SSE frame.
+
+    Servers may either send raw text (``data: hello``) or a JSON payload
+    ``data: {"text": "hello"}``. We accept both and prefer the JSON
+    ``text`` / ``data`` field when present.
+    """
+    raw = evt.data
+    if not raw:
+        return ""
+    stripped = raw.strip()
+    if stripped.startswith("{") and stripped.endswith("}"):
+        try:
+            obj = evt.json()
+        except Exception:
+            return raw
+        if isinstance(obj, dict):
+            for key in ("text", "data", "line"):
+                v = obj.get(key)
+                if isinstance(v, str):
+                    return v
+        return raw
+    return raw