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

snowflake/cli_sandbox/__init__.py

@@ -0,0 +1,13 @@
+"""The ``snow sandbox`` command group, as a Snowflake CLI plugin.
+
+Installed with the ``cli`` extra and discovered by the CLI through the
+``snowflake.cli.plugin.command`` entry point in `plugin_spec`. It is a thin front
+end: `commands` parses argv and formats output, `_adapter` borrows the CLI's own
+resolved connection for credentials, and everything underneath is the ordinary
+``snowflake.sandbox`` API — so the CLI and the library cannot drift apart in
+behaviour.
+
+Nothing here is importable public API; depend on ``snowflake.sandbox`` instead.
+"""
+
+from __future__ import annotations

snowflake/cli_sandbox/_adapter.py

@@ -0,0 +1,170 @@
+"""Borrow the ``snow`` CLI's own connection for this ``snow sandbox`` process.
+
+When a ``snow sandbox`` command needs Snowflake credentials, we let the CLI
+resolve and open the connection with *its* logic — the global ``--connection`` /
+``-c`` / ``--account`` / ``--user`` flags, ``connections.toml`` plus the CLI's own
+env overrides, and every auth method (password, key-pair, OAuth, external
+browser). We then adopt the resulting live connection object as this process's
+credential authority (REST bearer + stage PUT run on it).
+
+The resolved `Config` is held **here**, in the plugin, and every SDK call the
+plugin makes passes it explicitly as ``connection=`` (see `sdk_connection`). A
+process-wide default is legitimate at this layer — a CLI invocation *is* a process —
+and keeping it here is what lets the library itself stay free of ambient credential
+state. `use_snow_connection` is therefore called once per command, before any SDK
+call; `sdk_connection` is what the call sites read.
+
+So the plugin does no connection handling of its own: it is CLI parity by
+construction, and the SDK's own resolver is used only outside the CLI (the
+programmatic path) — or here, when this module resolved nothing and
+`sdk_connection` returns ``None``. An explicit token/PAT already in the environment —
+a bare ``SNOWFLAKE_PAT``, or the ``SNOWFLAKE_ACCOUNT_URL`` + ``SNOWFLAKE_TOKEN`` that
+Snowflake injects into a running sandbox — short-circuits this when no connection was
+named, so the CLI stays usable without a configured connection or a live SQL session.
+But an explicit ``-c``/``--connection`` wins over that env token: naming a connection is
+taken as meaning it, so the flag is honoured rather than dropped.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from snowflake.sandbox.config import Config, ConnectionLike
+
+# The connection this ``snow sandbox`` process runs against, resolved once by
+# `use_snow_connection` and passed explicitly as ``connection=`` on every SDK call the
+# plugin makes. ``None`` means "nothing adopted" — either the user opted out with
+# ``--no-snow-connection`` or an env token short-circuited us — and is forwarded as
+# ``connection=None``, which is exactly the SDK's own env/``connections.toml``
+# resolution.
+#
+# Deliberately NOT a ContextVar: a CLI invocation is the whole process, and a scope
+# that is never exited leaks process-wide (it contaminated 62 unrelated tests when
+# tried). Deliberately NOT in the SDK: that is the ambient credential state this
+# package exists to keep out of the library.
+_connection: Config | None = None
+
+
+def _env_token_var() -> str | None:
+    """Name of the env token that short-circuits the CLI connection, or ``None``.
+
+    ``SNOWFLAKE_PAT`` takes precedence over ``SNOWFLAKE_TOKEN`` (same order
+    `_make_connection`/`_from_env` read them), so report whichever wins.
+    """
+    import os
+
+    if os.environ.get("SNOWFLAKE_PAT"):
+        return "SNOWFLAKE_PAT"
+    if os.environ.get("SNOWFLAKE_TOKEN"):
+        return "SNOWFLAKE_TOKEN"
+    return None
+
+
+def _env_has_token() -> bool:
+    return _env_token_var() is not None
+
+
+def use_snow_connection(connection: str | None = None) -> None:
+    """Adopt the active ``snow`` CLI connection for this process.
+
+    No-op when an explicit token/PAT is already in the environment (bare PAT, or
+    Snowflake-injected inside a sandbox) — that credential wins and needs no live
+    session, and `sdk_connection` stays ``None`` so SDK calls resolve it from the
+    environment themselves. Otherwise open the CLI's connection and hold its resolved
+    `Config`, which `sdk_connection` then hands to every SDK call.
+
+    ``connection`` names the connection to use. These commands run with
+    ``requires_connection=False`` (so the CLI does not inject its own global
+    ``-c/--connection``); the plugin exposes its own ``--connection`` and resolves
+    it here. When omitted, the default connection
+    (``SNOWFLAKE_DEFAULT_CONNECTION_NAME`` or the ``default_connection_name`` in
+    config.toml) is used — the same connection the CLI would pick.
+
+    An explicitly named connection wins over an ambient env token: a user who
+    passes ``-c/--connection`` is naming the account they mean, so the flag is
+    honoured even when ``SNOWFLAKE_PAT``/``SNOWFLAKE_TOKEN`` is set. The env token
+    only short-circuits when no connection was named — the common in-sandbox case
+    where Snowflake injects the credential and there is no configured connection.
+    """
+    from snowflake.cli.api.cli_global_context import fork_cli_context, get_cli_context
+
+    if connection is not None:
+        # Resolve the NAMED connection with the CLI's own machinery, by forking the
+        # context with its connection_name -- not with a parser of our own. Every
+        # auth method the CLI supports (key-pair incl. inline private_key, every
+        # OAuth flow, MFA, SSO/Okta authenticator URLs, session/master token,
+        # workload identity) and every endpoint field (host, port, protocol, region,
+        # an underscore account identifier) is then handled by the CLI, because it is
+        # the CLI doing the handling. This runs regardless of any env token: an
+        # explicit -c beats the ambient credential.
+        with fork_cli_context(connection_overrides={"connection_name": connection}):
+            _adopt_cli_connection(get_cli_context().connection)
+        return
+
+    if _env_token_var() is not None:
+        # No connection named, so the ambient token/PAT (a bare PAT, or the
+        # host + token Snowflake injects inside a sandbox) wins and needs no live
+        # session: the SDK reads it from the environment directly.
+        return
+
+    _adopt_cli_connection(get_cli_context().connection)
+
+
+def sdk_connection() -> Config | None:
+    """The connection every SDK call in this process must be made with.
+
+    Pass the result as ``connection=`` — ``list_sandboxes(connection=sdk_connection())``
+    — on every SDK entry point the plugin calls. ``None`` is a valid, meaningful value:
+    it means nothing was adopted (``--no-snow-connection``, or an env token
+    short-circuited `use_snow_connection`), and ``connection=None`` is the SDK's own
+    env/``connections.toml`` resolution.
+
+    Read it at the call site rather than threading it down: `sdk_connection` is the one
+    place the answer lives, so a command that resolved a connection cannot end up
+    making a call against a different one.
+    """
+    return _connection
+
+
+def reset_sdk_connection() -> None:
+    """Forget the adopted connection.
+
+    A real ``snow sandbox`` invocation never needs this — the process exits. It exists
+    so tests (and anything driving the plugin in-process) can get a clean slate instead
+    of inheriting the previous command's credentials.
+    """
+    global _connection
+    _connection = None
+
+
+def _adopt_cli_connection(conn: ConnectionLike) -> None:
+    """Adopt the CLI-built connection as this process's credential authority.
+
+    Resolves it ONCE, here, into the `Config` the plugin then hands to every SDK call.
+    The live connection object is carried on that Config: no credential *derivation*
+    happens: the SDK reads what it needs off it (its REST bearer ``rest.token``, read
+    live so a connector renewal is picked up, and its endpoint ``host``/``port``) and
+    reuses the SAME connection for a later stage PUT. So every auth method the CLI
+    resolved (password, key-pair, OAuth, external browser, session/master token) just
+    works and nothing re-authenticates with a mis-derived scheme.
+
+    Resolving once rather than per call also keeps the SDK's transport pool to a single
+    entry for the process: it keys a caller-supplied connection on the identity of the
+    session object, which this Config holds.
+    """
+    global _connection
+
+    from snowflake.sandbox.config import resolve_config
+    from snowflake.sandbox.exceptions import SandboxError
+
+    token = getattr(getattr(conn, "rest", None), "token", None)
+    host = getattr(conn, "host", None)
+    if not token or not host:
+        raise SandboxError(
+            "the active snow connection did not yield a session token; check the "
+            "connection (`snow connection test`) or set SNOWFLAKE_PAT."
+        )
+    # The Config holds the connection (keeping a reference alive) and nothing closes
+    # it — the CLI owns its lifetime.
+    _connection = resolve_config(connection=conn)

snowflake/cli_sandbox/_common.py

@@ -0,0 +1,77 @@
+"""Shared leaves for the ``snow sandbox`` command modules.
+
+Help-text constants and the small helpers several commands read
+(``_parse_env``, ``_apply_connection``, ``_resolve_sandbox_id``). They live here —
+a dependency-free leaf importable by ``commands`` and the split-out command
+modules (``_shell_command``, ``_get_command``) alike — so those modules never
+have to import each other, keeping the CLI package acyclic.
+"""
+
+from __future__ import annotations
+
+import typer
+
+_ENV_HELP = "Runtime env var passed into the container as KEY=VAL; repeatable."
+_USE_SNOW_CONN_HELP = (
+    "Use the active snow connection's credentials (default). --no-snow-connection "
+    "falls back to SDK env/config resolution (SNOWFLAKE_PAT, ~/.snowflake). An "
+    "explicit SNOWFLAKE_PAT/TOKEN in the environment is used either way."
+)
+_CONN_HELP = (
+    "Named connection from ~/.snowflake to use (defaults to your default "
+    "connection). Provided by the plugin because `snow sandbox` commands do not "
+    "take the CLI's global -c."
+)
+
+
+def _parse_env(value: str) -> tuple[str, str]:
+    if "=" not in value:
+        raise typer.BadParameter(f"expected KEY=VALUE, got {value!r}")
+    k, _, v = value.partition("=")
+    if not k:
+        raise typer.BadParameter(f"empty key in {value!r}")
+    return k, v
+
+
+def _apply_connection(
+    use_snow_connection: bool,
+    connection: str | None = None,
+) -> None:
+    """Resolve the active snow connection for this process.
+
+    Does nothing when the user opted out with --no-snow-connection, and is a no-op
+    when an explicit SNOWFLAKE_PAT/TOKEN is set (the adapter short-circuits), so it
+    is safe to call on every credentialed command.
+
+    ``connection`` names the connection to use — the plugin's own --connection/-c,
+    since the CLI does not inject one on these requires_connection=False commands.
+    Omitted means the default connection.
+
+    Stores nothing in the SDK: it resolves into the plugin's own
+    `_adapter.sdk_connection`, which every SDK call downstream passes as
+    ``connection=``. Call it before the first SDK call in a command.
+    """
+    if use_snow_connection:
+        from snowflake.cli_sandbox._adapter import use_snow_connection as _use
+
+        _use(connection=connection)
+
+
+def _resolve_sandbox_id(identifier: str) -> str:
+    """Map a user-facing sandbox reference to its underlying app id.
+
+    The name is the identity users see, so ``shell``/``logs`` accept it. An app id
+    (``cntr_`` prefix) passes through; any other value is a vanity NAME resolved via
+    the by-name lookup (unique per owner).
+
+    Call after _apply_connection: the lookup runs against the connection that resolved,
+    so the name is looked up in the account the rest of the command will act on. Doing
+    it the other way round finds a name in the default account and then operates on it
+    in another.
+    """
+    if not identifier or identifier.startswith("cntr_"):
+        return identifier
+    from snowflake.cli_sandbox._adapter import sdk_connection
+    from snowflake.sandbox.sync_client import get_sandbox_by_name
+
+    return get_sandbox_by_name(identifier, connection=sdk_connection()).id

snowflake/cli_sandbox/_egress_flags.py

@@ -0,0 +1,121 @@
+"""`--eai` / `--secret` / `--no-default-egress`: the egress surface for the CLI.
+
+The SDK has had `Egress` and `Secret` since the beginning, but nothing in the CLI
+could reach them — so an EAI-granted or credential-bearing sandbox required writing
+Python, even though `run` is the primary verb. This module is the translation layer:
+flag strings in, an `Egress` out.
+
+Spellings, chosen to match what the CLI already does elsewhere:
+
+* ``--eai NAME`` — repeatable, like ``--env``. Names an External Access Integration
+  whose network rules the platform resolves; that is the only way to reach a host
+  outside the built-in baseline.
+* ``--secret ENV_VAR=DB.SCHEMA.NAME@host[,host]`` — repeatable. The env var comes
+  first because that is the order ``--env KEY=VAL`` established, and it is the half
+  you read inside the sandbox. Omitting ``ENV_VAR=`` takes the SDK's derived
+  ``CNG_SECRET_<FQN>`` name. The ``@host`` scopes where the credential may be sent and
+  is **always required**: the SDK refuses an unscoped credential rather than let it
+  egress anywhere. (Its error text says ``host=`` becomes optional once an EAI supplies
+  the hosts — that is not implemented; an EAI does not relax the requirement, verified
+  against `compile_egress`.)
+* ``--no-default-egress`` — the closed level: Snowflake and its stages only, no
+  baseline package-manager hosts. An ``--eai`` still grants on top of it, which is
+  why this is not called ``--no-egress``.
+"""
+
+from __future__ import annotations
+
+import typer
+
+from snowflake.sandbox.egress import Egress, compile_egress
+from snowflake.sandbox.secret import Secret
+
+__all__ = ["EAI_HELP", "NO_DEFAULT_EGRESS_HELP", "SECRET_HELP", "build_egress"]
+
+EAI_HELP = (
+    "External Access Integration granting egress to hosts outside the built-in "
+    "baseline; repeatable. Your role needs USAGE on it, and its network rules must be "
+    "readable or the platform refuses the create."
+)
+
+SECRET_HELP = (
+    "Snowflake SECRET to inject, as ENV_VAR=DB.SCHEMA.NAME@host[,host]; repeatable. "
+    "The @host scopes where the credential may be sent and is required. Omit ENV_VAR= "
+    "for the derived CNG_SECRET_<FQN> name."
+)
+
+NO_DEFAULT_EGRESS_HELP = (
+    "Close the default egress: Snowflake and its stages only, dropping the baseline "
+    "package-manager hosts. An --eai still grants on top of this."
+)
+
+
+def _parse_secret(raw: str) -> Secret:
+    """One ``--secret`` value into a `Secret`, or a `typer.BadParameter`."""
+    spec = raw.strip()
+    if not spec:
+        raise typer.BadParameter("--secret needs a value, e.g. GH_TOKEN=MYDB.SECRETS.GH")
+
+    env_var: str | None = None
+    if "=" in spec:
+        # Split on the FIRST '=' only: an FQN cannot contain one, so everything after
+        # it belongs to the secret reference.
+        env_var, spec = spec.split("=", 1)
+        env_var, spec = env_var.strip(), spec.strip()
+        if not env_var:
+            raise typer.BadParameter(f"empty env var in {raw!r} (expected ENV_VAR=DB.SCHEMA.NAME)")
+
+    hosts: tuple[str, ...] = ()
+    if "@" in spec:
+        # Split on the LAST '@' so a host list is unambiguous; FQNs have no '@'.
+        spec, host_part = spec.rsplit("@", 1)
+        hosts = tuple(h.strip() for h in host_part.split(",") if h.strip())
+        if not hosts:
+            raise typer.BadParameter(f"empty host in {raw!r} (expected ...@github.com)")
+
+    fqn = spec.strip()
+    if fqn.count(".") != 2:
+        raise typer.BadParameter(
+            f"{fqn!r} is not a fully-qualified SECRET name; write DB.SCHEMA.NAME (from {raw!r})."
+        )
+    try:
+        return Secret.from_name(fqn, env_var=env_var, hosts=hosts or None)
+    except (ValueError, TypeError) as exc:  # the SDK validates host coverage
+        raise typer.BadParameter(f"{raw!r}: {exc}") from exc
+
+
+def build_egress(
+    *,
+    eai: list[str] | None,
+    secret: list[str] | None,
+    no_default_egress: bool,
+) -> tuple[Egress | None, tuple[Secret, ...]]:
+    """`(egress, secrets)` for these flags — both empty when none were given.
+
+    Two values because that is the shape `Sandbox.create()` takes: secrets are authored
+    alongside egress but are their own argument (they nest under egress only on the
+    wire, where `compile_egress` puts them).
+
+    `None` rather than an empty `Egress` when nothing was passed, so the create body
+    stays byte-identical to today's for every existing caller.
+    """
+    integrations = tuple(dict.fromkeys(name.strip() for name in (eai or []) if name.strip()))
+    secrets = tuple(_parse_secret(raw) for raw in (secret or []))
+    if not integrations and not secrets and not no_default_egress:
+        return None, ()
+
+    egress = Egress(
+        external_access_integrations=integrations,
+        **({"allow_default_egress": False} if no_default_egress else {}),
+    )
+    # Validate now, not at create time. `compile_egress` is pure and is what the client
+    # runs anyway -- but it runs *after* connecting, so a mistyped host pair otherwise
+    # costs a Snowflake round trip before saying so. These are usage errors; they should
+    # cost nothing.
+    from snowflake.sandbox.exceptions import SandboxError
+
+    try:
+        compile_egress(egress, secrets)
+    except SandboxError as exc:
+        raise typer.BadParameter(str(exc)) from exc
+    return egress, secrets

snowflake/cli_sandbox/_get_command.py

@@ -0,0 +1,109 @@
+"""``snow sandbox get`` and ``snow sandbox cancel``.
+
+Split out of ``commands.py`` unchanged: the two callbacks that address a single
+sandbox by name/id — ``get`` (show status, optionally block for the result) and
+``cancel`` (destroy it). ``commands`` imports these and registers them onto the
+shared ``app`` in the original order; the shared helpers they read come from
+``_common``.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+from typing import Any
+
+import typer
+
+from snowflake.cli_sandbox._common import (
+    _CONN_HELP,
+    _USE_SNOW_CONN_HELP,
+    _apply_connection,
+    _resolve_sandbox_id,
+)
+
+
+def get(
+    sandbox: str = typer.Argument(..., help="Sandbox name (or id) whose status to show."),
+    wait: float | None = typer.Option(
+        None,
+        "--wait",
+        "-w",
+        metavar="SECONDS",
+        help="Block up to SECONDS for a terminal result, then print it.",
+    ),
+    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:
+    """Show a sandbox's status; with --wait, block for its result.
+
+    Example:
+        snow sandbox get SANDBOX
+        snow sandbox get SANDBOX --wait 120
+    """
+    from snowflake.cli_sandbox._adapter import sdk_connection
+    from snowflake.sandbox.exceptions import SandboxError
+    from snowflake.sandbox.jobs import Job
+
+    _apply_connection(use_snow_connection, connection)
+    j = Job.from_id(_resolve_sandbox_id(sandbox), connection=sdk_connection())
+    try:
+        if wait is not None:
+            result = asyncio.run(j.get(timeout=wait))
+            typer.echo(f"sandbox:    {sandbox}")
+            typer.echo(f"status:     {result.status}")
+            typer.echo(f"exit_code:  {result.exit_code}")
+            typer.echo(f"result:     {json.dumps(result.result)}")
+        else:
+            status = asyncio.run(j.status())
+            typer.echo(f"sandbox: {sandbox}")
+            typer.echo(f"status:  {status}")
+    except SandboxError as exc:
+        typer.echo(f"error: {exc}", err=True)
+        raise typer.Exit(1) from exc
+
+
+def _cancel_command(
+    sandbox: str = typer.Argument(..., help="Sandbox name (or id) to cancel."),
+    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:
+    """Cancel a sandbox by destroying its container.
+
+    Example:
+        snow sandbox cancel SANDBOX
+    """
+    from snowflake.cli_sandbox._adapter import sdk_connection
+    from snowflake.sandbox.exceptions import SandboxError
+    from snowflake.sandbox.jobs import Job
+
+    _apply_connection(use_snow_connection, connection)
+    try:
+        asyncio.run(Job.from_id(_resolve_sandbox_id(sandbox), connection=sdk_connection()).cancel())
+    except SandboxError as exc:
+        typer.echo(f"error: {exc}", err=True)
+        raise typer.Exit(1) from exc
+    typer.echo(f"cancelled {sandbox}")
+
+
+_CANCEL_HELP = "Cancel (destroy) a running sandbox.\n\nExample:\n    snow sandbox cancel SANDBOX\n"