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

snowflake/sandbox/secret.py

@@ -0,0 +1,244 @@
+"""Brokered Snowflake SECRETs.
+
+The real value never enters the sandbox: Snowflake injects a placeholder as an
+environment variable and the egress proxy substitutes the real value on
+outbound requests to the secret's authorized host.
+
+    from snowflake.sandbox import Secret, broker_env_var
+"""
+
+from __future__ import annotations
+
+import re
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass, field
+
+from snowflake.sandbox._env import (
+    _ENV_VAR_RE,
+    _RESERVED_ENV_PREFIX,
+    _reserved_env_reason,
+)
+from snowflake.sandbox._hosts import common_parent, host_covered_by_all
+from snowflake.sandbox.exceptions import SandboxError
+
+__all__ = [
+    "Secret",
+    "broker_env_var",
+    "validate_secret_entries",
+    "MAX_SECRETS_PER_APP",
+]
+
+
+# Snowflake's per-app cap on egress secrets.
+MAX_SECRETS_PER_APP = 16
+
+
+def broker_env_var(fqn: str, *, key: str | None = None) -> str:
+    """The environment variable a secret's placeholder lands in without `env_var`.
+
+    Non-alphanumeric characters become underscores and the name is uppercased, so
+    `MYDB.SECRETS.GH_TOKEN` becomes `CNG_SECRET_MYDB_SECRETS_GH_TOKEN`.
+
+    Password-type secrets expose two values; pass `key="password"` or `key="username"`
+    for the suffixed name. Naming the variable explicitly with
+    `Secret.from_name(..., env_var=...)` is usually clearer, which leaves this useful
+    mainly for introspection.
+    """
+    base = _RESERVED_ENV_PREFIX + re.sub(r"[^A-Za-z0-9]", "_", fqn).upper()
+    if key is None:
+        return base
+    suffix = {"password": "_PASSWORD", "username": "_USERNAME"}.get(key.lower())
+    if suffix is None:
+        raise SandboxError(f"unknown secret key {key!r} (expected 'password' or 'username')")
+    return base + suffix
+
+
+@dataclass(frozen=True)
+class Secret:
+    """A Snowflake SECRET brokered to the sandbox, referenced by name.
+
+    The real value never enters the sandbox. Snowflake fetches the secret,
+    injects a randomly generated placeholder as an environment variable, and
+    the egress proxy substitutes the real value on outbound requests to `host`.
+
+    Example:
+        secret = Secret.from_name(
+            "MYDB.SECRETS.GH_TOKEN",
+            env_var="GH_TOKEN",
+            host="github.com",
+        )
+        with Sandbox.create(secrets=[secret]) as sb:
+            sb.exec(["git", "clone", "https://github.com/acme/repo"])
+    """
+
+    fqn: str
+    env_var: str | None = None
+    hosts: tuple[str, ...] = field(default_factory=tuple)
+    # Singular convenience/alias. Folded into ``hosts`` by __post_init__, which
+    # then mirrors the canonical value back so reading either name works.
+    host: str | None = None
+
+    def __post_init__(self) -> None:
+        hosts = self.hosts
+        if self.host is not None:
+            if hosts:
+                raise SandboxError(f"secret {self.fqn!r}: pass host= or hosts=, not both")
+            hosts = (self.host,)
+        object.__setattr__(self, "hosts", tuple(hosts or ()))
+        # ``host`` stays readable for the single-host case (the common one).
+        object.__setattr__(self, "host", self.hosts[0] if len(self.hosts) == 1 else None)
+
+    def scope(self) -> str | None:
+        """The single host pattern that covers every host in ``hosts``.
+
+        The wire carries one ``allowed_host`` per entry, and each entry mints its
+        own dummy, so a list is only expressible when one of its own members
+        already covers the rest (a bare host covers its subdomains). Returns
+        ``None`` when no host is declared; raises when the list cannot collapse —
+        widening to a shared parent would grant more than was asked for.
+        """
+        if not self.hosts:
+            return None
+        if len(self.hosts) == 1:
+            return self.hosts[0]
+        for candidate in self.hosts:
+            others = [h for h in self.hosts if h != candidate]
+            if host_covered_by_all(others, candidate):
+                return candidate
+        base = self.env_var or "TOKEN"
+        options = []
+        parent = common_parent(self.hosts)
+        if parent is not None:
+            options.append(
+                f"(1) host={parent!r} — one env var, but also authorizes every other "
+                f"host under {parent}"
+            )
+            options.append(f"(2) host={'*.' + parent!r} — the same, minus the {parent} apex itself")
+        options.append(
+            f"({len(options) + 1}) one entry per host, exactly these and nothing more, "
+            f"each with its own env var: "
+            + " and ".join(
+                f"Secret.from_name({self.fqn!r}, env_var='{base}_{n}', host={h!r})"
+                for n, h in enumerate(self.hosts, 1)
+            )
+        )
+        raise SandboxError(
+            f"secret {self.fqn!r} lists hosts {list(self.hosts)} that no single one "
+            f"covers. One env var receives one dummy, and a dummy is swapped at one "
+            f"host pattern, so a list only collapses when a listed host already covers "
+            f"the rest. Pick one: " + "; ".join(options) + "."
+        )
+
+    @staticmethod
+    def from_name(
+        fqn: str,
+        *,
+        env_var: str | None = None,
+        host: str | None = None,
+        hosts: Sequence[str] | None = None,
+    ) -> Secret:
+        """Reference a Snowflake SECRET by fully-qualified name.
+
+        Parameters
+        ----------
+        fqn:
+            Fully-qualified SECRET name, ``db.schema.name``.  The caller's role
+            needs USAGE on it.
+        env_var:
+            Env var to receive the placeholder (e.g. ``"GH_TOKEN"``). Omit
+            for the derived ``CNG_SECRET_<FQN>`` name.
+        host:
+            Egress host this credential is authorized at — a bare host (covers
+            subdomains) or ``"*.suffix"`` (subdomains only).
+        hosts:
+            Several hosts, when one of them covers the rest (e.g.
+            ``["github.com", "api.github.com"]``). Mutually exclusive with
+            ``host``. For unrelated hosts, declare the SECRET once per host with
+            its own ``env_var`` — see the secrets guide.
+        """
+        return Secret(
+            fqn=fqn, env_var=env_var, host=host, hosts=tuple(hosts) if hosts is not None else ()
+        )
+
+    def resolved_env_var(self) -> str:
+        """The env var name this secret will land under in the container."""
+        return self.env_var or broker_env_var(self.fqn)
+
+    def to_wire(self) -> dict[str, str]:
+        """This secret as one ``egress.secrets`` entry.
+
+        Emits on ``is not None``, not truthiness, so an explicitly-empty
+        ``env_var``/``host`` reaches validation and is reported rather than
+        silently falling back to the platform default.
+        """
+        entry: dict[str, str] = {"fqn": self.fqn}
+        scope = self.scope()
+        if scope is not None:
+            entry["allowed_host"] = scope
+        if self.env_var is not None:
+            entry["env_var"] = self.env_var
+        return entry
+
+
+def validate_secret_entries(
+    entries: Sequence[Mapping[str, object]], *, has_eai: bool = False
+) -> None:
+    """Check secret entries against Snowflake's server-side rules.
+
+    Mirrors the platform's own egress-secret validation so a misconfiguration
+    fails here rather than when the sandbox is launched.  Raises `SandboxError`.
+    """
+    if len(entries) > MAX_SECRETS_PER_APP:
+        raise SandboxError(
+            f"{len(entries)} secrets exceeds the platform maximum of "
+            f"{MAX_SECRETS_PER_APP} per sandbox"
+        )
+
+    seen: dict[str, str] = {}
+    for entry in entries:
+        fqn = str(entry.get("fqn") or "")
+        if not fqn:
+            raise SandboxError("Secret.fqn is required")
+
+        host = str(entry.get("allowed_host") or "")
+        if not host and not has_eai:
+            raise SandboxError(
+                f"secret {fqn!r} needs host= (the egress host its credential is "
+                "authorized at). A bare host covers its subdomains, so one host= is "
+                "usually enough; host= becomes optional once an External Access "
+                "Integration can supply the hosts."
+            )
+        # The proxy matches on a bare, port-stripped hostname, so a scheme or
+        # port would silently never match and the secret would never inject.
+        if host and any(c in host for c in "/: \t"):
+            raise SandboxError(
+                f"secret {fqn!r} host {host!r} must be a bare hostname — "
+                "no scheme, port, or path (e.g. 'api.github.com', '*.github.com')"
+            )
+
+        declared = entry.get("env_var")
+        if declared is not None:
+            name = str(declared)
+            if not _ENV_VAR_RE.match(name):
+                raise SandboxError(
+                    f"secret {fqn!r} env_var {name!r} is not a legal environment variable name"
+                )
+            reserved = _reserved_env_reason(name)
+            if reserved is not None:
+                # A brokered secret whose placeholder lands in PATH / LD_PRELOAD /
+                # a proxy or TLS var / SNOWFLAKE_* could disable egress confinement,
+                # the MITM-CA trust, or redirect the sandbox's Snowflake identity —
+                # the CNG_SECRET_ prefix was the only name this used to reject.
+                raise SandboxError(f"secret {fqn!r} env_var {reserved}")
+            env_var = name
+        else:
+            env_var = broker_env_var(fqn)
+
+        # Env var names are the only uniqueness rule — two secrets MAY share a host.
+        if env_var in seen:
+            raise SandboxError(
+                f"secrets {seen[env_var]!r} and {fqn!r} both resolve to env var "
+                f"{env_var!r} — give one an explicit env_var. (Sharing a host is fine; "
+                "only the env var names must differ.)"
+            )
+        seen[env_var] = fqn

