[SNOW-2409721] Add SPCS_TOKEN to SF login requests originating from SPCS containers (#2714)

sfc-gh-mpenumathsa · web-flow · commit 2586881ec535 · 2025-12-12T14:09:09.000+01:00
diff --git a/DESCRIPTION.md b/DESCRIPTION.md
@@ -10,6 +10,7 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
 - v4.2.0(TBD)
   - Added support for async I/O. Asynchronous version of connector is available via `snowflake.connector.aio` module.
   - Added `SnowflakeCursor.stats` property to expose granular DML statistics (rows inserted, deleted, updated, and duplicates) for operations like CTAS where `rowcount` is insufficient.
+  - Added support for injecting SPCS service identifier token (`SPCS_TOKEN`) into login requests when present in SPCS containers.
 
 - v4.1.1(TBD)
   - Relaxed pandas dependency requirements for Python below 3.12.
diff --git a/src/snowflake/connector/_utils.py b/src/snowflake/connector/_utils.py
@@ -1,12 +1,16 @@
 from __future__ import annotations
 
+import logging
+import os
 import string
 from enum import Enum
 from inspect import stack
 from random import choice
 from threading import Timer
 from uuid import UUID
 
+logger = logging.getLogger(__name__)
+
 
 class TempObjectType(Enum):
     TABLE = "TABLE"
@@ -96,3 +100,30 @@ def get_application_path() -> str:
         return outermost_frame.filename
     except Exception:
         return "unknown"
+
+
+_SPCS_TOKEN_ENV_VAR_NAME = "SF_SPCS_TOKEN_PATH"
+_SPCS_TOKEN_DEFAULT_PATH = "/snowflake/session/spcs_token"
+
+
+def get_spcs_token() -> str | None:
+    """Return the SPCS token read from the configured path, or None.
+
+    The path is determined by the SF_SPCS_TOKEN_PATH environment variable,
+    falling back to ``/snowflake/session/spcs_token`` when unset.
+
+    Any I/O errors or missing/empty files are treated as \"no token\" and
+    will not cause authentication to fail.
+    """
+    path = os.getenv(_SPCS_TOKEN_ENV_VAR_NAME) or _SPCS_TOKEN_DEFAULT_PATH
+    try:
+        if not os.path.isfile(path):
+            return None
+        with open(path, encoding="utf-8") as f:
+            token = f.read().strip()
+        if not token:
+            return None
+        return token
+    except Exception as exc:  # pragma: no cover - best-effort logging only
+        logger.debug("Failed to read SPCS token from %s: %s", path, exc)
+        return None
diff --git a/src/snowflake/connector/aio/auth/_auth.py b/src/snowflake/connector/aio/auth/_auth.py
@@ -108,6 +108,8 @@ async def authenticate(
         )
 
         body = copy.deepcopy(body_template)
+        # Add SPCS token if present, independent of authenticator type.
+        self._add_spcs_token_to_body(body)
         # updating request body
         await auth_instance.update_body(body)
 
@@ -221,6 +223,8 @@ async def post_request_wrapper(self, url, headers, body) -> None:
             ):
                 body = copy.deepcopy(body_template)
                 body["inFlightCtx"] = ret["data"].get("inFlightCtx")
+                # Add SPCS token to the follow-up login request as well.
+                self._add_spcs_token_to_body(body)
                 # final request to get tokens
                 ret = await self._rest._post_request(
                     url,
@@ -261,6 +265,8 @@ async def post_request_wrapper(self, url, headers, body) -> None:
                     else None
                 )
                 body["data"]["CHOSEN_NEW_PASSWORD"] = password_callback()
+                # Add SPCS token to the password change login request as well.
+                self._add_spcs_token_to_body(body)
                 # New Password input
                 ret = await self._rest._post_request(
                     url,
diff --git a/src/snowflake/connector/auth/_auth.py b/src/snowflake/connector/auth/_auth.py
@@ -17,7 +17,7 @@
     load_pem_private_key,
 )
 
