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

snowflake/sandbox/connect.py

@@ -0,0 +1,288 @@
+"""Snowflake connector helper for code running *inside* a Snowflake/gVisor sandbox.
+
+`get_snowflake_connection()` is the single public entry point, and it is supported
+public API — it is what workload code inside a sandbox calls to reach Snowflake. It
+reads the injected env vars and returns a live ``snowflake.connector`` connection.
+
+**Its contract is what the platform provides, not the caller's intent.** The credential
+is resolved in order: ``SNOWFLAKE_PAT``/``SNOWFLAKE_TOKEN`` in the environment, then a
+token file (``SNOWFLAKE_TOKEN_FILE_PATH``, else ``/snowflake/session/token``) with
+``SNOWFLAKE_ACCOUNT``, then the platform-rendered ``$SNOWFLAKE_HOME/connections.toml``.
+Which one is present varies by deployment and is not the caller's business. It also
+relies on the platform's egress-proxy CA at ``/etc/cng/proxy-ca.crt``. The token is **not** a real credential and is **not**
+injected by Snowflake: it is a per-app placeholder minted by Snowflake's ``AppCredential``, and the
+platform's outbound egress proxy swaps it for the account's real token on the way out
+(the same brokering that keeps secrets out of the sandbox). None of that exists outside
+a Snowflake sandbox, so this is **not** a general-purpose connector factory: calling it from
+laptop or service code that is not running inside a sandbox is unsupported and will
+either fail or, worse, widen that process's TLS trust (see `_patch_proxy_ca`) for no
+reason. Use ``snowflake.connector.connect`` directly there.
+
+So the narrowness here is a scope of *deployment*, not a scope of *seriousness*: this
+is in-sandbox-only, not dev/test-only, and in-sandbox workload code is a production
+caller. Please do not re-add a "dev / test use only" disclaimer — it would contradict
+this module's place in the package's public ``__all__``.
+
+This module has no framework dependency (no crewai, no langchain, no agents
+layer) so it can be imported from any sandbox workload regardless of which
+agent framework — or no framework at all — the workload uses.
+
+``snowflake-connector-python`` is a **core** dependency of the SDK, so it is
+always installed. Importing this module stays cheap regardless: the connector is
+loaded lazily inside `get_snowflake_connection()`, not at import time.
+"""
+
+from __future__ import annotations
+
+import os
+import pathlib
+from typing import Any
+
+__all__ = ["get_snowflake_connection"]
+
+# Connection parameters a caller may NOT override through get_snowflake_connection's
+# **kwargs. These define *who* the connection authenticates as and *where* it
+# connects; letting a caller (often an LLM assembling kwargs) redirect them would
+# send the sandbox's brokered credential to an arbitrary host or swap the auth
+# scheme. Everything else (warehouse, database, schema, role, session params, ...)
+# is passed through.
+_PROTECTED_CONNECT_KWARGS = frozenset(
+    {
+        "account",
+        "token",
+        "password",
+        "authenticator",
+        "host",
+        # connection_name / connections_file_path / token_file_path became
+        # live once the connections.toml path started passing connection_name:
+        # a caller-supplied TOML with a [default] entry would otherwise redirect
+        # account/host/token/authenticator wholesale.
+        "connection_name",
+        "connections_file_path",
+        "token_file_path",
+    }
+)
+
+# Default connection timeout (seconds). Without one the connector blocks
+# indefinitely when egress is closed; a sandbox under a restrictive Egress policy
+# would hang instead of failing.
+_DEFAULT_LOGIN_TIMEOUT_S = 30
+
+# Where a system CA bundle may live, in preference order. Debian/Ubuntu first,
+# then RHEL/Fedora, then the BSD/macOS layout Alpine also uses.
+_SYSTEM_CA_CANDIDATES = (
+    "/etc/ssl/certs/ca-certificates.crt",
+    "/etc/pki/tls/certs/ca-bundle.crt",
+    "/etc/ssl/cert.pem",
+)
+
+
+def _system_ca_bundle() -> str | None:
+    """Path to the OS's CA bundle, or ``None`` if none of the known ones exist.
+
+    Falls back to ``certifi`` (a hard dependency of the connector) so a slim /
+    distroless / RHEL / Alpine image without the Debian bundle still resolves a
+    real trust store rather than silently dropping the system CAs.
+    """
+    for path in _SYSTEM_CA_CANDIDATES:
+        if os.path.exists(path):
+            return path
+    try:
+        import certifi
+
+        return certifi.where()
+    except Exception:  # noqa: BLE001 - certifi absent/broken; no system bundle
+        return None
+
+
+def _patch_proxy_ca() -> None:
+    """Trust the Snowflake hostagent MITM CA alongside the system bundle.
+
+    When code runs inside a gVisor sandbox, the Snowflake hostagent intercepts
+    all outbound HTTPS and re-signs with its own CA at /etc/cng/proxy-ca.crt.
+    snowflake-connector-python uses requests/urllib3 internally; without this
+    patch every connector HTTPS call raises an SSL verification error.
+
+    **This widens TLS trust process-wide**: it sets ``REQUESTS_CA_BUNDLE`` and
+    ``SSL_CERT_FILE`` for the whole process, so *every* library that reads those
+    env vars — not just this connection — will additionally trust the Snowflake MITM CA.
+    That is intended inside a Snowflake sandbox (all egress is proxied) but is why this
+    module is not for general use outside one.
+
+    Called once per process: if REQUESTS_CA_BUNDLE / SSL_CERT_FILE is already set,
+    or the proxy CA file doesn't exist (outside a Snowflake sandbox), this is a no-op.
+    """
+    if os.environ.get("REQUESTS_CA_BUNDLE") or os.environ.get("SSL_CERT_FILE"):
+        return
+    proxy_ca = "/etc/cng/proxy-ca.crt"
+    if not os.path.exists(proxy_ca):
+        return
+    import tempfile
+
+    sys_ca = _system_ca_bundle()
+    if sys_ca is None:
+        # Trusting ONLY the proxy CA would silently drop every public root, so a
+        # Snowflake TLS handshake to a non-proxied host would fail with an opaque
+        # SSL error. Fail loudly instead — the fix (install ca-certificates or set
+        # REQUESTS_CA_BUNDLE) is actionable.
+        raise RuntimeError(
+            "get_snowflake_connection: no system CA bundle found (looked in "
+            f"{', '.join(_SYSTEM_CA_CANDIDATES)} and certifi). Cannot trust the "
+            "Snowflake proxy CA without dropping the system roots. Install a CA bundle "
+            "(e.g. ca-certificates) or set REQUESTS_CA_BUNDLE to a bundle that "
+            "already includes /etc/cng/proxy-ca.crt."
+        )
+    with open(sys_ca) as f:
+        combined = f.read()
+    with open(proxy_ca) as f:
+        combined += "\n" + f.read()
+    # delete=False so the bundle outlives this block; the context manager just
+    # closes the handle, and tmp.name stays valid for the env vars below.
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".pem", delete=False) as tmp:
+        tmp.write(combined)
+    os.environ["REQUESTS_CA_BUNDLE"] = tmp.name
+    os.environ["SSL_CERT_FILE"] = tmp.name
+
+
+# Where the platform leaves the sandbox's session token when it names no other path,
+# and the connection entry it renders. Both are platform contracts, not guesses:
+# sandbox-server's _render_snowflake_connection writes exactly one [default] entry.
+_DEFAULT_TOKEN_FILE = "/snowflake/session/token"
+_PLATFORM_CONNECTION_NAME = "default"
+
+
+def _resolve_token() -> str | None:
+    """The sandbox's session token, by whichever mechanism the platform used.
+
+    The env VALUE comes first (unchanged). Otherwise the token is a FILE: cng-explore
+    #1571 names it in SNOWFLAKE_TOKEN_FILE_PATH, and the well-known path is the
+    fallback. It is a per-app AppCredential placeholder, swapped for the real token by
+    the egress proxy on the way out — so reading it here is not handling a secret.
+    """
+    env_value = os.environ.get("SNOWFLAKE_PAT") or os.environ.get("SNOWFLAKE_TOKEN")
+    if env_value:
+        return env_value
+    path = os.environ.get("SNOWFLAKE_TOKEN_FILE_PATH") or _DEFAULT_TOKEN_FILE
+    try:
+        token = pathlib.Path(path).read_text(encoding="utf-8").strip()
+    except (OSError, UnicodeDecodeError):
+        return None
+    return token or None
+
+
+def _host_from_account_url() -> str | None:
+    """SNOWFLAKE_ACCOUNT_URL as a bare host, since the connector wants no scheme."""
+    url = os.environ.get("SNOWFLAKE_ACCOUNT_URL", "").strip()
+    if not url:
+        return None
+    return url.split("://", 1)[-1].rstrip("/").split("/", 1)[0] or None
+
+
+def _connections_toml() -> pathlib.Path | None:
+    """The platform-rendered connections.toml, if SNOWFLAKE_HOME holds one."""
+    home = os.environ.get("SNOWFLAKE_HOME", "").strip()
+    if not home:
+        return None
+    path = pathlib.Path(home) / "connections.toml"
+    return path if path.exists() else None
+
+
+def _make_connection(**kwargs: Any) -> Any:
+    """Open a new Snowflake connection from env vars, merging ``kwargs`` overrides."""
+    protected = _PROTECTED_CONNECT_KWARGS & kwargs.keys()
+    if protected:
+        raise ValueError(
+            "get_snowflake_connection does not accept "
+            f"{sorted(protected)} as keyword overrides: the account, host, and "
+            "credential come from the sandbox's injected environment and must not "
+            "be redirected. Pass only non-credential options (warehouse, database, "
+            "schema, role, session parameters, ...)."
+        )
+    import snowflake.connector
+
+    account = os.environ.get("SNOWFLAKE_ACCOUNT")
+    token = _resolve_token()
+    if account and token:
+        params: dict[str, Any] = {
+            "account": account,
+            "token": token,
+            "authenticator": "oauth",
+            "login_timeout": _DEFAULT_LOGIN_TIMEOUT_S,
+        }
+        # SNOWFLAKE_HOST first, but ACCOUNT_URL is what today's platform actually sets
+        # and the host is not optional: without it the connector builds
+        # <account>.snowflakecomputing.com and fails 540001 on any non-default region.
+        host = os.environ.get("SNOWFLAKE_HOST") or _host_from_account_url()
+        if host:
+            params["host"] = host
+        for k, env in (
+            ("user", "SNOWFLAKE_USER"),
+            ("role", "SNOWFLAKE_ROLE"),
+            ("warehouse", "SNOWFLAKE_WAREHOUSE"),
+        ):
+            v = os.environ.get(env)
+            if v:
+                params[k] = v
+        params.update(kwargs)
+        return snowflake.connector.connect(**params)
+
+    # The platform's own rendered config. connection_name is passed explicitly because
+    # ANY kwarg otherwise suppresses the connector's default-connection resolution, and
+    # the caller may well pass warehouse= — which failed as 251005 "User is empty".
+    if _connections_toml() is not None:
+        # Merge like the branch above rather than splatting **kwargs alongside
+        # hard-coded keywords: a caller passing login_timeout= (documented as
+        # overridable) would otherwise be a duplicate-keyword TypeError.
+        params = {
+            "connection_name": _PLATFORM_CONNECTION_NAME,
+            "login_timeout": _DEFAULT_LOGIN_TIMEOUT_S,
+        }
+        params.update(kwargs)
+        return snowflake.connector.connect(**params)
+
+    raise RuntimeError(
+        "get_snowflake_connection found no usable credential. Looked for: "
+        "SNOWFLAKE_PAT/SNOWFLAKE_TOKEN in the environment; a token file at "
+        f"$SNOWFLAKE_TOKEN_FILE_PATH or {_DEFAULT_TOKEN_FILE} (with SNOWFLAKE_ACCOUNT "
+        f"set); and a [{_PLATFORM_CONNECTION_NAME}] entry in "
+        "$SNOWFLAKE_HOME/connections.toml. This helper is for code running INSIDE a "
+        "Snowflake sandbox, where the platform provides one of those."
+    )
+
+
+def get_snowflake_connection(**kwargs: Any) -> Any:
+    """Return a Snowflake connection for code running inside a Snowflake sandbox.
+
+    Resolves the credential the way the platform provides it, in order: the
+    ``SNOWFLAKE_PAT``/``SNOWFLAKE_TOKEN`` env value; a token file named by
+    ``SNOWFLAKE_TOKEN_FILE_PATH`` or at ``/snowflake/session/token``, with
+    ``SNOWFLAKE_ACCOUNT``; then a ``[default]`` entry in the platform-rendered
+    ``$SNOWFLAKE_HOME/connections.toml``. Deployments differ in which they set, so
+    reading only one made this raise where a credential was in fact present.
+
+    The token is Snowflake's ``AppCredential`` placeholder, swapped for the real account
+    token by the platform's outbound egress proxy — not a credential Snowflake hands the
+    container. Non-credential keyword arguments override the corresponding default, e.g.
+    ``get_snowflake_connection(warehouse="MY_WH", database="MY_DB")``.
+
+    The identity-bearing parameters — ``account``, ``host``, ``token``,
+    ``password``, ``authenticator`` — cannot be overridden via ``**kwargs`` and
+    raise ``ValueError`` if passed: they come from the sandbox's injected
+    environment and must not be redirected to another host or auth scheme.
+
+    The platform's egress-proxy CA (``/etc/cng/proxy-ca.crt``) is automatically
+    trusted so HTTPS verification succeeds inside gVisor. **This widens TLS trust
+    process-wide** — it sets ``REQUESTS_CA_BUNDLE``/``SSL_CERT_FILE`` for the whole
+    process (system roots + the Snowflake CA), affecting every library that reads them,
+    which is why this helper is for in-sandbox use only. Outside a Snowflake sandbox the
+    proxy CA file doesn't exist and normal system CA verification is used. If no
+    system CA bundle can be found it raises rather than trust only the proxy CA
+    (which would drop the public roots).
+
+    A default ``login_timeout`` is applied so a closed-egress sandbox fails fast
+    instead of hanging; override it with ``login_timeout=`` if needed.
+
+    Each call returns a *new* connection; callers are responsible for closing it.
+    """
+    _patch_proxy_ca()
+    return _make_connection(**kwargs)