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

snowflake/sandbox/sync_exec_stream.py

@@ -0,0 +1,238 @@
+"""Public ``SyncExecStream`` iterator over SSE-framed exec output.
+
+Returned by `Sandbox.exec_stream()`. Wraps the raw SSE event stream
+from ``_sync_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.
+
+This is the synchronous counterpart to ``ExecStream``.
+
+The public stream API does NOT raise on non-zero exit. Users check
+``stream.exit_code`` after iterating. ``Sandbox.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 Iterator
+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__ = ["SyncExecStream"]
+
+
+class SyncExecStream(Iterator[StreamLine]):
+    """Synchronous iterator over the lines of a streaming exec.
+
+    Properties (populated after the iterator drains, or on early
+    `close()`):
+
+    * ``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: Iterator[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
+        self._terminated = False
+
+    def __iter__(self) -> Self:
+        return self
+
+    def __next__(self) -> StreamLine:
+        if self._closed:
+            raise StopIteration
+        while True:
+            try:
+                evt = next(self._source)
+            except StopIteration:
+                self._closed = True
+                if not self._terminated:
+                    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)
+                self._terminated = True
+                self._closed = True
+                self._maybe_close_source()
+                self._warn_if_gaps()
+                raise StopIteration
+            if event_name in ("stdout", "stderr"):
+                stream: StreamName = "stderr" if event_name == "stderr" else "stdout"
+                data = _extract_data(evt)
+                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
+                self._maybe_close_source()
+                raise SandboxExecError(f"exec stream failed server-side: {message}", exit_code=-1)
+            if event_name == "gap":
+                self.gaps += 1
+                continue
+            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)
+
+    def close(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
+        self._terminated = True
+        self._closed = True
+        self._maybe_close_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 ``close()``. On the truncation and server-error
+        paths an exception already carries the diagnosis, and warning there would fire
+        *while* one is being handled -- which, in the async twin, breaks generator
+        teardown under warnings-as-errors. Kept symmetrical so the two agree.
+        """
+        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,
+        )
+
+    def _maybe_close_source(self) -> None:
+        close_method = getattr(self._source, "close", None)
+        if close_method is not None:
+            try:
+                close_method()
+            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."""
+    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."""
+    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

snowflake/sandbox/sync_files.py

@@ -0,0 +1,377 @@
+"""Sync byte transfer over the container's ``/files`` HTTP route, for ``Sandbox``.
+
+This is the synchronous counterpart to ``files.py`` — all operations block until
+complete instead of returning coroutines. See ``files.py`` for the detailed
+documentation of each operation's behavior, constraints, and platform wiring.
+
+The key differences from the async module:
+- Uses ``SyncTransport`` instead of ``Transport``
+- Uses ``time.sleep()`` instead of ``asyncio.sleep()``
+- Direct blocking I/O instead of ``asyncio.to_thread()``
+
+The module split mirrors the async side one-for-one: `_sync_fs_ops` holds the
+in-container directory operations and `_sync_watch` the change monitor.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import os
+import time
+import uuid
+from pathlib import Path
+from typing import TYPE_CHECKING, TypeVar
+
+from snowflake.sandbox._sync_transport import SyncTransport
+
+# Constants and the pure (non-coroutine) helpers are single-sourced from the async
+# module so the two clients cannot drift -- `_is_bad_gateway` (what counts as a
+# retryable 502) and `_prune_empty_dirs` (undo a failed download's mkdir) have no
+# async in them, so the sync side imports rather than re-copies. Only the retry
+# *driver* is a genuine twin below (it must ``time.sleep`` instead of ``await``).
+from snowflake.sandbox._upload_plan import UploadPlan, plan_directory
+from snowflake.sandbox.exceptions import (
+    SandboxError,
+    SandboxFileTooLargeError,
+    SandboxTransportError,
+)
+from snowflake.sandbox.files import (
+    _FILES_RETRY_ATTEMPTS,
+    _FILES_RETRY_BASE_DELAY,
+    MAX_FILE_BYTES,
+    _is_bad_gateway,
+    _join,
+    _prune_empty_dirs,
+    _require_absolute,
+    _resolve_local_read,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from snowflake.sandbox.sync_client import Sandbox
+
+_T = TypeVar("_T")
+
+__all__ = [
+    "upload_file",
+    "upload_dir",
+    "download_file",
+    "stage_put",
+    "stage_get",
+    "read_text",
+    "read_bytes",
+    "write_text",
+    "write_bytes",
+]
+
+
+def _retry_502(attempt: Callable[[], _T]) -> _T:
+    """Run ``attempt`` up to `_FILES_RETRY_ATTEMPTS` times, retrying only a transient
+    502 Bad Gateway (`_is_bad_gateway`) with exponential backoff — sync twin of
+    ``files._retry_502``.
+
+    The only difference from the async driver is ``time.sleep`` in place of
+    ``await asyncio.sleep``; the retry policy (which errors are terminal, the
+    backoff schedule) lives in the shared `_is_bad_gateway` predicate and the shared
+    constants, so the two cannot drift.
+    """
+    for attempt_no in range(_FILES_RETRY_ATTEMPTS):
+        try:
+            return attempt()
+        except SandboxTransportError as exc:
+            if not _is_bad_gateway(exc) or attempt_no == _FILES_RETRY_ATTEMPTS - 1:
+                raise
+            time.sleep(_FILES_RETRY_BASE_DELAY * (2**attempt_no))
+    raise AssertionError(
+        "unreachable: the loop returns or raises every iteration"
+    )  # pragma: no cover
+
+
+def _files_request_with_retry(
+    t: SyncTransport,
+    method: str,
+    endpoint: str,
+    *,
+    params: dict[str, str] | None = None,
+    content: bytes | None = None,
+    extra_headers: dict[str, str] | None = None,
+) -> None:
+    """Issue a ``/files`` request, retrying a transient 502 (sync version)."""
+
+    def _attempt() -> None:
+        t.request(method, endpoint, params=params, content=content, extra_headers=extra_headers)
+
+    _retry_502(_attempt)
+
+
+def upload_file(
+    sandbox: Sandbox,
+    local: str | Path,
+    remote: str,
+    *,
+    transport: SyncTransport | None = None,
+) -> None:
+    """Upload a local file into the sandbox at ``remote`` (sync version).
+
+    See ``files.upload_file`` for full documentation.
+    """
+    if not sandbox.id:
+        raise SandboxError("sandbox must be created before uploading files")
+    src = _resolve_local_read(local)
+    size = src.stat().st_size
+    if size > MAX_FILE_BYTES:
+        raise SandboxFileTooLargeError(
+            f"{src} is {size} bytes, over the {MAX_FILE_BYTES} byte limit for "
+            "upload_file; mount a stage and write through the mount instead"
+        )
+    t = transport or sandbox._transport
+    dest = _require_absolute(remote)
+
+    def _attempt() -> None:
+        # Reopened per attempt so a retry re-sends the whole file, never a stream
+        # the previous try already consumed.
+        with src.open("rb") as body:
+            t.request(
+                "PUT",
+                f"containers/{sandbox.id}/files",
+                params={"path": dest},
+                content=body,
+                extra_headers={"Content-Type": "application/octet-stream"},
+            )
+
+    _retry_502(_attempt)
+
+
+def download_file(
+    sandbox: Sandbox,
+    remote: str,
+    local: str | Path,
+    *,
+    transport: SyncTransport | None = None,
+) -> None:
+    """Download ``remote`` from the sandbox to ``local`` (sync version).
+
+    See ``files.download_file`` for full documentation.
+    """
+    if not sandbox.id:
+        raise SandboxError("sandbox must be created before downloading files")
+    dst = Path(local).expanduser().resolve()
+    if dst.is_dir():
+        raise SandboxError(f"local target is a directory, expected a file path: {dst}")
+    t = transport or sandbox._transport
+    path = _require_absolute(remote)
+
+    highest_existing = dst.parent
+    while not highest_existing.exists():
+        highest_existing = highest_existing.parent
+    dst.parent.mkdir(parents=True, exist_ok=True)
+    # Unique per call (pid + random): two concurrent downloads to the same local
+    # path in one process must not share a temp file and race on replace/unlink.
+    tmp = dst.with_name(f"{dst.name}.download-{os.getpid()}-{uuid.uuid4().hex}")
+
+    def _attempt() -> None:
+        # `wb` truncates any leftover from a prior failed try, so each attempt
+        # writes the whole body afresh; on success os.replace consumes tmp.
+        with tmp.open("wb") as out:
+            resp = t.download(f"containers/{sandbox.id}/files", out.write, params={"path": path})
+            if "json" in resp.headers.get("Content-Type", "").lower():
+                raise SandboxError(
+                    f"expected file bytes for {path}, got a JSON response "
+                    f"({resp.headers.get('Content-Type')})"
+                )
+        os.replace(tmp, dst)
+
+    try:
+        _retry_502(_attempt)
+    except BaseException:
+        # One teardown for every failure -- a non-502, an exhausted retry budget, a
+        # JSON body, or a KeyboardInterrupt. Drop the partial temp file and prune only
+        # the directories this call created (see `_prune_empty_dirs`).
+        with contextlib.suppress(OSError):
+            tmp.unlink()
+        _prune_empty_dirs(dst.parent, highest_existing)
+        raise
+
+
+def stage_put(
+    sandbox: Sandbox,
+    local: str,
+    stage_path: str,
+    *,
+    transport: SyncTransport | None = None,
+) -> None:
+    """Not implemented — the ``stage/put`` route returns 501 (sync version).
+
+    See ``files.stage_put`` for full documentation.
+    """
+    if not sandbox.id:
+        raise SandboxError("sandbox must be created before stage I/O")
+    t = transport or sandbox._transport
+    t.request(
+        "POST",
+        f"containers/{sandbox.id}/stage/put",
+        json_body={"local": local, "stage": stage_path},
+    )
+
+
+def stage_get(
+    sandbox: Sandbox,
+    stage_path: str,
+    local: str,
+    *,
+    transport: SyncTransport | None = None,
+) -> None:
+    """Not implemented — the ``stage/get`` route returns 501 (sync version).
+
+    See ``files.stage_get`` for full documentation.
+    """
+    if not sandbox.id:
+        raise SandboxError("sandbox must be created before stage I/O")
+    t = transport or sandbox._transport
+    t.request(
+        "POST",
+        f"containers/{sandbox.id}/stage/get",
+        json_body={"stage": stage_path, "local": local},
+    )
+
+
+# ---- in-container read / write over the /files byte route ----------------
+
+
+def read_bytes(sandbox: Sandbox, path: str, *, transport: SyncTransport | None = None) -> bytes:
+    """Return the file at ``path`` as bytes (sync version).
+
+    See ``files.read_bytes`` for full documentation.
+    """
+    if not sandbox.id:
+        raise SandboxError("sandbox must be created before reading files")
+    t = transport or sandbox._transport
+    p = _require_absolute(path)
+
+    def _attempt() -> bytes:
+        # Fresh buffer per attempt: a retried read must not append to bytes a prior
+        # try already collected.
+        buf = bytearray()
+
+        def _sink(chunk: bytes) -> None:
+            buf.extend(chunk)
+            if len(buf) > MAX_FILE_BYTES:
+                raise SandboxFileTooLargeError(
+                    f"{p} exceeds the {MAX_FILE_BYTES} byte read limit; mount a stage "
+                    "and read through the mount instead"
+                )
+
+        resp = t.download(f"containers/{sandbox.id}/files", _sink, params={"path": p})
+        if "json" in resp.headers.get("Content-Type", "").lower():
+            raise SandboxError(
+                f"expected file bytes for {p}, got a JSON response "
+                f"({resp.headers.get('Content-Type')})"
+            )
+        return bytes(buf)
+
+    return _retry_502(_attempt)
+
+
+def read_text(
+    sandbox: Sandbox,
+    path: str,
+    *,
+    encoding: str = "utf-8",
+    transport: SyncTransport | None = None,
+) -> str:
+    """Return the file at ``path`` decoded as text (sync version)."""
+    raw = read_bytes(sandbox, path, transport=transport)
+    return raw.decode(encoding)
+
+
+def write_bytes(
+    sandbox: Sandbox,
+    path: str,
+    data: bytes,
+    *,
+    transport: SyncTransport | None = None,
+) -> None:
+    """Write ``data`` to ``path`` inside the sandbox (sync version).
+
+    See ``files.write_bytes`` for full documentation.
+    """
+    if not sandbox.id:
+        raise SandboxError("sandbox must be created before writing files")
+    if len(data) > MAX_FILE_BYTES:
+        raise SandboxFileTooLargeError(
+            f"{len(data)} bytes is over the {MAX_FILE_BYTES} byte limit for "
+            "write_bytes; mount a stage and write through the mount instead"
+        )
+    t = transport or sandbox._transport
+    _files_request_with_retry(
+        t,
+        "PUT",
+        f"containers/{sandbox.id}/files",
+        params={"path": _require_absolute(path)},
+        content=bytes(data),
+        extra_headers={"Content-Type": "application/octet-stream"},
+    )
+
+
+def write_text(
+    sandbox: Sandbox,
+    path: str,
+    data: str,
+    *,
+    encoding: str = "utf-8",
+    transport: SyncTransport | None = None,
+) -> None:
+    """Write text to ``path``, encoded ``utf-8`` by default (sync version)."""
+    write_bytes(sandbox, path, data.encode(encoding), transport=transport)
+
+
+def upload_dir(
+    sandbox: Sandbox,
+    local_dir: str | Path,
+    remote_dir: str,
+    *,
+    exclude: list[str] | None = None,
+    include: list[str] | None = None,
+    allow_credential_files: list[str] | None = None,
+    dry_run: bool = False,
+    on_file: Callable[[int, int, str], None] | None = None,
+    transport: SyncTransport | None = None,
+) -> UploadPlan:
+    """Upload a local directory tree into the sandbox, returning what it did (sync).
+
+    See ``_SyncFilesMixin.upload_dir`` for the full contract.
+    """
+    if not sandbox.id:
+        raise SandboxError("sandbox must be created before uploading files")
+    dest_root = _require_absolute(remote_dir).rstrip("/") or "/"
+    plan = plan_directory(
+        local_dir,
+        dest_root=dest_root,
+        exclude=exclude,
+        include=include,
+        allow_credential_files=allow_credential_files,
+        max_file_bytes=MAX_FILE_BYTES,
+        strip_root=True,
+    )
+    if dry_run:
+        return plan
+    # Checked after the `dry_run` return below, not before it: a dry run sends no
+    # bytes, and raising there denied the one call whose whole purpose is previewing
+    # the transfer any chance to report `plan.oversized` back to the caller.
+    if plan.oversized:
+        listed = ", ".join(f"{s.rel} ({s.size} bytes)" for s in plan.oversized[:5])
+        more = f", and {len(plan.oversized) - 5} more" if len(plan.oversized) > 5 else ""
+        raise SandboxFileTooLargeError(
+            f"{len(plan.oversized)} file(s) exceed the {MAX_FILE_BYTES} byte limit "
+            f"for upload_dir: {listed}{more}. Mount a stage for bulk data, or pass "
+            f"exclude= to skip them."
+        )
+    total = plan.file_count
+    for index, item in enumerate(plan.selected, start=1):
+        # One request per file: there is no bulk route. `upload_file` creates parent
+        # directories itself, so no make_directory round trips.
+        upload_file(sandbox, item.source, _join(dest_root, item.rel), transport=transport)
+        if on_file is not None:
+            on_file(index, total, item.rel)
+    return plan