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

snowflake/sandbox/files.py

@@ -0,0 +1,547 @@
+"""Byte transfer over the container's ``/files`` HTTP route, for ``AsyncSandbox``.
+
+The transfer surface treats a container as its own filesystem namespace; the path
+inside the sandbox is a regular absolute POSIX path.
+
+Automatic retries
+-----------------
+The byte-transfer functions (``write_bytes``, ``write_text``, ``upload_file``,
+``read_bytes``, ``read_text``, ``download_file``) automatically retry transient
+502 Bad Gateway errors with exponential backoff. This handles a platform timing
+edge case where the container's internal HTTP server may not be fully ready
+immediately after ``AsyncSandbox.create()`` returns. Retries are transparent to
+callers — most users will never see a 502. The SDK retries up to 3 times with
+delays of 0.5s, 1s, and 2s between attempts.
+
+What is wired, and how:
+
+* **Byte transfer** — ``upload_file`` / ``download_file`` / ``read_bytes`` /
+  ``read_text`` / ``write_bytes`` / ``write_text`` ride the real
+  ``PUT/GET /containers/{id}/files?path=/abs/name`` route, which the control
+  plane forwards to the container's own ``/files`` endpoint. Bytes stream through
+  the control plane; a single transfer is capped at ``MAX_FILE_BYTES`` (mount a
+  stage for anything larger).
+* **Stage I/O** — ``stage_put`` / ``stage_get`` are **not** wired on any backend
+  yet: the route returns 501. Mount a stage (``StageMount``) for stage-backed
+  files.
+
+Two other mechanisms reach a sandbox filesystem, and each has its own module now,
+because neither rides this route and their path semantics differ from it:
+
+* `_fs_ops` — ``list_files`` / ``stat`` / ``make_directory`` / ``remove``, an
+  in-container ``python3`` helper driven by ``exec``. Unlike every transfer here,
+  those do **not** require an absolute path.
+* `_watch` — ``watch``, a ``ctypes``/``inotify(7)`` monitor driven by the
+  resumable ``exec_stream``. It does require an absolute path, and calls
+  `_require_absolute` here, so the rule is stated once.
+
+The file path travels as a query parameter rather than extra URL path segments.
+Snowflake's REST layer validates every path segment against a narrow charset and
+forwards the query string verbatim, so as segments a filename containing a space
+(or ``#``, ``+``, ``,``, ``@``, or any non-ASCII byte) is rejected before it
+reaches the sandbox, and a leading ``/`` cannot be expressed at all.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import os
+import uuid
+from pathlib import Path
+from typing import TYPE_CHECKING, TypeVar
+
+from snowflake.sandbox._transport import Transport
+from snowflake.sandbox._upload_plan import UploadPlan, plan_directory
+from snowflake.sandbox.exceptions import (
+    SandboxError,
+    SandboxFileTooLargeError,
+    SandboxTransportError,
+)
+
+# Retry config for /files: the container's internal HTTP server may not be ready
+# immediately after create() returns "ready" (Snowflake controller ready != app server ready).
+# A 502 Bad Gateway on the first attempt is retried with exponential backoff.
+_FILES_RETRY_ATTEMPTS = 3
+_FILES_RETRY_BASE_DELAY = 0.5  # seconds
+
+# Ceiling on one transfer, matching what the platform will actually carry:
+# Snowflake's nginx front end caps request bodies at 30 MB, so a larger upload is
+# rejected before it reaches any sandbox code. Checked locally so the caller gets a
+# named error instead of an opaque 413 from a hop they cannot see.
+#
+# Two distinct server-side caps exist and only one applies here: the ``/files``
+# route has its OWN 25 MiB body limit, so a transfer up to this size is accepted —
+# the smaller 4 MiB cap is the JSON-body limit for routes like ``/exec`` and does
+# not apply to the streamed ``/files`` body. Keeping the client ceiling at 25 MiB
+# matches the route's own limit exactly.
+MAX_FILE_BYTES = 25 * 1024 * 1024
+
+if TYPE_CHECKING:
+    from collections.abc import Awaitable, Callable
+
+    from snowflake.sandbox.client import AsyncSandbox
+
+_T = TypeVar("_T")
+
+__all__ = [
+    "upload_file",
+    "upload_dir",
+    "download_file",
+    "stage_put",
+    "stage_get",
+    "read_text",
+    "read_bytes",
+    "write_text",
+    "write_bytes",
+]
+
+
+def _resolve_local_read(local: str | Path) -> Path:
+    """The local path to read bytes from, or a stdlib error naming the real problem.
+
+    A directory gets `IsADirectoryError` naming `upload_dir`, not the
+    `FileNotFoundError` this used to raise for it: the folder plainly exists, so
+    "local file not found" sent people looking for a typo instead of telling them
+    the verb they wanted. Both are stdlib errors rather than `SandboxError`s for the
+    same reason -- a local-filesystem miss is not a sandbox failure.
+    """
+    p = Path(local).expanduser().resolve()
+    if p.is_dir():
+        raise IsADirectoryError(
+            f"{p} is a directory; use upload_dir() to upload a folder, or name a file inside it"
+        )
+    if not p.is_file():
+        raise FileNotFoundError(f"local file not found: {p}")
+    return p
+
+
+def _require_absolute(remote: str) -> str:
+    """Validate the in-sandbox path.
+
+    Absolute only: a relative path would have to resolve against something, and the
+    candidates (the exec working directory, ``HOME``, the extracted code directory)
+    differ from each other and from what the caller's shell would pick, so the
+    ambiguity is refused rather than guessed. Note there is no ``/workspace`` in the
+    sandbox image, and ``/tmp`` is RAM-backed -- its contents do not survive the
+    sandbox being recreated.
+    """
+    if not remote:
+        raise SandboxError("remote path is required")
+    if not remote.startswith("/"):
+        raise SandboxError(
+            f"remote path must be absolute, got {remote!r}. There is no default "
+            "directory to resolve it against."
+        )
+    return remote
+
+
+def _join(dest_root: str, rel: str) -> str:
+    """Join an in-sandbox directory to a path below it, without a doubled slash.
+
+    ``dest_root`` is already normalised to have no trailing slash *except* when it is
+    the root itself, where f-string concatenation produced ``//a.txt``. POSIX
+    resolves that to the right file, but it went out on the wire and into anything
+    that compares the paths back.
+    """
+    return f"{dest_root.rstrip('/')}/{rel}"
+
+
+def _is_bad_gateway(exc: SandboxTransportError) -> bool:
+    """True for the transient 502 the ``/files`` route retries.
+
+    The container's internal HTTP server can briefly answer 502 right after
+    ``create()`` returns "ready" -- the Snowflake controller reports ready before
+    the in-container app server is. The transport surfaces that as
+    ``SandboxTransportError("HTTP 502: ...")``, so the retryable condition is
+    recognised from the message rather than a status attribute the error does not
+    carry. Shared by every byte-transfer function (and the sync twin imports it) so
+    "what counts as a retryable 502" is decided in exactly one place.
+    """
+    text = str(exc)
+    return "502" in text or "Bad Gateway" in text
+
+
+async def _retry_502(attempt: Callable[[], Awaitable[_T]]) -> _T:
+    """Run ``attempt`` up to `_FILES_RETRY_ATTEMPTS` times, retrying only a transient
+    502 Bad Gateway (`_is_bad_gateway`) with exponential backoff.
+
+    Anything else -- a non-502 transport error, a validation error, a cancellation --
+    propagates immediately; once the budget is spent the final 502 is re-raised. The
+    ``attempt`` callable owns its own per-try setup and teardown (reopening the source
+    file, allocating a fresh read buffer), so a retried transfer re-frames its request
+    body rather than re-sending a half-consumed stream.
+    """
+    for attempt_no in range(_FILES_RETRY_ATTEMPTS):
+        try:
+            return await attempt()
+        except SandboxTransportError as exc:
+            # A non-502, or the last attempt's 502, is terminal: re-raise as-is.
+            if not _is_bad_gateway(exc) or attempt_no == _FILES_RETRY_ATTEMPTS - 1:
+                raise
+            await asyncio.sleep(_FILES_RETRY_BASE_DELAY * (2**attempt_no))
+    raise AssertionError(
+        "unreachable: the loop returns or raises every iteration"
+    )  # pragma: no cover
+
+
+def _prune_empty_dirs(leaf: Path, stop: Path) -> None:
+    """Remove ``leaf`` and its parents up to (but not including) ``stop``, innermost
+    first, stopping at the first directory that is non-empty or already gone.
+
+    Undoes the ``dst.parent.mkdir(parents=True)`` a failed `download_file` made:
+    ``mkdir`` can create several levels, so removing only ``dst.parent`` would strand
+    the rest of a tree the caller never had. ``stop`` is the topmost directory that
+    already existed and is never removed. A directory that now holds other content is
+    left intact -- the guard is "remove only what this call created and still owns".
+    """
+    pruned = leaf
+    while pruned != stop:
+        try:
+            pruned.rmdir()
+        except OSError:
+            break  # non-empty, or already gone: stop and leave what remains
+        pruned = pruned.parent
+
+
+async def _files_request_with_retry(
+    t: Transport,
+    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 (see `_retry_502`)."""
+
+    async def _attempt() -> None:
+        await t.request(
+            method, endpoint, params=params, content=content, extra_headers=extra_headers
+        )
+
+    await _retry_502(_attempt)
+
+
+async def upload_file(
+    sandbox: AsyncSandbox,
+    local: str | Path,
+    remote: str,
+    *,
+    transport: Transport | None = None,
+) -> None:
+    """Upload a local file into the sandbox at ``remote``, which must be absolute.
+
+    Above `MAX_FILE_BYTES` use a stage mount: Snowflake's REST front end caps
+    request bodies below any interesting file size, so this path is for source
+    files, configs and small artifacts rather than bulk data.
+
+    Streamed from disk, not buffered. httpx reads an open file's size and frames
+    the request with ``Content-Length`` -- required, because the sandbox container's
+    server cannot decode a chunked body -- while still sending it incrementally. The
+    file is reopened per attempt by the transport's rewind, so an idempotent retry
+    re-sends it byte-for-byte.
+    """
+    if not sandbox.id:
+        raise SandboxError("sandbox must be created before uploading files")
+    src = _resolve_local_read(local)
+    size = await asyncio.to_thread(lambda: 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)
+
+    async 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:
+            await t.request(
+                "PUT",
+                f"containers/{sandbox.id}/files",
+                params={"path": dest},
+                content=body,
+                extra_headers={"Content-Type": "application/octet-stream"},
+            )
+
+    await _retry_502(_attempt)
+
+
+async def download_file(
+    sandbox: AsyncSandbox,
+    remote: str,
+    local: str | Path,
+    *,
+    transport: Transport | None = None,
+) -> None:
+    """Download ``remote`` (absolute) from the sandbox to ``local``.
+
+    Written to a temp file beside the destination and renamed on success, so a
+    failure part-way leaves neither a truncated file where a whole one is expected
+    nor an empty directory tree the caller did not have before.
+    """
+    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)
+
+    # The topmost directory that already existed; everything below it up to
+    # dst.parent is what this call is about to create, and all this call may
+    # remove on failure.
+    highest_existing = dst.parent
+    while not highest_existing.exists():
+        highest_existing = highest_existing.parent
+    await asyncio.to_thread(lambda: 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}")
+
+    async 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 = await t.download(
+                f"containers/{sandbox.id}/files", out.write, params={"path": path}
+            )
+            # The container answers a rejection with a JSON body under a non-2xx,
+            # which the transport has already raised on. A 2xx carrying JSON means
+            # something upstream substituted its own response, and writing that to
+            # disk as file content is how a "successful" download ends up being an
+            # error message -- so it is refused rather than delivered.
+            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')})"
+                )
+        await asyncio.to_thread(os.replace, tmp, dst)
+
+    try:
+        await _retry_502(_attempt)
+    except BaseException:
+        # One teardown for every failure -- a non-502, an exhausted retry budget, a
+        # JSON body, or a cancellation. 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
+
+
+async def stage_put(
+    sandbox: AsyncSandbox,
+    local: str,
+    stage_path: str,
+    *,
+    transport: Transport | None = None,
+) -> None:
+    """Not implemented — the ``stage/put`` route returns 501 on every backend.
+
+    *Would* server-side PUT from ``local`` (inside the sandbox) to ``stage_path``
+    (a Snowflake stage path like ``@my_stage/file.parquet``) without bytes
+    traversing the caller's memory. No backend serves this yet; mount a stage with
+    `StageMount` and write through the mount instead.
+
+    Raises:
+        SandboxNotImplementedError: always, until a backend wires the route
+            (the 501 maps to this permanent "not supported" type, not a
+            retryable transport error).
+    """
+    if not sandbox.id:
+        raise SandboxError("sandbox must be created before stage I/O")
+    t = transport or sandbox._transport
+    await t.request(
+        "POST",
+        f"containers/{sandbox.id}/stage/put",
+        json_body={"local": local, "stage": stage_path},
+    )
+
+
+async def stage_get(
+    sandbox: AsyncSandbox,
+    stage_path: str,
+    local: str,
+    *,
+    transport: Transport | None = None,
+) -> None:
+    """Not implemented — the ``stage/get`` route returns 501 on every backend.
+
+    *Would* server-side GET from ``stage_path`` to ``local`` (inside the sandbox).
+    No backend serves this yet; mount a stage with `StageMount` and read through
+    the mount instead.
+
+    Raises:
+        SandboxNotImplementedError: always, until a backend wires the route
+            (the 501 maps to this permanent "not supported" type, not a
+            retryable transport error).
+    """
+    if not sandbox.id:
+        raise SandboxError("sandbox must be created before stage I/O")
+    t = transport or sandbox._transport
+    await 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 ----------------
+
+
+async def read_bytes(
+    sandbox: AsyncSandbox, path: str, *, transport: Transport | None = None
+) -> bytes:
+    """Return the file at ``path`` (absolute) as bytes.
+
+    ``GET /containers/{id}/files?path=`` streamed into memory. A file larger than
+    ``MAX_FILE_BYTES`` raises `SandboxFileTooLargeError` mid-stream — the client
+    guard trips first, and the server enforces its own 413 as a backstop — so
+    mount a stage for bulk data. A missing file surfaces as `SandboxNotFoundError`.
+
+    Retries on 502 Bad Gateway: the container's internal HTTP server may not be
+    ready immediately after create() returns (Snowflake ready != app server ready).
+    """
+    if not sandbox.id:
+        raise SandboxError("sandbox must be created before reading files")
+    t = transport or sandbox._transport
+    p = _require_absolute(path)
+
+    async 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 = await t.download(f"containers/{sandbox.id}/files", _sink, params={"path": p})
+        # A 2xx carrying JSON is an upstream error page substituted for the file body,
+        # not file content -- refuse it rather than hand back an error message as bytes
+        # (mirrors download_file()).
+        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 await _retry_502(_attempt)
+
+
+async def read_text(
+    sandbox: AsyncSandbox,
+    path: str,
+    *,
+    encoding: str = "utf-8",
+    transport: Transport | None = None,
+) -> str:
+    """Return the file at ``path`` (absolute) decoded as text (``utf-8`` default)."""
+    raw = await read_bytes(sandbox, path, transport=transport)
+    return raw.decode(encoding)
+
+
+async def write_bytes(
+    sandbox: AsyncSandbox,
+    path: str,
+    data: bytes,
+    *,
+    transport: Transport | None = None,
+) -> None:
+    """Write ``data`` to ``path`` (absolute) inside the sandbox.
+
+    ``PUT /containers/{id}/files?path=`` with the bytes as the request body. Refused
+    locally when over ``MAX_FILE_BYTES`` (the same cap the route enforces) so the
+    caller gets a named error rather than an opaque 413 from a hop it cannot see.
+
+    Retries on 502 Bad Gateway: the container's internal HTTP server may not be
+    ready immediately after create() returns (Snowflake ready != app server ready).
+    """
+    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
+    await _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"},
+    )
+
+
+async def write_text(
+    sandbox: AsyncSandbox,
+    path: str,
+    data: str,
+    *,
+    encoding: str = "utf-8",
+    transport: Transport | None = None,
+) -> None:
+    """Write text to ``path`` (absolute), encoded ``utf-8`` by default."""
+    await write_bytes(sandbox, path, data.encode(encoding), transport=transport)
+
+
+async def upload_dir(
+    sandbox: AsyncSandbox,
+    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: Transport | None = None,
+) -> UploadPlan:
+    """Upload a local directory tree into the sandbox, returning what it did.
+
+    See ``_FilesMixin.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 = await asyncio.to_thread(
+        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.
+        await 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