-from .._utils import get_application_path
+from .._utils import get_application_path, get_spcs_token
 from ..compat import urlencode
 from ..constants import (
     DAY_IN_SECONDS,
@@ -95,6 +95,17 @@ def __init__(self, rest) -> None:
         self._rest = rest
         self._token_cache: TokenCache | None = None
 
+    def _add_spcs_token_to_body(self, body: dict[Any, Any]) -> None:
+        """Inject SPCS_TOKEN into the login request body when available.
+
+        This reads the SPCS token from the path specified by SF_SPCS_TOKEN_PATH,
+        or from ``/snowflake/session/spcs_token`` when the env var is unset.
+        """
+        spcs_token = get_spcs_token()
+        if spcs_token is not None:
+            # Ensure the \"data\" envelope exists and add the token.
+            body.setdefault("data", {})["SPCS_TOKEN"] = spcs_token
+
     @staticmethod
     def base_auth_data(
         user,
@@ -205,6 +216,8 @@ def authenticate(
         )
 
         body = copy.deepcopy(body_template)
+        # Add SPCS token if present, independent of authenticator type.
+        self._add_spcs_token_to_body(body)
         # updating request body
         auth_instance.update_body(body)
 
@@ -323,6 +336,8 @@ def post_request_wrapper(self, url, headers, body) -> None:
             ):
                 body = copy.deepcopy(body_template)
                 body["inFlightCtx"] = ret["data"].get("inFlightCtx")
+                # Add SPCS token to the follow-up login request as well.
+                self._add_spcs_token_to_body(body)
                 # final request to get tokens
                 ret = self._rest._post_request(
                     url,
@@ -363,6 +378,8 @@ def post_request_wrapper(self, url, headers, body) -> None:
                     else None
                 )
                 body["data"]["CHOSEN_NEW_PASSWORD"] = password_callback()
+                # Add SPCS token to the password change login request as well.
+                self._add_spcs_token_to_body(body)
                 # New Password input
                 ret = self._rest._post_request(
                     url,
diff --git a/test/unit/aio/test_spcs_token_async.py b/test/unit/aio/test_spcs_token_async.py
@@ -0,0 +1,150 @@
+from __future__ import annotations
+
+import json
+from unittest import mock
+
+import pytest
+
+import snowflake.connector.aio
+
+
+@pytest.mark.skipolddriver
+async def test_spcs_token_included_in_login_request_async(monkeypatch):
+    """Verify that SPCS_TOKEN is injected into async login request body when present."""
+
+    custom_path = "/custom/path/to/spcs_token"
+    monkeypatch.setenv("SF_SPCS_TOKEN_PATH", custom_path)
+    monkeypatch.setattr(
+        "snowflake.connector._utils.os.path.isfile",
+        lambda path: path == custom_path,
+        raising=False,
+    )
+    mock_open = mock.mock_open(read_data="TEST_SPCS_TOKEN_ASYNC")
+    monkeypatch.setattr("snowflake.connector._utils.open", mock_open, raising=False)
+
+    captured_requests: list[tuple[str, dict]] = []
+
+    async def mock_post_request(url, headers, body, **kwargs):
+        captured_requests.append((url, json.loads(body)))
+        return {
+            "success": True,
+            "message": None,
+            "data": {
+                "token": "TOKEN_ASYNC",
+                "masterToken": "MASTER_TOKEN_ASYNC",
+            },
+        }
+
+    with mock.patch(
+        "snowflake.connector.aio._network.SnowflakeRestful._post_request",
+        side_effect=mock_post_request,
+    ):
+        conn = snowflake.connector.aio.SnowflakeConnection(
+            account="testaccount",
+            user="testuser",
+            password="testpwd",
+            host="testaccount.snowflakecomputing.com",
+        )
+        await conn.connect()
+        assert conn._rest.token == "TOKEN_ASYNC"
+        assert conn._rest.master_token == "MASTER_TOKEN_ASYNC"
+        await conn.close()
+
+    # Exactly one login-request should have been sent for this simple flow
+    login_bodies = [body for (url, body) in captured_requests if "login-request" in url]
+    assert len(login_bodies) == 1
+    body = login_bodies[0]
+    assert body["data"]["SPCS_TOKEN"] == "TEST_SPCS_TOKEN_ASYNC"
+
+
+@pytest.mark.skipolddriver
+async def test_spcs_token_not_included_when_file_missing_async(monkeypatch):
+    """Verify that SPCS_TOKEN is not added to async login request when file does not exist."""
+
+    monkeypatch.delenv("SF_SPCS_TOKEN_PATH", raising=False)
+
+    captured_requests: list[tuple[str, dict]] = []
+
+    async def mock_post_request(url, headers, body, **kwargs):
+        captured_requests.append((url, json.loads(body)))
+        return {
+            "success": True,
+            "message": None,
+            "data": {
+                "token": "TOKEN_ASYNC",
+                "masterToken": "MASTER_TOKEN_ASYNC",
+            },
+        }
+
+    with mock.patch(
+        "snowflake.connector.aio._network.SnowflakeRestful._post_request",
+        side_effect=mock_post_request,
+    ):
+        conn = snowflake.connector.aio.SnowflakeConnection(
+            account="testaccount",
+            user="testuser",
+            password="testpwd",
+            host="testaccount.snowflakecomputing.com",
+        )
+        await conn.connect()
+        assert conn._rest.token == "TOKEN_ASYNC"
+        assert conn._rest.master_token == "MASTER_TOKEN_ASYNC"
+        await conn.close()
+
+    # Exactly one login-request should have been sent for this simple flow
+    login_bodies = [body for (url, body) in captured_requests if "login-request" in url]
+    assert len(login_bodies) == 1
+    body = login_bodies[0]
+    assert "SPCS_TOKEN" not in body["data"]
+
+
+@pytest.mark.skipolddriver
+async def test_spcs_token_default_path_used_when_env_unset_async(
+    monkeypatch,
+):
+    """When SF_SPCS_TOKEN_PATH is not set, default path should be used (async)."""
+
+    monkeypatch.delenv("SF_SPCS_TOKEN_PATH", raising=False)
+
+    default_path = "/snowflake/session/spcs_token"
+    monkeypatch.setattr(
+        "snowflake.connector._utils.os.path.isfile",
+        lambda path: path == default_path,
+        raising=False,
+    )
+    mock_open = mock.mock_open(read_data="DEFAULT_PATH_SPCS_TOKEN_ASYNC")
+    monkeypatch.setattr("snowflake.connector._utils.open", mock_open, raising=False)
+
+    captured_requests: list[tuple[str, dict]] = []
+
+    async def mock_post_request(url, headers, body, **kwargs):
+        captured_requests.append((url, json.loads(body)))
+        return {
+            "success": True,
+            "message": None,
+            "data": {
+                "token": "TOKEN_ASYNC",
+                "masterToken": "MASTER_TOKEN_ASYNC",
+            },
+        }
+
+    with mock.patch(
+        "snowflake.connector.aio._network.SnowflakeRestful._post_request",
+        side_effect=mock_post_request,
+    ):
+        conn = snowflake.connector.aio.SnowflakeConnection(
+            account="testaccount",
+            user="testuser",
+            password="testpwd",
+            host="testaccount.snowflakecomputing.com",
+        )
+        await conn.connect()
+        assert conn._rest.token == "TOKEN_ASYNC"
+        assert conn._rest.master_token == "MASTER_TOKEN_ASYNC"
+        await conn.close()
+
+    # Exactly one login-request should have been sent for this simple flow
+    login_bodies = [body for (url, body) in captured_requests if "login-request" in url]
+    assert len(login_bodies) == 1
+    body = login_bodies[0]
+    assert body["data"]["SPCS_TOKEN"] == "DEFAULT_PATH_SPCS_TOKEN_ASYNC"
diff --git a/test/unit/test_spcs_token.py b/test/unit/test_spcs_token.py