snowflake/sandbox/session_app.py

@@ -0,0 +1,244 @@
+"""``SessionApp`` — the ``@app.session`` stateful keep-alive class surface.
+
+A decorated class becomes a `SessionApp` whose ``.session(key=...)`` deploys a
+long-lived daemon. The ``enter`` decorator marks the cold-start hook on that
+class, and ``_session_runner_shim`` generates the in-container driver that
+dispatches to it.
+"""
+
+from __future__ import annotations
+
+import tempfile
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+from snowflake.sandbox._assemble import Bundle
+from snowflake.sandbox._deploy_spec import DeploySpec
+from snowflake.sandbox._runtime._protocol import _ENTER_FLAG
+from snowflake.sandbox._runtime._shims import _session_runner_shim
+from snowflake.sandbox.egress import Egress
+from snowflake.sandbox.function import FunctionSpec
+from snowflake.sandbox.image import Image
+from snowflake.sandbox.mount import StageMount
+from snowflake.sandbox.secret import Secret
+from snowflake.sandbox.types import MemoryTier
+
+if TYPE_CHECKING:
+    from snowflake.sandbox.app import App
+    from snowflake.sandbox.warm_session import Session
+
+__all__ = ["SessionApp", "SessionSpec", "enter"]
+
+
+def enter(method: Callable[..., Any]) -> Callable[..., Any]:
+    """Mark a method as the cold-start hook of an ``@app.session`` class.
+
+    The decorated method runs **once** when the session daemon cold-starts
+    (before any message is processed) — the place to build warm state
+    (``self.history = []``, open a client, etc.). It is a no-op tag: the method
+    is returned unchanged with a flag attribute set, and the generated session
+    shim looks the flag up at startup. A method named ``enter`` or ``setup`` is
+    also recognised by name as a fallback, so ``@enter`` is optional.
+
+    Usage:
+
+        @app.session(image="sandbox-base")
+        class SlackTriage:
+            @enter
+            def setup(self):
+                self.history = []
+
+            def on_message(self, msg: str) -> str:
+                self.history.append(msg)
+                return str(len(self.history))
+    """
+    try:
+        setattr(method, _ENTER_FLAG, True)
+    except AttributeError:
+        pass
+    return method
+
+
+@dataclass(frozen=True)
+class SessionSpec:
+    """All ``@app.session`` decorator metadata for a registered class.
+
+    Mirrors `FunctionSpec` but for a stateful Session daemon. ``entry``
+    is always the generated session-runner shim (there is no process-mode for a
+    Session), so it is not part of the spec.
+    """
+
+    image: Image
+    memory: MemoryTier = "4g"
+    cpu: float | None = None
+    bundle: Bundle | None = None
+    secrets: tuple[Secret, ...] = field(default_factory=tuple)
+    env: Mapping[str, str] = field(default_factory=dict)
+    egress: Egress | None = None
+    timeout_s: float = 3600.0
+    code_stage: str | None = None
+    stage_mounts: tuple[StageMount, ...] = field(default_factory=tuple)
+
+
+# Name of the generated daemon shim inside the bundle.
+_SESSION_RUNNER_FILENAME = "__session_runner__.py"
+
+
+class SessionApp:
+    """Returned by ``@app.session`` — a deployable stateful Session class.
+
+    Wraps the user's class + its `SessionSpec`. Calling
+    `session()` generates the daemon shim, bundles it alongside the user's
+    module, and deploys via `agent_session()` — returning a live
+    `Session`. ``__call__`` constructs the class locally (no sandbox),
+    so the class is still unit-testable directly.
+    """
+
+    def __init__(self, cls: type, spec: SessionSpec, app: App) -> None:
+        self._cls = cls
+        self._spec = spec
+        self._app = app
+        # Preserve the class identity for introspection / repr.
+        self.__name__ = getattr(cls, "__name__", "SessionApp")
+        self.__qualname__ = getattr(cls, "__qualname__", self.__name__)
+        self.__doc__ = getattr(cls, "__doc__", None)
+        self.__module__ = getattr(cls, "__module__", __name__)
+        self.__wrapped__ = cls
+
+    @property
+    def spec(self) -> SessionSpec:
+        return self._spec
+
+    @property
+    def app(self) -> App:
+        return self._app
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        """Construct the underlying class locally (no sandbox)."""
+        return self._cls(*args, **kwargs)
+
+    def local(self, *args: Any, **kwargs: Any) -> Any:
+        """Alias for ``__call__`` — construct the class in-process."""
+        return self._cls(*args, **kwargs)
+
+    async def session(
+        self,
+        *,
+        key: str,
+        env_overrides: Mapping[str, str] | None = None,
+        force_fresh: bool = False,
+    ) -> Session:
+        """Create-or-resume a keep-alive `Session` for this class.
+
+        Generates the daemon shim (the inline mailbox loop), bundles it with the
+        user's module + secrets via a generated ``DeploySpec``, and deploys it
+        detached through `agent_session()` keyed by *key*. Returns the live
+        `Session`; drive it with ``await s.send(msg)``.
+
+        Mirrors function-mode bundling: the user's module must be importable
+        from the bundle root. When no ``bundle=`` is given, the bundle root
+        defaults to the directory containing the class's source file so the
+        module is importable inside the container.
+        """
+        from snowflake.sandbox.warm_session import agent_session
+
+        with tempfile.TemporaryDirectory(prefix="sandbox_session_") as tmp:
+            tmp_path = Path(tmp)
+            self._write_session_shim(tmp_path)
+            session = await agent_session(
+                self.deploy_spec(extra_env=env_overrides),
+                key=key,
+                shim_dir=tmp_path,
+                force_fresh=force_fresh,
+                timeout=self._spec.timeout_s,
+            )
+        return session
+
+    def session_sync(
+        self,
+        *,
+        key: str,
+        env_overrides: Mapping[str, str] | None = None,
+        force_fresh: bool = False,
+    ) -> Session:
+        """Create-or-resume a keep-alive `Session` for this class — synchronous
+        counterpart of `session`.
+
+        Same daemon deploy as `session`, blocking instead of awaiting: generates
+        the mailbox-loop shim, bundles it with the user's module via the same
+        `deploy_spec`, and deploys it detached through `agent_session_sync()` keyed
+        by *key*. Returns the live `Session`; drive it with ``s.send_sync(msg)``.
+        Safe with no running event loop and never starts one.
+
+        Example:
+            session = SlackTriage.session_sync(key="triage")
+            reply = session.send_sync("what changed in prod yesterday?")
+        """
+        from snowflake.sandbox.warm_session import agent_session_sync
+
+        with tempfile.TemporaryDirectory(prefix="sandbox_session_") as tmp:
+            tmp_path = Path(tmp)
+            self._write_session_shim(tmp_path)
+            session = agent_session_sync(
+                self.deploy_spec(extra_env=env_overrides),
+                key=key,
+                shim_dir=tmp_path,
+                force_fresh=force_fresh,
+                timeout=self._spec.timeout_s,
+            )
+        return session
+
+    def _write_session_shim(self, dest: Path) -> None:
+        """Write the generated session-runner shim into *dest*.
+
+        The one non-I/O step both `session` and `session_sync` share: resolve the
+        user class's module + name and render the mailbox-loop driver into
+        ``dest/__session_runner__.py``. Kept as a helper so the two calling styles
+        cannot generate different shims.
+        """
+        module = getattr(self._cls, "__module__", None) or "__main__"
+        name = getattr(self._cls, "__name__", None) or "Session"
+        (dest / _SESSION_RUNNER_FILENAME).write_text(_session_runner_shim(module, name))
+
+    def deploy_spec(self, *, extra_env: Mapping[str, str] | None = None) -> DeploySpec:
+        """The `DeploySpec` for this Session's daemon.
+
+        The entry is the generated mailbox-loop shim; the bundle is the user's
+        module directory so the class is importable in the container.
+        """
+        spec = self._spec
+        bundle = spec.bundle or Bundle.from_dir(self._bundle_default_root())
+        fn_spec = FunctionSpec(
+            image=spec.image,
+            memory=spec.memory,
+            cpu=spec.cpu,
+            bundle=bundle,
+            secrets=spec.secrets,
+            egress=spec.egress,
+            env=spec.env,
+            stage_mounts=spec.stage_mounts,
+            timeout_s=spec.timeout_s,
+            entry=("python", _SESSION_RUNNER_FILENAME),
+            code_stage=spec.code_stage,
+        )
+        return fn_spec.to_deploy_spec(
+            project_name=self._app.name,
+            entry=["python", _SESSION_RUNNER_FILENAME],
+            extra_env=extra_env,
+            bundle=bundle,
+        )
+
+    def _bundle_default_root(self) -> str:
+        """Resolve the bundle root: explicit ``bundle.root`` or the class's
+        source-file directory (so the user's module is importable)."""
+        if self._spec.bundle is not None:
+            return self._spec.bundle.root
+        try:
+            import inspect
+
+            src = inspect.getsourcefile(self._cls) or inspect.getfile(self._cls)
+            return str(Path(src).resolve().parent)
+        except (TypeError, OSError):
+            return "."