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

snowflake/sandbox/_session_registry.py

@@ -0,0 +1,222 @@
+"""Connector-session registry + credential renewal — split out of ``config.py``.
+
+Lifecycle is delegated to snowflake-connector-python. The token is read LIVE
+from ``connection.rest.token`` (``Config.resolve_pat``), and the connection is
+minted with ``client_session_keep_alive=True`` so the connector heartbeats the
+session alive server-side and renews it in place — the SDK picks that up on the
+next request without any snapshot going stale.
+
+When a request still comes back with a reauth signal (the session was truly
+gone — 390111 — or the master token / OAuth grant lapsed faster than the
+heartbeat), the transport calls the reactive path below as a backstop:
+``refresh_config`` renews off the master token first (one
+``/session/token-request``, no re-login) and falls back to a full reconnect via
+the carried ``connector_kwargs``. ``refresh_credential`` is the process-wide
+entry point; it de-duplicates a concurrent burst (via ``_refresh_lock``) so the
+fleet of failing requests triggers exactly ONE re-mint and the rest adopt its
+result. Crucially it does the connector I/O WITHOUT holding ``_lock`` — that
+lock only guards the resolution memo's read/swap — so a refresh (up to a 30 s
+reconnect) never blocks ``current_config()`` on the request hot path.
+
+Every name here is re-exported from ``config`` (``from ._session_registry
+import ...``) so ``snowflake.sandbox.config.<name>`` keeps resolving. The shared
+state these touch (the ``_session_registry`` list, ``_lock``, ``_resolution_cache``,
+``_refresh_lock``) and ``Config`` / ``refresh_connection_config`` stay in
+``config`` and are imported lazily inside the functions that need them — a
+module-scope import back into ``config`` would close a cycle (``config`` imports
+this module to re-export it).
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from snowflake.sandbox.config import Config
+
+
+def _close_session_registry() -> None:
+    from snowflake.sandbox.config import _session_registry
+
+    conns = list(_session_registry)
+    _session_registry.clear()
+    for conn in conns:
+        close = getattr(conn, "close", None)
+        if close is None:
+            continue
+        try:
+            close()
+        except Exception:  # noqa: BLE001 - best-effort; a dead session is fine
+            pass
+
+
+def _close_one_session(conn: object) -> bool:
+    """Close *conn* only if the SDK opened it. Returns whether it was closed.
+
+    Ownership is already recorded by how the Config was built:
+    ``_config_from_connection`` appends the session it minted to ``_session_registry``,
+    while ``_config_from_live_connection`` deliberately does not -- a connection or
+    Snowpark ``Session`` the caller opened stays theirs. So membership in the registry IS
+    the ownership test, and a caller-supplied connection is never closed here no matter
+    what a ``close=True`` asks for.
+    """
+    from snowflake.sandbox.config import _session_registry
+
+    if conn is None:
+        return False
+    for i, known in enumerate(_session_registry):
+        if known is conn:
+            del _session_registry[i]
+            _close_connection(conn)
+            return True
+    return False
+
+
+def _renew_or_reconnect(cfg: Config) -> Config | None:
+    """Re-mint *cfg*'s token. Returns the Config to install, or ``None``.
+
+    Returns *cfg* UNCHANGED when the token was renewed in place on the existing
+    connection (the cheap master-token path — ``resolve_pat`` already reflects the
+    new token, so nothing needs replacing). Returns a NEW Config carrying a fresh
+    connection when a full reconnect was required (390111 / master dead). Returns
+    ``None`` when neither was possible, so the caller surfaces the original auth
+    error instead of retrying an identical credential.
+    """
+    conn = cfg.connection
+    if conn is not None:
+        rest = getattr(conn, "rest", None)
+        renew = getattr(rest, "_renew_session", None)
+        if callable(renew):
+            before = getattr(rest, "token", None)
+            try:
+                # Master-token round-trip; mutates rest.token in place. Raises
+                # ReauthenticationRequest for 390111/390114/390115 (non-renewable)
+                # — caught here so we fall through to a full reconnect.
+                renew()
+            except Exception:  # noqa: BLE001 - non-renewable / connector API drift
+                pass
+            else:
+                after = getattr(rest, "token", None)
+                if after and after != before:
+                    return cfg  # renewed in place; resolve_pat() sees it live
+    return _reconnect(cfg)
+
+
+def _reconnect(cfg: Config) -> Config | None:
+    """Full re-login via the carried ``connector_kwargs``; ``None`` if unavailable.
+
+    This is what naming the connection again would do — re-read the TOML and
+    authenticate with the connection's declared method — but on demand, keyed to a
+    single expired session rather than run on a timer. The old connection is closed
+    and dropped from the registry (its keep-alive heartbeat thread and sockets
+    would otherwise accumulate one-per-reconnect for the process lifetime); the
+    fresh one replaces ``cfg.connection`` so ``resolve_pat`` reads it and the next
+    renewal has a live master token again.
+    """
+    kwargs = cfg.connector_kwargs
+    if not kwargs:
+        return None
+    import dataclasses
+
+    import snowflake.connector
+
+    from snowflake.sandbox.config import _session_registry, refresh_connection_config
+
+    refresh_connection_config()  # re-point CONFIG_MANAGER at the right home first
+    conn = snowflake.connector.connect(**kwargs)
+    rest = getattr(conn, "rest", None)
+    token = getattr(rest, "token", None) if rest is not None else None
+    if not token:
+        # Do not retain a connection we cannot use; close it so its heartbeat
+        # thread does not linger.
+        _close_connection(conn)
+        return None
+    _session_registry.append(conn)  # keep alive so its token stays valid
+    _retire_connection(cfg.connection)  # close+drop the one we are replacing
+    return dataclasses.replace(cfg, pat=token, connection=conn)
+
+
+def _close_connection(conn: object) -> None:
+    close = getattr(conn, "close", None)
+    if close is None:
+        return
+    try:
+        close()
+    except Exception:  # noqa: BLE001 - best-effort; a dead session is fine
+        pass
+
+
+def _retire_connection(conn: object | None) -> None:
+    """Close *conn* and drop it from the registry (bounds reconnect growth)."""
+    from snowflake.sandbox.config import _session_registry
+
+    if conn is None:
+        return
+    try:
+        _session_registry.remove(conn)
+    except ValueError:
+        pass
+    _close_connection(conn)
+
+
+def refresh_config(cfg: Config, stale_pat: str) -> Config:
+    """Return *cfg* with its token re-minted after ``stale_pat`` stopped working.
+
+    Returns *cfg* UNCHANGED when its live token is no longer ``stale_pat`` (another
+    path already refreshed it), when it was renewed in place, or when there is
+    nothing to refresh (an env/PAT Config with no connection and no reconnect
+    kwargs). The caller detects a real reconnect by identity (``result is not cfg``)
+    and/or a changed ``resolve_pat()``.
+
+    Does the connector I/O (a token-request round-trip, or a reconnect) itself, so
+    an ``async`` caller must run it off the event loop (``asyncio.to_thread``), and
+    a caller sharing it across threads must serialize (see ``refresh_credential`` /
+    the transport's per-instance refresh lock).
+    """
+    if cfg.resolve_pat() != stale_pat:
+        return cfg
+    new = _renew_or_reconnect(cfg)
+    return new if new is not None else cfg
+
+
+def refresh_credential(stale_pat: str) -> str | None:
+    """Re-mint the AMBIENT config's token; the transport's reauth hook.
+
+    Serves the transport that resolves through `current_config` on every request (the
+    one a caller who named no connection gets). A transport bound to an explicit
+    ``connection=`` carries its own Config and refreshes it through `refresh_config`
+    under its own lock instead.
+
+    De-duplicated against the token the caller actually sent: the first of a
+    concurrent burst re-mints (holding ``_refresh_lock``, NOT ``_lock``, so config
+    readers are never blocked on the network) and installs the result; every racer
+    then sees the live token has changed and adopts it without minting again.
+    Returns the fresh token, or ``None`` when nothing could be refreshed (a config with
+    no connection/reconnect kwargs) so the transport surfaces the original auth error
+    rather than retrying an identical credential.
+    """
+    from snowflake.sandbox.config import _refresh_lock, _replace_cached_config, current_config
+
+    cur = current_config()
+    if cur.resolve_pat() != stale_pat:
+        return cur.resolve_pat()  # already refreshed elsewhere
+    with _refresh_lock:
+        # Re-read + re-check under the refresh lock: another thread may have
+        # re-minted while we waited, in which case adopt its token.
+        cur = current_config()
+        current = cur.resolve_pat()
+        if current != stale_pat:
+            return current
+        new = _renew_or_reconnect(cur)  # connector I/O — no lock held here
+        if new is None:
+            return None  # nothing to refresh — let the auth error surface
+        # Install the reconnect where the ambient config is resolved FROM, so the next
+        # request serves it instead of re-minting. A miss means the Config we renewed is
+        # no longer the one on offer — ``close_connections`` / ``refresh_connection_config``
+        # dropped the memo mid-refresh, or the ambient config came from the bare
+        # environment and was never memoised — so retire our now-orphan reconnect rather
+        # than leak it, and adopt whatever resolves now.
+        if new is not cur and not _replace_cached_config(cur, new):
+            _retire_connection(new.connection)
+            return current_config().resolve_pat()
+        return new.resolve_pat()

snowflake/sandbox/_sse.py

@@ -0,0 +1,160 @@
+"""Loop-agnostic Server-Sent-Events parsing shared by the transport twins.
+
+The parsed-frame type, the single-frame parser, the open-response frame iterator,
+and the initial-open retry policy for a stream that failed on first contact. Used by
+both the async ``_transport`` and the sync ``_sync_transport`` (whose own frame
+iterator is the blocking mirror of ``_iter_sse_frames`` here). A dependency-free
+leaf so both transports import one copy.
+"""
+
+from __future__ import annotations
+
+import json
+import random
+from collections.abc import AsyncIterator
+from typing import Any
+
+import httpx
+
+from snowflake.sandbox._retry import (
+    DEFAULT_COLD_START_DELAYS_S,
+    DEFAULT_NOT_FOUND_DELAYS_S,
+    compute_backoff,
+)
+from snowflake.sandbox._transport_errors import (
+    SandboxSSEDecodeError,
+    _rate_limit_delay,
+    _says_not_ready,
+)
+
+
+class SSEEvent:
+    """One parsed Server-Sent-Events frame."""
+
+    __slots__ = ("data", "event", "id")
+
+    def __init__(self, event: str, data: str, id: str | None = None) -> None:  # noqa: A002
+        self.event = event
+        self.data = data
+        # The frame's `id:` field, when the server sent one. Both resumable
+        # streams -- exec (reattached from `last_event_id` after the ~180s
+        # ingress GOAWAY; see client._exec_stream_frames) and shell (cut by the
+        # gateway at its stream timeout and resumed from the last id seen) --
+        # advance a cursor from this id. Optional because non-resumable frames
+        # (e.g. the `session`/`exit` bookkeeping frames) may omit it. This parser
+        # used to drop the field entirely, which made resumption impossible to
+        # implement without the server duplicating the cursor into `data`.
+        self.id = id
+
+    def json(self) -> Any:
+        """Parse ``data`` as JSON.
+
+        An **empty** ``data`` (a frame that carried no ``data:`` field) returns
+        ``None`` without attempting a parse -- distinct from the literal string
+        ``"null"``, which is valid JSON and also decodes to ``None`` but goes
+        through the parser. Malformed non-empty ``data`` raises ``SandboxError``
+        (a typed member of the hierarchy) rather than the bare
+        ``json.JSONDecodeError`` a caller doing ``except SandboxError`` would miss.
+        """
+        if self.data == "":
+            return None
+        try:
+            return json.loads(self.data)
+        except json.JSONDecodeError as exc:
+            raise SandboxSSEDecodeError(f"SSE frame data is not valid JSON: {exc}") from exc
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, SSEEvent):
+            return NotImplemented
+        return self.event == other.event and self.data == other.data and self.id == other.id
+
+    def __hash__(self) -> int:
+        return hash((self.event, self.data, self.id))
+
+    def __repr__(self) -> str:  # pragma: no cover - debug
+        return f"SSEEvent(event={self.event!r}, data={self.data!r}, id={self.id!r})"
+
+
+def parse_sse_event(buf: str) -> SSEEvent | None:
+    """Parse a single SSE event (text terminated by a blank line).
+
+    Returns ``None`` when the block carries no recognized field -- a blank block,
+    or one made only of comment/keep-alive lines (``: keep-alive``). A keep-alive
+    frame previously produced a phantom ``SSEEvent("message", "")`` that every
+    consumer had to filter (or emit a spurious blank line per heartbeat); it is
+    now dropped here.
+    """
+    if not buf.strip():
+        return None
+    event = "message"
+    data_lines: list[str] = []
+    event_id: str | None = None
+    saw_field = False
+    for line in buf.splitlines():
+        if not line or line.startswith(":"):
+            continue
+        field, _, value = line.partition(":")
+        if value.startswith(" "):
+            value = value[1:]
+        if field == "event":
+            event = value
+            saw_field = True
+        elif field == "data":
+            data_lines.append(value)
+            saw_field = True
+        elif field == "id":
+            event_id = value
+            saw_field = True
+    if not saw_field:
+        return None
+    return SSEEvent(event=event, data="\n".join(data_lines), id=event_id)
+
+
+async def _iter_sse_frames(resp: httpx.Response) -> AsyncIterator[SSEEvent]:
+    """Yield the parsed frames of an open SSE response.
+
+    Only frames terminated by a blank line are dispatched. A leftover, un-blank-
+    terminated buffer at end-of-stream is a **truncated** frame -- the stream was
+    cut mid-frame (an h2 GOAWAY at the 180s ingress cut, a 900s shell timeout) --
+    and is discarded rather than delivered as if complete. Yielding it corrupted
+    the data at exactly the resume seam: the SSE spec says an event is dispatched
+    on the blank line, so a partial trailing block is not an event.
+    """
+    buf: list[str] = []
+    async for line in resp.aiter_lines():
+        if line == "":
+            if buf:
+                evt = parse_sse_event("\n".join(buf))
+                buf = []
+                if evt is not None:
+                    yield evt
+        else:
+            buf.append(line)
+
+
+def _sse_open_retry_delay(
+    status: int,
+    resp: httpx.Response,
+    attempt: int,
+    rng: random.Random,
+    *,
+    retry_not_found: bool = True,
+) -> float | None:
+    """Delay before re-opening a stream that failed on first contact (SSE/download).
+
+    ``None`` means "do not retry" -- the status is terminal, the caller opted out
+    of the 404 create-race wait, or the budget is spent. Only statuses that prove
+    the server did not start the work are retried, since these routes are POSTs.
+    A 429 is capped at the shared rate-limit budget (it used to honor Retry-After
+    forever) and, on the last allowed attempt, its ``None`` return sends the caller
+    to the ``SandboxRateLimitError`` mapping rather than looping.
+    """
+    if status == 404:
+        if not retry_not_found:
+            return None
+        return compute_backoff(attempt, delays=DEFAULT_NOT_FOUND_DELAYS_S, rng=rng)
+    if status == 429:
+        return _rate_limit_delay(resp, attempt, rng)
+    if status in (409, 503) and _says_not_ready(resp):
+        return compute_backoff(attempt, delays=DEFAULT_COLD_START_DELAYS_S, rng=rng)
+    return None

snowflake/sandbox/_stage.py

@@ -0,0 +1,270 @@
+"""Snowflake stage operations for code and image bundles.
+
+Resolves credentials, opens a connector connection, and PUTs a bundle. The
+connector PUT is the only upload path that works for every stage type (``@~/``,
+``@%table``, named); the Files REST API is limited to named stages.
+
+`_validate_stage_identifier` is the SQL-injection boundary for every path that
+names a stage, so it lives here with the code that interpolates one.
+
+Shared by the async and sync clients and by ``image``: the functions need only a
+``.config``, so they take any transport that carries one rather than naming a
+concrete class.
+"""
+
+from __future__ import annotations
+
+import re
+from typing import TYPE_CHECKING, Any, Protocol
+
+from snowflake.sandbox.exceptions import SandboxError
+
+if TYPE_CHECKING:
+    from snowflake.sandbox.config import Config
+    from snowflake.sandbox.mount import StageMount
+
+
+class _ConfigCarrier(Protocol):
+    """Anything that can supply a resolved `Config` — either transport does.
+
+    Stated as a protocol rather than ``Transport | SyncTransport`` because that is
+    the whole requirement, and naming both would make this module depend on both
+    for nothing.
+    """
+
+    @property
+    def config(self) -> Config: ...
+
+
+__all__ = [
+    "_resolve_stage_auth",
+    "_env_has_stage_credential",
+    "_connect_for_stage",
+    "_validate_stage_identifier",
+    "_safe_stage_filename",
+    "_resolve_stage_fqn",
+    "_put_file_to_stage",
+    "_parse_toml_stage_mounts",
+]
+
+
+# A Snowflake identifier part: an unquoted name (letters/digits/_/$, not starting
+# with a digit) or a double-quoted identifier ("" escapes an inner quote). A
+# fully-qualified stage is 1-3 dot-separated parts. Anything else — a space, a
+# quote-and-space (`S LOCATION='s3://attacker/'`), a semicolon, parentheses — is
+# rejected, so a `stage=`/`code_stage=` from an agent or config cannot break out
+# of `CREATE STAGE`/`PUT`/`LIST`/`GET_PRESIGNED_URL` and run attacker SQL as the
+# caller's role.
+_STAGE_ID_PART = re.compile(r'^(?:[A-Za-z_][A-Za-z0-9_$]*|"(?:[^"]|"")*")$')
+
+
+# Chars allowed in the stage-object filename we interpolate into PUT/LIST. The
+# hash16 is the content key; the dir-name prefix is cosmetic, so a name with a
+# quote/slash/space (a local directory basename is attacker-influenceable) is
+# squashed to underscores rather than left to break the `'file://...'` literal.
+_UNSAFE_FILENAME_CHARS = re.compile(r"[^A-Za-z0-9._-]")
+
+
+def _resolve_stage_auth(transport: _ConfigCarrier | None) -> tuple[str, str, dict[str, Any]]:
+    """``(account, host, connect_kwargs)`` for a connector stage PUT, resolved from
+    the environment (what Snowflake injects into a sandbox) first, then the SDK's config.
+    Raises `SandboxError` when neither host+token nor host+user/password is set."""
+    import os
+
+    from snowflake.sandbox.config import effective_config
+
+    _cfg = transport.config if transport else effective_config()
+    account_url = os.environ.get("SNOWFLAKE_ACCOUNT_URL", "").rstrip("/")
+    if not account_url:
+        try:
+            account_url = _cfg.base_url.rstrip("/")
+        except RuntimeError:
+            account_url = ""
+    token = os.environ.get("SNOWFLAKE_TOKEN") or os.environ.get("SNOWFLAKE_PAT") or (_cfg.pat or "")
+    _user = os.environ.get("SNOWFLAKE_USER", "")
+    _password = os.environ.get("SNOWFLAKE_PASSWORD", "")
+    if not account_url or not (token or (_user and _password)):
+        raise SandboxError(
+            "uploading local files to a stage needs a host and a credential: "
+            'name a connection (connection="my_connection", a ~/.snowflake default, '
+            "or SNOWFLAKE_HOST + SNOWFLAKE_PAT), or set "
+            "SNOWFLAKE_ACCOUNT_URL plus SNOWFLAKE_TOKEN (OAuth/Snowflake) or "
+            "SNOWFLAKE_USER + SNOWFLAKE_PASSWORD (local dev)"
+        )
+    _host = account_url.split("//")[-1].split(":")[0].rstrip("/").lower()
+    _account = (
+        _host[: -len(".snowflakecomputing.com")]
+        if _host.endswith(".snowflakecomputing.com")
+        else _host
+    )
+    if token:
+        connect_kwargs: dict[str, Any] = {"authenticator": "oauth", "token": token}
+    else:
+        connect_kwargs = {"user": _user, "password": _password}
+    return _account, _host, connect_kwargs
+
+
+def _env_has_stage_credential() -> bool:
+    """Whether the environment alone carries a usable stage credential.
+
+    Snowflake injects ``SNOWFLAKE_ACCOUNT_URL`` + ``SNOWFLAKE_TOKEN`` into sandbox
+    containers, and that has always taken precedence over the SDK's config; this
+    keeps that order intact.
+    """
+    import os
+
+    if not os.environ.get("SNOWFLAKE_ACCOUNT_URL", "").strip():
+        return False
+    if os.environ.get("SNOWFLAKE_TOKEN") or os.environ.get("SNOWFLAKE_PAT"):
+        return True
+    return bool(os.environ.get("SNOWFLAKE_USER") and os.environ.get("SNOWFLAKE_PASSWORD"))
+
+
+def _connect_for_stage(transport: _ConfigCarrier | None) -> tuple[Any, bool]:
+    """Return ``(connection, owns)`` for a stage PUT. The connector PUT is the only
+    upload path that works for every stage type (``@~/``, ``@%table``, named); the
+    Files REST API is limited to named stages.
+
+    When the config holds a live connection the connector already opened — a named
+    connection the SDK resolved, or one a caller handed to ``connection=<obj>`` (the
+    ``snow`` CLI does this) — and the environment carries no injected stage
+    credential, that connection is REUSED: ``owns`` is False and the caller must NOT
+    close it, because closing it would drop the session the REST transport renews its
+    bearer token from. This is the single-authority path: the connection that logged
+    in does the PUT too, with its own credential and role, so nothing is re-derived.
+
+    Otherwise a fresh connection is opened and the caller OWNS it (``owns`` True):
+    via the named connection's own connect kwargs (``connector_kwargs``), or from the
+    environment (what Snowflake injects into a sandbox — ``_resolve_stage_auth``).
+    Deriving auth from a bare token as ``authenticator="oauth"`` is correct only for
+    such an injected OAuth token, never for a minted *session* token (which fails
+    ``250001 Invalid OAuth access token``) — which is why a connection-backed config
+    is reused directly above and never reaches that derivation.
+    """
+    from snowflake.sandbox.config import effective_config
+
+    # A Transport may carry no config at all, so this cannot assume an object.
+    cfg = transport.config if transport is not None else effective_config()
+
+    # Single-authority reuse: the live connection already authenticated (its role,
+    # its credential), and the environment has no injected stage credential that must
+    # take precedence. Reuse it directly — the caller must not close it.
+    live = cfg.connection if cfg is not None else None
+    if live is not None and not _env_has_stage_credential():
+        return live, False  # owns=False: shared connection, caller must not close it
+
+    try:
+        import snowflake.connector
+    except ImportError as exc:
+        # snowflake-connector-python is a core dependency, so this should not
+        # happen in a normal install; guard defensively and name the real package.
+        raise SandboxError(
+            "uploading local files requires snowflake-connector-python: "
+            "pip install snowflake-connector-python"
+        ) from exc
+
+    conn_kwargs = cfg.connector_kwargs if cfg is not None else None
+    if conn_kwargs and not _env_has_stage_credential():
+        # connector_kwargs re-opens the named connection (connection_name=...), so
+        # the connector's CONFIG_MANAGER must point at the same home it was
+        # resolved from — it is a process-global that caches its last read.
+        from snowflake.sandbox.config import refresh_connection_config
+
+        refresh_connection_config()
+        return snowflake.connector.connect(**conn_kwargs), True
+
+    account, host, connect_kwargs = _resolve_stage_auth(transport)
+    return snowflake.connector.connect(account=account, host=host, **connect_kwargs), True
+
+
+def _validate_stage_identifier(name: str) -> str:
+    """Return *name* if it is a well-formed (optionally dotted) Snowflake stage
+    identifier, else raise `SandboxError`. This is the injection boundary for
+    every SQL path that names a stage (see `_STAGE_ID_PART`)."""
+    raw = name.strip()
+    if not raw:
+        raise SandboxError("stage name must not be empty")
+    parts = raw.split(".")
+    if len(parts) > 3 or not all(_STAGE_ID_PART.match(p) for p in parts):
+        raise SandboxError(
+            f"stage name {name!r} is not a valid Snowflake identifier: expected "
+            "DB.SCHEMA.NAME or a bare NAME (letters, digits, '_' or '$'; a "
+            'double-quoted "identifier" is also accepted). Names with spaces, '
+            "quotes, or SQL punctuation are rejected to prevent injection."
+        )
+    return raw
+
+
+def _safe_stage_filename(name: str) -> str:
+    """Squash *name* to the charset safe to interpolate into a PUT/LIST path."""
+    return _UNSAFE_FILENAME_CHARS.sub("_", name)
+
+
+def _resolve_stage_fqn(cur: Any, name: str) -> str:
+    """Resolve *name* to a fully-qualified stage and ensure it exists.
+
+    A ``DB.SCHEMA.NAME`` is used as-is; a bare name is qualified with the
+    connection's current database + schema. ``CREATE STAGE IF NOT EXISTS`` is then
+    run (via ``IDENTIFIER(?)`` bind) so the copy staging stage is auto-created.
+    *name* is validated as a Snowflake identifier first, so it cannot inject SQL.
+    Raises `SandboxError` if a bare name is given without a current db/schema."""
+    _validate_stage_identifier(name)
+    if name.count(".") >= 2:
+        fqn = name
+    else:
+        cur.execute("SELECT CURRENT_DATABASE(), CURRENT_SCHEMA()")
+        row = cur.fetchone() or (None, None)
+        db, schema = row[0], row[1]
+        if not db or not schema:
+            raise SandboxError(
+                f"copy_local needs a current database and schema to create the "
+                f"staging stage {name!r}. Set them on your connection (in "
+                f"connections.toml, or with USE DATABASE/USE SCHEMA), or pass "
+                f"stage='DB.SCHEMA.NAME'."
+            )
+        fqn = f"{db}.{schema}.{name}"
+    # IDENTIFIER(%s) binds the (validated) name as data, never as raw SQL text.
+    # %s is pyformat, the connector's default paramstyle.
+    cur.execute("CREATE STAGE IF NOT EXISTS IDENTIFIER(%s)", (fqn,))
+    return fqn
+
+
+def _put_file_to_stage(cur: Any, fqn: str, filename: str, zip_bytes: bytes) -> None:
+    """PUT *zip_bytes* as *filename* into stage *fqn* on an already-open cursor,
+    LIST-guarded so an already-present content-addressed object is not re-uploaded.
+
+    *fqn* must already be a validated identifier (see `_resolve_stage_fqn`);
+    *filename* is squashed to a safe charset so it cannot break out of the PUT
+    ``'file://...'`` literal or the LIST path."""
+    import os
+    import tempfile
+
+    _validate_stage_identifier(fqn)
+    safe_name = _safe_stage_filename(filename)
+    cur.execute(f"LIST @{fqn}/{safe_name}")
+    if cur.fetchone():
+        return
+    with tempfile.TemporaryDirectory() as _tmpdir:
+        local_zip = os.path.join(_tmpdir, safe_name)
+        with open(local_zip, "wb") as _f:
+            _f.write(zip_bytes)
+        cur.execute(f"PUT 'file://{local_zip}' @{fqn}/ AUTO_COMPRESS=FALSE OVERWRITE=TRUE")
+
+
+def _parse_toml_stage_mounts(raw: list[Any]) -> list[StageMount] | None:
+    """Parse ``[[stage_mounts]]`` TOML entries into `StageMount` objects."""
+    if not raw:
+        return None
+    from snowflake.sandbox.mount import StageMount
+
+    mounts = []
+    for entry in raw:
+        kind = entry.get("kind", "stage")
+        name = entry.get("name") or entry.get("stage_name", "")
+        mount_path = entry.get("mount_path", "")
+        readonly = bool(entry.get("readonly", False))
+        if kind == "workspace":
+            mounts.append(StageMount.from_workspace(name, mount_path=mount_path, readonly=readonly))
+        else:
+            mounts.append(StageMount.from_stage(name, mount_path=mount_path, readonly=readonly))
+    return mounts or None