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

snowflake/sandbox/__init__.py

@@ -0,0 +1,317 @@
+"""Snowflake sandboxes — a Modal-style ergonomic Python client over Snowflake.
+
+The sandbox is the primitive — create it, run code in it, reconnect by id.
+``Sandbox.create()`` is the factory (like Modal's ``Sandbox.create``); reconnect
+with ``Sandbox.connect(name)`` / ``get_sandbox_by_name``:
+
+    from snowflake.sandbox import Sandbox
+
+    # Sync API (default) — blocking calls, no async/await needed
+    with Sandbox.create() as sb:
+        result = sb.exec(["python", "main.py"])
+
+    # Async API — for concurrent workloads
+    from snowflake.sandbox import AsyncSandbox
+
+    async with await AsyncSandbox.create() as sb:
+        result = await sb.exec(["python", "main.py"])
+
+The public surface is lazy-imported (via module ``__getattr__``) so a bare
+``import snowflake.sandbox`` stays fast and pulls neither ``httpx`` nor
+``pydantic`` until you actually construct something.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from snowflake.sandbox._version import __version__
+from snowflake.sandbox.exceptions import (
+    SandboxAuthError,
+    SandboxConflictError,
+    SandboxContractWarning,
+    SandboxError,
+    SandboxExecError,
+    SandboxExecTimeoutError,
+    SandboxFileTooLargeError,
+    SandboxNotFoundError,
+    SandboxNotImplementedError,
+    SandboxNotReadyError,
+    SandboxRateLimitError,
+    SandboxTransportError,
+    SandboxValidationError,
+)
+from snowflake.sandbox.types import (
+    ExecResult,
+    FileInfo,
+    FileWatchEvent,
+    FileWatchEventType,
+    StreamLine,
+)
+
+__all__ = [
+    # package metadata
+    "__version__",
+    # core primitive
+    "Sandbox",
+    "AsyncSandbox",
+    "list_sandboxes",
+    "get_sandbox",
+    "get_sandbox_by_name",
+    "destroy_sandbox",
+    "list_sandboxes_async",
+    "get_sandbox_async",
+    "get_sandbox_by_name_async",
+    "destroy_sandbox_async",
+    "ExecResult",
+    "RunResult",
+    "StreamLine",
+    "ExecStream",
+    "SyncExecStream",
+    "Shell",
+    "SyncShell",
+    "ShellClosed",
+    "FileInfo",
+    "FileWatchEvent",
+    "FileWatchEventType",
+    # deploy-your-code layer
+    "App",
+    "Function",
+    "Image",
+    "Secret",
+    "Egress",
+    "Bundle",
+    "StageMount",
+    "McpServer",
+    "McpServerInfo",
+    "McpConnectionStatus",
+    "list_mcp_servers",
+    "mcp_server_status",
+    "mcp_token_env_var",
+    "broker_env_var",
+    "compile_egress",
+    "validate_secret_entries",
+    "validate_mcp_servers",
+    # detached execution (jobs) + deploy results
+    "Job",
+    "SyncJob",
+    "deploy_async",
+    "deploy_async_sync",
+    "DeployResult",
+    "DeployPlan",
+    # warm-daemon side
+    "session_loop",
+    "Session",
+    "agent_session",
+    "agent_session_sync",
+    "SessionApp",
+    # config / connections
+    "Config",
+    "ConnectionLike",
+    "using",
+    "current_config",
+    "resolve_config",
+    "close_connections",
+    "refresh_connection_config",
+    "get_snowflake_connection",
+    # transport lifecycle
+    "shutdown",
+    # exceptions
+    "SandboxError",
+    "SandboxAuthError",
+    "SandboxValidationError",
+    "SandboxConflictError",
+    "SandboxNotFoundError",
+    "SandboxNotReadyError",
+    "SandboxNotImplementedError",
+    "SandboxExecError",
+    "SandboxFileTooLargeError",
+    "SandboxExecTimeoutError",
+    "SandboxRateLimitError",
+    "SandboxTransportError",
+    # warnings
+    "SandboxContractWarning",
+]
+
+
+# Public attribute name -> (submodule, attribute). Resolved lazily on first
+# access so ``import snowflake.sandbox`` doesn't pull httpx / pydantic.
+_LAZY: dict[str, tuple[str, str]] = {
+    "Sandbox": ("snowflake.sandbox.sync_client", "Sandbox"),
+    "AsyncSandbox": ("snowflake.sandbox.client", "AsyncSandbox"),
+    # Module-level sandbox helpers. Sync-by-default: the flagship `Sandbox`
+    # is blocking, so the bare names resolve to the *sync* `sync_client`
+    # implementations — ``from snowflake.sandbox import get_sandbox`` returns a
+    # `Sandbox`, not an un-awaited coroutine. This is the spelling
+    # ``docs/reference/sandbox.md`` has documented all along; the table pointed at
+    # the async `client` module, which was the drift.
+    "list_sandboxes": ("snowflake.sandbox.sync_client", "list_sandboxes"),
+    "get_sandbox": ("snowflake.sandbox.sync_client", "get_sandbox"),
+    "get_sandbox_by_name": ("snowflake.sandbox.sync_client", "get_sandbox_by_name"),
+    "destroy_sandbox": ("snowflake.sandbox.sync_client", "destroy_sandbox"),
+    # Async twins, top-level so an async caller is not pushed into importing the
+    # `client` submodule by hand. The ``_async`` suffix matches this repo's
+    # convention for an async accessor whose default-named twin is sync
+    # (`Session.status` / `Session.status_async`); the attribute in `client` keeps
+    # its bare name, only the exported alias carries the suffix.
+    "list_sandboxes_async": ("snowflake.sandbox.client", "list_sandboxes"),
+    "get_sandbox_async": ("snowflake.sandbox.client", "get_sandbox"),
+    "get_sandbox_by_name_async": ("snowflake.sandbox.client", "get_sandbox_by_name"),
+    "destroy_sandbox_async": ("snowflake.sandbox.client", "destroy_sandbox"),
+    "ExecStream": ("snowflake.sandbox.exec_stream", "ExecStream"),
+    "SyncExecStream": ("snowflake.sandbox.sync_exec_stream", "SyncExecStream"),
+    # Lazy like ExecStream, and for the same reason: shell.py imports asyncio at
+    # module scope, which was 16ms of the package's 38ms import — 42% of it — spent
+    # by every caller whether or not they open a shell.
+    "Shell": ("snowflake.sandbox.shell", "Shell"),
+    "SyncShell": ("snowflake.sandbox.sync_shell", "SyncShell"),
+    "ShellClosed": ("snowflake.sandbox.shell", "ShellClosed"),
+    "RunResult": ("snowflake.sandbox.jobs", "RunResult"),
+    "App": ("snowflake.sandbox.app", "App"),
+    "Function": ("snowflake.sandbox.function", "Function"),
+    "Image": ("snowflake.sandbox.image", "Image"),
+    "Secret": ("snowflake.sandbox.secret", "Secret"),
+    "Egress": ("snowflake.sandbox.egress", "Egress"),
+    "broker_env_var": ("snowflake.sandbox.secret", "broker_env_var"),
+    "Bundle": ("snowflake.sandbox._assemble", "Bundle"),
+    "StageMount": ("snowflake.sandbox.mount", "StageMount"),
+    "McpServer": ("snowflake.sandbox.mcp", "McpServer"),
+    "McpServerInfo": ("snowflake.sandbox.mcp", "McpServerInfo"),
+    "McpConnectionStatus": ("snowflake.sandbox._mcp_status", "McpConnectionStatus"),
+    "list_mcp_servers": ("snowflake.sandbox._mcp_discovery", "list_mcp_servers"),
+    "mcp_server_status": ("snowflake.sandbox._mcp_status", "mcp_server_status"),
+    "mcp_token_env_var": ("snowflake.sandbox.mcp", "mcp_token_env_var"),
+    "compile_egress": ("snowflake.sandbox.egress", "compile_egress"),
+    "validate_secret_entries": ("snowflake.sandbox.secret", "validate_secret_entries"),
+    "validate_mcp_servers": ("snowflake.sandbox.mcp", "validate_mcp_servers"),
+    # Detached-execution handles/results a public method hands back, exported so a
+    # caller can annotate them (parity with `RunResult`, which `sb.wait()` returns).
+    "Job": ("snowflake.sandbox.jobs", "Job"),
+    "SyncJob": ("snowflake.sandbox.jobs", "SyncJob"),
+    "deploy_async": ("snowflake.sandbox.jobs", "deploy_async"),
+    "deploy_async_sync": ("snowflake.sandbox.jobs", "deploy_async_sync"),
+    "DeployResult": ("snowflake.sandbox.deploy", "DeployResult"),
+    "DeployPlan": ("snowflake.sandbox.deploy", "DeployPlan"),
+    "session_loop": ("snowflake.sandbox.warm_session", "session_loop"),
+    "Session": ("snowflake.sandbox.warm_session", "Session"),
+    "agent_session": ("snowflake.sandbox.warm_session", "agent_session"),
+    "agent_session_sync": ("snowflake.sandbox.warm_session", "agent_session_sync"),
+    "SessionApp": ("snowflake.sandbox.session_app", "SessionApp"),
+    "Config": ("snowflake.sandbox.config", "Config"),
+    "ConnectionLike": ("snowflake.sandbox.config", "ConnectionLike"),
+    "using": ("snowflake.sandbox.config", "using"),
+    "resolve_config": ("snowflake.sandbox.config", "resolve_config"),
+    "current_config": ("snowflake.sandbox.config", "current_config"),
+    "close_connections": ("snowflake.sandbox.config", "close_connections"),
+    "refresh_connection_config": ("snowflake.sandbox.config", "refresh_connection_config"),
+    "get_snowflake_connection": ("snowflake.sandbox.connect", "get_snowflake_connection"),
+    "shutdown": ("snowflake.sandbox.lifecycle", "shutdown"),
+}
+
+
+def configure(*_args: Any, **_kwargs: Any) -> Any:
+    """Removed. Name the connection explicitly instead — see the message below.
+
+    A raising shim rather than nothing at all, because nothing at all is unhelpful in the
+    form callers actually hit. ``from snowflake.sandbox import configure`` does not surface
+    a module ``__getattr__``'s AttributeError: CPython's IMPORT_FROM swallows it and raises
+    a bare ``ImportError: cannot import name 'configure'``, so a migrating caller would get
+    no guidance at all. Keeping the name resolvable and failing at the CALL puts the message
+    exactly where their code breaks.
+
+    Deliberately absent from ``__all__``: it is not public surface, it is a signpost. It
+    restores no behaviour -- in particular not the per-role slots whose collision is #247.
+    """
+    raise NotImplementedError(
+        "snowflake.sandbox.configure() was removed: it stored credentials in a "
+        "process-wide slot, which no first-party Snowflake SDK does, and its per-role "
+        "slots silently collided (#247). Name the connection explicitly instead — "
+        'Sandbox.create(connection="my_connection") per call, `with using("my_connection"):` '
+        "per block, or connection=Config(account=..., pat=...) for raw credentials. For a "
+        "process-wide default set SNOWFLAKE_DEFAULT_CONNECTION_NAME, or "
+        "default_connection_name in ~/.snowflake/config.toml."
+    )
+
+
+def __getattr__(name: str) -> Any:
+    """Lazy-resolve public symbols so bare imports stay fast."""
+    target = _LAZY.get(name)
+    if target is None:
+        raise AttributeError(f"module 'snowflake.sandbox' has no attribute {name!r}")
+    module_name, attr_name = target
+    import importlib
+
+    module = importlib.import_module(module_name)
+    value = getattr(module, attr_name)
+    globals()[name] = value
+    return value
+
+
+def __dir__() -> list[str]:
+    return sorted(set(globals()) | set(_LAZY))
+
+
+if TYPE_CHECKING:  # pragma: no cover -- for type checkers / IDEs only
+    # Every public name in `_LAZY` is declared here so a downstream type checker
+    # (this package ships py.typed) resolves it to its real type rather than the
+    # `Any` a module-level ``__getattr__`` alone yields. Kept in lockstep with
+    # `_LAZY` by test_package_exports.py::test_every_lazy_name_is_statically_typed.
+    from snowflake.sandbox._assemble import Bundle
+    from snowflake.sandbox._mcp_discovery import list_mcp_servers
+    from snowflake.sandbox._mcp_status import McpConnectionStatus, mcp_server_status
+    from snowflake.sandbox.app import App
+    from snowflake.sandbox.client import AsyncSandbox
+    from snowflake.sandbox.client import (
+        destroy_sandbox as destroy_sandbox_async,
+    )
+    from snowflake.sandbox.client import (
+        get_sandbox as get_sandbox_async,
+    )
+    from snowflake.sandbox.client import (
+        get_sandbox_by_name as get_sandbox_by_name_async,
+    )
+    from snowflake.sandbox.client import (
+        list_sandboxes as list_sandboxes_async,
+    )
+    from snowflake.sandbox.config import (
+        Config,
+        ConnectionLike,
+        close_connections,
+        current_config,
+        refresh_connection_config,
+        resolve_config,
+        using,
+    )
+    from snowflake.sandbox.connect import get_snowflake_connection
+    from snowflake.sandbox.deploy import DeployPlan, DeployResult
+    from snowflake.sandbox.egress import Egress, compile_egress
+    from snowflake.sandbox.exec_stream import ExecStream
+    from snowflake.sandbox.function import Function
+    from snowflake.sandbox.image import Image
+    from snowflake.sandbox.jobs import Job, RunResult, SyncJob, deploy_async, deploy_async_sync
+    from snowflake.sandbox.lifecycle import shutdown
+    from snowflake.sandbox.mcp import (
+        McpServer,
+        McpServerInfo,
+        mcp_token_env_var,
+        validate_mcp_servers,
+    )
+    from snowflake.sandbox.mount import StageMount
+    from snowflake.sandbox.secret import Secret, broker_env_var, validate_secret_entries
+    from snowflake.sandbox.session_app import SessionApp
+    from snowflake.sandbox.shell import Shell, ShellClosed
+    from snowflake.sandbox.sync_client import (
+        Sandbox,
+        destroy_sandbox,
+        get_sandbox,
+        get_sandbox_by_name,
+        list_sandboxes,
+    )
+    from snowflake.sandbox.sync_exec_stream import SyncExecStream
+    from snowflake.sandbox.sync_shell import SyncShell
+    from snowflake.sandbox.warm_session import (
+        Session,
+        agent_session,
+        agent_session_sync,
+        session_loop,
+    )

snowflake/sandbox/__main__.py

@@ -0,0 +1,225 @@
+"""``python -m snowflake.sandbox`` — attach a terminal to a sandbox.
+
+    python -m snowflake.sandbox ssh                 # pick from a list
+    python -m snowflake.sandbox ssh cntr_9f2a1b     # straight in
+    python -m snowflake.sandbox list
+
+The user-facing entry point is `snow sandbox ssh`, in this SDK's snowflake-cli
+extension; that command calls straight into `main()` here so the two cannot drift.
+This module form stays useful without the `cli` extra installed.
+
+Named `ssh` rather than `shell` for the verb people already reach for. It is not
+the SSH protocol -- no server, no port 22, no keys, no scp -- just a pty in the
+sandbox over the REST path. The library API keeps `shell()`/`Shell`, because that
+object really is a pty shell session and naming it `ssh` would claim a protocol it
+does not speak.
+
+Stdlib only, deliberately — the SDK's `cli` extra pulls in snowflake-cli, and a
+terminal should not need it.
+"""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import os
+import sys
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from snowflake.sandbox.client import AsyncSandbox
+
+# Statuses worth attaching to. A pending sandbox has no shell yet and a dead one
+# never will, so offering either in a picker is offering a guaranteed failure.
+_ATTACHABLE = frozenset({"ready"})
+
+
+def main(argv: list[str] | None = None) -> int:
+    """Run one ``python -m snowflake.sandbox`` invocation and return its exit code.
+
+    Also the entry point the ``snow sandbox ssh`` CLI extension calls, so the two
+    front ends cannot drift.
+    """
+    parser = argparse.ArgumentParser(
+        prog="python -m snowflake.sandbox",
+        description="Attach a terminal to a Snowflake sandbox.",
+    )
+    sub = parser.add_subparsers(dest="command", required=True)
+
+    shell = sub.add_parser("ssh", help="open an interactive terminal in a sandbox")
+    shell.add_argument(
+        "sandbox_id",
+        nargs="?",
+        help="sandbox to attach to; omit to choose from a list",
+    )
+    shell.add_argument("--cwd", help="working directory for the shell")
+    shell.add_argument(
+        "--no-sanitize",
+        action="store_true",
+        help="pass the sandbox's output through byte-exact, including clipboard-write "
+        "sequences (see Shell.attach)",
+    )
+
+    sub.add_parser("list", help="list sandboxes")
+
+    args = parser.parse_args(argv)
+    try:
+        if args.command == "list":
+            return asyncio.run(_cmd_list())
+        return asyncio.run(_cmd_ssh(args))
+    except KeyboardInterrupt:
+        # 130 is the shell convention for SIGINT, and this command is expected to
+        # be used from a shell.
+        return 130
+
+
+async def _cmd_list() -> int:
+    rows = await _fetch()
+    if not rows:
+        print("No sandboxes.", file=sys.stderr)
+        return 0
+    print(_render(rows))
+    return 0
+
+
+async def _cmd_ssh(args: argparse.Namespace) -> int:
+    from snowflake.sandbox.client import get_sandbox
+
+    if args.sandbox_id:
+        sandbox = await get_sandbox(args.sandbox_id)
+    else:
+        chosen = await _choose()
+        if chosen is None:
+            return 1
+        sandbox = chosen
+
+    size = _terminal_size()
+    sh = await sandbox.shell(rows=size[0], cols=size[1], cwd=args.cwd)
+    print(
+        f"Connected to {sandbox.name or sandbox.id} (session {sh.session_id}). "
+        "Exit the shell to disconnect.",
+        file=sys.stderr,
+    )
+    try:
+        code = await sh.attach(sanitize=not args.no_sanitize)
+    finally:
+        await sh.close()
+    # Exit with the shell's own code so `$?` means what a user expects; None (a
+    # code we could not read) is reported as failure rather than silent success.
+    return 0 if code == 0 else (code if isinstance(code, int) and code > 0 else 1)
+
+
+async def _fetch() -> list[AsyncSandbox]:
+    from snowflake.sandbox.client import list_sandboxes
+
+    return [s async for s in list_sandboxes()]
+
+
+async def _choose() -> AsyncSandbox | None:
+    """List sandboxes and let the caller pick one."""
+    rows = await _fetch()
+    if not rows:
+        print(
+            "No sandboxes to attach to. Create one first, or pass an id.",
+            file=sys.stderr,
+        )
+        return None
+
+    attachable = [s for s in rows if s.status in _ATTACHABLE]
+    if not attachable:
+        print("No sandbox is ready to attach to:\n", file=sys.stderr)
+        print(_render(rows), file=sys.stderr)
+        return None
+
+    if len(attachable) == 1:
+        # Nothing to choose between. Say which one, so the user is never left
+        # wondering what they just attached to.
+        only = attachable[0]
+        print(f"One sandbox available: {only.name or only.id}", file=sys.stderr)
+        return only
+
+    if not sys.stdin.isatty():
+        # A prompt with nothing to read from would hang. Print the choices and
+        # let the caller re-run with an id.
+        print(
+            "Several sandboxes are available; pass one as an argument "
+            "(stdin is not a terminal, so there is nothing to prompt):\n",
+            file=sys.stderr,
+        )
+        print(_render(attachable), file=sys.stderr)
+        return None
+
+    print(_render(attachable, numbered=True), file=sys.stderr)
+    return _prompt(attachable)
+
+
+def _prompt(rows: list[AsyncSandbox]) -> AsyncSandbox | None:
+    """Read a selection. Accepts an index, a (unique) id prefix, or an exact name."""
+    while True:
+        try:
+            raw = input(f"Attach to [1-{len(rows)}, an id, or a name, q to quit]: ").strip()
+        except EOFError:
+            return None
+        if raw in ("q", "quit", ""):
+            return None
+        if raw.isdigit():
+            n = int(raw)
+            if 1 <= n <= len(rows):
+                return rows[n - 1]
+            print(f"  {n} is out of range.", file=sys.stderr)
+            continue
+        # An id, or enough of one to be unambiguous -- ids are long hex and
+        # nobody wants to type all of it. An exact name match too: _render leads
+        # with NAME and this is the picker, so the thing people read off the
+        # screen has to be the thing they can type back in. Names aren't unique
+        # (see get_sandbox_by_name), so a duplicate just falls through to the
+        # "matches N" branch below like an ambiguous id prefix does.
+        needle = raw.casefold()
+        matches = [
+            s
+            for s in rows
+            if s.id == raw or s.id.startswith(raw) or (s.name or "").casefold() == needle
+        ]
+        if len(matches) == 1:
+            return matches[0]
+        if not matches:
+            print(f"  no sandbox matches {raw!r}.", file=sys.stderr)
+        else:
+            print(
+                f"  {raw!r} matches {len(matches)}: " + ", ".join(s.name or s.id for s in matches),
+                file=sys.stderr,
+            )
+
+
+def _render(rows: list[AsyncSandbox], *, numbered: bool = False) -> str:
+    """A plain aligned table. No colour and no dependency: this output is as
+    likely to be piped as read."""
+    header = ["NAME", "ID", "STATUS", "IMAGE", "MEMORY"]
+    body = [
+        [s.name or "-", s.id, str(s.status), s.image or "-", str(s.memory or "-")] for s in rows
+    ]
+    if numbered:
+        header = ["#", *header]
+        body = [[str(i), *r] for i, r in enumerate(body, 1)]
+
+    widths = [max(len(h), *(len(r[i]) for r in body)) for i, h in enumerate(header)]
+    lines = ["  ".join(h.ljust(w) for h, w in zip(header, widths, strict=True)).rstrip()]
+    lines += ["  ".join(c.ljust(w) for c, w in zip(r, widths, strict=True)).rstrip() for r in body]
+    return "\n".join(lines)
+
+
+def _terminal_size() -> tuple[int, int]:
+    """(rows, cols) of the real terminal, so the remote pty starts the right size.
+
+    Creating the session at 24x80 and resizing after would make every curses app
+    redraw on connect; asking for the right size up front avoids that.
+    """
+    try:
+        size = os.get_terminal_size(sys.stdout.fileno())
+    except OSError:
+        return 24, 80
+    return size.lines or 24, size.columns or 80
+
+
+if __name__ == "__main__":
+    sys.exit(main())