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

snowflake/cli_sandbox/commands.py

@@ -0,0 +1,556 @@
+"""``snow sandbox`` commands (init / run / ssh / list / get / logs / …).
+
+The command callbacks live in modules that all share this one ``app``: the
+shell/ssh block in ``_shell_command``, the get/cancel block in ``_get_command``,
+the ``run`` verb in ``_run_command``, and the rest here. The shared help-text
+constants and small helpers they read (``_apply_connection``,
+``_resolve_sandbox_id``) live in the dependency-free ``_common`` leaf, so the
+modules never import each other. This module imports the split-out callbacks and
+registers every command onto ``app`` in one place; the moved shell names are
+re-exported below so ``snowflake.cli_sandbox.commands.<name>`` keeps resolving.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import sys
+from typing import Any
+
+import typer
+from snowflake.cli.api.commands.snow_typer import SnowTyperFactory
+
+from snowflake.cli_sandbox._common import (
+    _CONN_HELP,
+    _USE_SNOW_CONN_HELP,
+    _apply_connection,
+    _resolve_sandbox_id,
+)
+from snowflake.cli_sandbox._get_command import _CANCEL_HELP, _cancel_command, get
+from snowflake.cli_sandbox._run_command import _RUN_HELP, _run_command
+from snowflake.cli_sandbox._shell_command import (  # noqa: F401  (re-export: import surface)
+    _CREATE_RETRY_ATTEMPTS,
+    _CREATE_RETRY_BACKOFF_S,
+    _SHELL_HELP,
+    _TRANSIENT_CREATE_MARKERS,
+    _attach_pty,
+    _emit_status,
+    _exec_streaming,
+    _looks_transient_create_error,
+    _shell_command,
+)
+
+app = SnowTyperFactory(
+    name="sandbox",
+    help="Run and manage Snowflake sandbox agents.",
+)
+
+
+@app.command(name="init", requires_connection=False)
+def init(
+    path: str | None = typer.Argument(None, help="Module to write (default: ./agent_app.py)."),
+    infer: bool = typer.Option(
+        False,
+        "--infer",
+        help="Use LLM to suggest Egress settings from project imports (advisory; degrades gracefully).",
+    ),
+    force: bool = typer.Option(
+        False,
+        "--force",
+        help="Overwrite an existing file.",
+    ),
+    **options: Any,
+) -> None:
+    """Scaffold a starter Python module that defines a deployable agent App.
+
+    Example:
+        snow sandbox init                  # creates ./agent_app.py
+        snow sandbox init my_agent.py      # custom name
+        snow sandbox init . --infer        # ask an LLM to suggest Egress settings
+    """
+    import os as _os
+    from pathlib import Path
+
+    from snowflake.sandbox._diagnostics import scaffold_module
+
+    out_path = Path(path).resolve() if path else Path(_os.getcwd()) / "agent_app.py"
+    if out_path.is_dir():
+        out_path = out_path / "agent_app.py"
+
+    if out_path.exists() and not force:
+        typer.echo(
+            f"{out_path} already exists. Use --force to overwrite.",
+            err=True,
+        )
+        raise typer.Exit(2)
+
+    text = scaffold_module(
+        project_name=out_path.stem.replace("_", "-"),
+        infer_dir=out_path.parent if infer else None,
+        source_dir=out_path.parent,
+    )
+    out_path.write_text(text)
+    typer.echo(str(out_path))
+
+    # The scaffold's entry= points at main.py when the directory has no runnable
+    # entry file. Without one, the obvious next step (`snow sandbox run
+    # agent_app.py --dry-run`) fails preflight ("entry script not found"). Drop a
+    # minimal runnable main.py alongside so init -> run is clean out of the box.
+    # Skip if any recognised entry file already exists (never overwrite one).
+    _entry_candidates = ("setup.sh", "main.py", "app.py", "agent.py", "__main__.py", "run.py")
+    scaffold_dir = out_path.parent
+    if not any((scaffold_dir / cand).exists() for cand in _entry_candidates):
+        main_py = scaffold_dir / "main.py"
+        main_py.write_text('print("hello from your Snowflake sandbox agent")\n')
+        typer.echo(str(main_py))
+
+
+@app.command(name="list", requires_connection=False)
+def list_(
+    use_snow_connection: bool = typer.Option(
+        True,
+        "--use-snow-connection/--no-snow-connection",
+        help=_USE_SNOW_CONN_HELP,
+    ),
+    connection: str | None = typer.Option(
+        None,
+        "--connection",
+        "-c",
+        help=_CONN_HELP,
+    ),
+    include_old_apps: bool = typer.Option(
+        False,
+        "--all",
+        "-a",
+        help=(
+            "Also show sandboxes stopped more than 24h ago and other terminal "
+            "states the backend still retains. By default only running/suspended "
+            "and recently-stopped (<=24h) sandboxes are listed."
+        ),
+    ),
+    **options: Any,
+) -> None:
+    """List your sandboxes: id, name, created_at, stopped_at, status, reason, and role.
+
+    The name is the sandbox's identity — unique per owner and the handle every other
+    command accepts (shell/logs). The app id is shown alongside it for debugging and
+    log correlation.
+
+    By default this shows the live set (running/suspended) plus sandboxes stopped
+    within the last 24h. Pass --all to also show older-stopped and other terminal
+    states the backend still retains (bounded by its retention window); it cannot
+    show sandboxes the backend has already reaped.
+
+    STATUS is the lifecycle bucket (pending/ready/dead/failed/unknown); REASON is
+    the backend's own word for the container's managed process
+    (starting/running/stopped/crashed/...), so a job still coming up reads
+    STATUS=pending REASON=starting and a live one STATUS=ready REASON=running.
+
+    REASON separates the two dead cases STATUS collapses -- a deliberate stop from a
+    crash: a sandbox you stopped reads dead/stopped, one that broke reads dead/crashed,
+    and one whose create never brought a container up reads failed/failed.
+
+    Example:
+        snow sandbox list
+    """
+    from snowflake.cli_sandbox._adapter import sdk_connection
+    from snowflake.sandbox.client import (
+        _epoch_seconds,
+        _format_created_at,
+        list_sandboxes,
+    )
+    from snowflake.sandbox.exceptions import SandboxError
+
+    _apply_connection(use_snow_connection, connection=connection)
+
+    async def _fetch() -> list[Any]:
+        rows = [
+            sb
+            async for sb in list_sandboxes(
+                include_old_apps=include_old_apps, connection=sdk_connection()
+            )
+        ]
+        # The list endpoint returns a stale create-time status (a job is stamped
+        # "starting" and never advanced), so probe each row's live status via the
+        # cheap GET /{id}/status route. Concurrent; the transport caps in-flight.
+        await asyncio.gather(*(sb.refresh_live_status() for sb in rows), return_exceptions=True)
+        return rows
+
+    try:
+        rows = asyncio.run(_fetch())
+    except SandboxError as exc:
+        typer.echo(f"error: {exc}", err=True)
+        raise typer.Exit(1) from exc
+
+    if not rows:
+        typer.echo("No sandboxes.", err=True)
+        return
+
+    # Order for the reader: rows with no readable created_at (shown as "-") first,
+    # then the rest oldest-first. `_epoch_seconds() is None` is exactly the "-"
+    # set, so the sort key agrees with the CREATED_AT column. `id` is the final
+    # tiebreaker: it orders the "-" group (which has no time to sort by) and
+    # settles equal timestamps, so the output is fully deterministic.
+    def _created_key(sb: Any) -> tuple[bool, float, str]:
+        secs = _epoch_seconds(sb.created_at)
+        return (secs is not None, secs if secs is not None else 0.0, sb.id)
+
+    rows.sort(key=_created_key)
+
+    header = ["ID", "NAME", "CREATED_AT", "STOPPED_AT", "STATUS", "REASON", "ROLE"]
+    body = [
+        [
+            sb.id,
+            sb.name or "-",
+            _format_created_at(sb.created_at),
+            # _format_created_at is a generic epoch->local-time formatter; a live
+            # sandbox has no stop time so this renders "-".
+            _format_created_at(sb.stopped_at),
+            str(sb.status),
+            sb.server_status or "-",
+            sb.role or "-",
+        ]
+        for sb in rows
+    ]
+    widths = [max(len(h), *(len(r[i]) for r in body)) for i, h in enumerate(header)]
+    typer.echo("  ".join(h.ljust(w) for h, w in zip(header, widths, strict=True)).rstrip())
+    for r in body:
+        typer.echo("  ".join(c.ljust(w) for c, w in zip(r, widths, strict=True)).rstrip())
+
+
+# `run` is the one verb for launching a workload — an image, a local file/dir, or a
+# .py App module (the target's shape is auto-detected); --detach makes it
+# long-running rather than run-to-completion. It subsumes the old `deploy` verb (the
+# App-module target is deploy's authoring path). The callback lives in _run_command;
+# register it here so the command order stays fixed alongside the others.
+app.command(name="run", requires_connection=False, help=_RUN_HELP)(_run_command)
+
+
+# `shell` is the primary verb; `ssh` is an alias (same callable, same options).
+# `ssh`'s original attach/picker behavior is preserved -- `shell` adds the ephemeral
+# run-and-teardown (--cmd / create-flag) path on top of it, and fills the one gap the
+# picker left: when there is nothing to attach to, it creates a fresh box rather than
+# dead-ending (with --new to force a fresh box regardless). The callback lives in
+# _shell_command; register it here so the command order stays fixed.
+app.command(name="shell", requires_connection=False, help=_SHELL_HELP)(_shell_command)
+app.command(
+    name="ssh",
+    requires_connection=False,
+    help="Alias for `shell`.\n\n" + _SHELL_HELP,
+)(_shell_command)
+
+
+async def _stream_process_logs(sb: Any, tail: int) -> None:
+    """Print the last `tail` lines of a container's managed-process output, then
+    follow new lines live until the process is terminal.
+
+    Backs `snow sandbox logs --follow`. stdout goes to stdout and stderr to
+    stderr, so a caller's `2>` separates them.
+    """
+    from snowflake.sandbox.log_stream import LogStream
+
+    since = 0
+    if tail > 0:
+        # Print the last N lines (both channels) first, then follow only what
+        # arrives after them -- classic `tail -f`. tail=0 follows from the start.
+        #
+        # This reads the SDK's private per-line accessor (`{ts, stream, text}`) on
+        # purpose: the whole JSON-line snapshot/parse path is a stopgap until the
+        # backend supports true log streaming, so it is deliberately NOT promoted to
+        # public API. `logs()` (public) returns flattened text and cannot split the
+        # stderr channel or give the ts watermark this needs.
+        snapshot = await sb._fetch_log_lines(tail=tail)
+        for ln in snapshot:
+            typer.echo(str(ln.get("text") or ""), err=(ln.get("stream") == "stderr"))
+        last_ts = max((int(ln.get("ts") or 0) for ln in snapshot), default=0)
+        since = last_ts + 1 if last_ts else 0
+
+    async def _drain(lines: Any, *, err: bool) -> None:
+        # LogStream is the SDK's per-channel poll iterator (source=process; ends
+        # when the process is terminal). The post-snapshot watermark keeps the tail
+        # lines above from being reprinted.
+        async for line in lines:
+            typer.echo(line, err=err)
+
+    await asyncio.gather(
+        _drain(LogStream(sb, "stdout", since_ts_ms=since), err=False),
+        _drain(LogStream(sb, "stderr", since_ts_ms=since), err=True),
+    )
+
+
+_LOGS_HELP = (
+    "Print or follow a sandbox's output.\n\n"
+    "By default this is the captured stdout/stderr of the sandbox's managed process, "
+    "which only exists for command containers (created with `command=`). Snapshot by "
+    "default; --follow tails the SAME source live until the process finishes, printing "
+    "stdout to stdout and stderr to stderr so `2>` separates them.\n\n"
+    "--file reads a file inside the sandbox instead. That is how you read a "
+    "`run --detach` local-code run: its command is backgrounded rather than being the "
+    "managed process, so its output is tee'd to a file and the managed-process log is "
+    "empty. --follow tails the file; --since does not apply to it.\n\n"
+    "Example:\n"
+    "    snow sandbox logs SANDBOX_ID\n"
+    "    snow sandbox logs SANDBOX_ID --tail 50\n"
+    "    snow sandbox logs SANDBOX_ID --follow\n"
+    "    snow sandbox logs SANDBOX_ID --file /tmp/sandbox-run.out --follow"
+)
+
+
+def _logs_command(
+    container_id: str = typer.Argument(
+        ...,
+        help="Sandbox (name or app id) whose captured output to read.",
+    ),
+    tail: int = typer.Option(
+        0,
+        "--tail",
+        metavar="N",
+        help="Print only the last N lines (0 = all). With --follow, print the last N, then stream new lines.",
+    ),
+    follow: bool = typer.Option(
+        False,
+        "--follow",
+        "-f",
+        help="Stream the output live until the process finishes (Ctrl-C to stop).",
+    ),
+    since: int = typer.Option(
+        0,
+        "--since",
+        metavar="MS",
+        help="Only lines at or after this server timestamp (Unix epoch ms). Ignored with --follow.",
+    ),
+    file: str | None = typer.Option(
+        None,
+        "--file",
+        metavar="PATH",
+        help="Read this file inside the sandbox instead of the managed-process log. "
+        "Use it for a `run --detach` local-code run, whose output is tee'd to "
+        "/tmp/sandbox-run.out. Not valid with --since.",
+    ),
+    use_snow_connection: bool = typer.Option(
+        True,
+        "--use-snow-connection/--no-snow-connection",
+        help=_USE_SNOW_CONN_HELP,
+    ),
+    connection: str | None = typer.Option(
+        None,
+        "--connection",
+        "-c",
+        help=_CONN_HELP,
+    ),
+    **options: Any,
+) -> None:
+    """Print or follow a sandbox's output — the managed-process log, or `--file`.
+
+    Backs `snow sandbox logs`. Snapshot by default; --follow tails the same source
+    live until the process finishes.
+
+    Example:
+        snow sandbox logs SANDBOX_ID --tail 50
+        snow sandbox logs SANDBOX_ID --follow
+        snow sandbox logs SANDBOX_ID --file /tmp/sandbox-run.out
+    """
+    import shlex
+
+    from snowflake.cli_sandbox._adapter import sdk_connection
+    from snowflake.sandbox.client import get_sandbox
+    from snowflake.sandbox.exceptions import SandboxError
+
+    if file is not None and since:
+        raise typer.BadParameter(
+            "--since applies to the managed-process log, not to --file (a file has no "
+            "server timestamps). Use --tail to limit the output instead."
+        )
+
+    _apply_connection(use_snow_connection, connection)
+
+    # Accept the sandbox NAME (what `list` shows) as well as the raw app id.
+    resolved_id = _resolve_sandbox_id(container_id)
+
+    async def _run() -> None:
+        sb = await get_sandbox(resolved_id, connection=sdk_connection())
+        if file is not None:
+            await _read_sandbox_file(sb, file, tail=tail, follow=follow)
+        elif follow:
+            await _stream_process_logs(sb, tail)
+        else:
+            typer.echo(await sb.logs(tail=tail, since_ts_ms=since))
+
+    async def _read_sandbox_file(sb: Any, path: str, *, tail: int, follow: bool) -> None:
+        """Print, or tail, a file inside the sandbox.
+
+        `--follow` shells out to `tail -f` over `exec_stream` because there is no
+        watch-a-file route; the snapshot path uses `read_text` so a plain read costs
+        one request and cannot be confused by a missing `tail` binary.
+        """
+        quoted = shlex.quote(path)
+        if follow:
+            await _exec_streaming(sb, f"tail -n {max(tail, 0)} -f {quoted}", None, None)
+            return
+        text = await sb.read_text(path)
+        if tail > 0:
+            text = "\n".join(text.splitlines()[-tail:])
+        typer.echo(text)
+
+    try:
+        asyncio.run(_run())
+    except KeyboardInterrupt:
+        # 130 is the shell convention for SIGINT; --follow is meant to be Ctrl-C'd.
+        raise typer.Exit(130) from None
+    except SandboxError as exc:
+        typer.echo(f"error: {exc}", err=True)
+        raise typer.Exit(1) from exc
+
+
+app.command(name="logs", requires_connection=False, help=_LOGS_HELP)(_logs_command)
+
+
+def _is_interactive() -> bool:
+    """Whether stdin is a terminal. Factored out so the --all confirmation gate is
+    testable without a real tty."""
+    return sys.stdin.isatty()
+
+
+_DELETE_HELP = """Delete (terminate) sandboxes and free their capacity.
+
+Destroys the named sandboxes, or with --all every sandbox this connection's role
+can see. Targets may be sandbox ids (cntr_...) or names. --all asks for
+confirmation first unless --yes is given, and refuses without --yes when stdin is
+not a terminal -- so a stray `--all` in a script cannot silently wipe everything.
+
+Deletion is role-scoped like every sandbox operation: you can only delete
+sandboxes owned by the role your connection is using.
+
+Example:
+    snow sandbox delete cntr_9f2a1b
+    snow sandbox delete my-box other-box
+    snow sandbox delete --all
+    snow sandbox delete --all --yes          # no prompt (scripting)
+"""
+
+
+def _resolve_delete_targets(*, all_: bool, ids: list[str], yes: bool) -> list[str]:
+    """Compute the sandbox ids to delete: every visible sandbox (``--all``, with a
+    confirmation) or the explicitly-named ones (names resolved to ids; a name that
+    no longer resolves is already gone, so it is skipped). Returns ``[]`` when there
+    is nothing to do; may ``Exit`` on a collection error or a declined confirmation."""
+    from snowflake.cli_sandbox._adapter import sdk_connection
+    from snowflake.sandbox.client import list_sandboxes
+    from snowflake.sandbox.exceptions import SandboxError, SandboxNotFoundError
+
+    if not all_:
+        # Resolve names to ids up front (a cntr_ id passes through). A name that no
+        # longer resolves is already gone -- the delete's desired end state.
+        targets: list[str] = []
+        for token in ids:
+            try:
+                targets.append(_resolve_sandbox_id(token))
+            except SandboxNotFoundError:
+                typer.echo(f"already gone {token}")
+        return targets
+
+    async def _collect() -> list[str]:
+        return [sb.id async for sb in list_sandboxes(connection=sdk_connection())]
+
+    try:
+        targets = asyncio.run(_collect())
+    except SandboxError as exc:
+        typer.echo(f"error: {exc}", err=True)
+        raise typer.Exit(1) from exc
+    if not targets:
+        typer.echo("No sandboxes to delete.", err=True)
+        return []
+    if not yes:
+        # Bulk + destructive: confirm interactively, and refuse outright in a
+        # non-interactive shell rather than read EOF as "yes".
+        if not _is_interactive():
+            typer.echo(
+                "refusing to delete --all without --yes in a non-interactive shell.",
+                err=True,
+            )
+            raise typer.Exit(2)
+        typer.confirm(f"Delete all {len(targets)} sandbox(es)?", abort=True)
+    return targets
+
+
+def _delete_command(
+    sandbox_ids: list[str] | None = typer.Argument(
+        None,
+        metavar="[SANDBOX_IDS]...",
+        help="Sandbox ids (cntr_...) or names to delete. Omit when using --all.",
+    ),
+    all_: bool = typer.Option(
+        False,
+        "--all",
+        help="Delete every sandbox this connection's role can see.",
+    ),
+    yes: bool = typer.Option(
+        False,
+        "--yes",
+        "-y",
+        help="Skip the confirmation prompt for --all.",
+    ),
+    use_snow_connection: bool = typer.Option(
+        True,
+        "--use-snow-connection/--no-snow-connection",
+        help=_USE_SNOW_CONN_HELP,
+    ),
+    connection: str | None = typer.Option(
+        None,
+        "--connection",
+        "-c",
+        help=_CONN_HELP,
+    ),
+    **options: Any,
+) -> None:
+    """Delete (terminate) one or more sandboxes, or all of them with --all."""
+    from snowflake.cli_sandbox._adapter import sdk_connection
+    from snowflake.sandbox.client import destroy_sandbox
+    from snowflake.sandbox.exceptions import SandboxError
+
+    ids = list(sandbox_ids or [])
+    if all_ and ids:
+        raise typer.BadParameter("pass either --all or explicit SANDBOX_IDS, not both")
+    if not all_ and not ids:
+        raise typer.BadParameter("provide one or more SANDBOX_IDS, or --all")
+
+    _apply_connection(use_snow_connection, connection=connection)
+
+    targets = _resolve_delete_targets(all_=all_, ids=ids, yes=yes)
+    if not targets:
+        return
+
+    async def _run() -> int:
+        ok = fail = 0
+        for sid in targets:
+            try:
+                # destroy_sandbox is idempotent -- a 404 counts as success, so an
+                # already-gone sandbox is reported deleted, not as a failure.
+                await destroy_sandbox(sid, connection=sdk_connection())
+                ok += 1
+                typer.echo(f"deleted {sid}")
+            except SandboxError as exc:
+                fail += 1
+                typer.echo(f"error: could not delete {sid}: {exc}", err=True)
+        if fail:
+            typer.echo(f"deleted {ok}, failed {fail}", err=True)
+            return 1
+        return 0
+
+    try:
+        code = asyncio.run(_run())
+    except SandboxError as exc:
+        typer.echo(f"error: {exc}", err=True)
+        raise typer.Exit(1) from exc
+    if code:
+        raise typer.Exit(code)
+
+
+app.command(name="delete", requires_connection=False, help=_DELETE_HELP)(_delete_command)
+
+
+# `get` and `cancel` each address a single sandbox by name/id -- `get` shows its
+# status (or blocks for a result with --wait), `cancel` destroys it. Both callbacks
+# live in _get_command. There is no `job` subgroup: a sandbox is managed on its own,
+# so these are plain top-level verbs alongside `list`/`logs`/`delete`.
+app.command(name="get", requires_connection=False)(get)
+app.command(name="cancel", requires_connection=False, help=_CANCEL_HELP)(_cancel_command)

snowflake/cli_sandbox/plugin_spec.py

@@ -0,0 +1,28 @@
+"""The entry point the Snowflake CLI loads this plugin through.
+
+Named by the ``snowflake.cli.plugin.command`` entry point in ``pyproject.toml``;
+the CLI calls `command_spec` during startup and mounts the returned typer app.
+Registering at the root path is what makes the group ``snow sandbox`` rather
+than a subcommand of an existing group.
+"""
+
+from __future__ import annotations
+
+from snowflake.cli.api.plugins.command import (
+    SNOWCLI_ROOT_COMMAND_PATH,
+    CommandSpec,
+    CommandType,
+    plugin_hook_impl,
+)
+
+from snowflake.cli_sandbox import commands
+
+
+@plugin_hook_impl
+def command_spec() -> CommandSpec:
+    """Describe the ``snow sandbox`` command group for the CLI's plugin loader."""
+    return CommandSpec(
+        parent_command_path=SNOWCLI_ROOT_COMMAND_PATH,
+        command_type=CommandType.COMMAND_GROUP,
+        typer_instance=commands.app.create_instance(),
+    )