diff --git a/DESCRIPTION.md b/DESCRIPTION.md index acd254db0..36e3c5f49 100644 --- 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 index dbdd2bc57..4d45e1914 100644 --- a/src/snowflake/connector/_utils.py +++ b/src/snowflake/connector/_utils.py @@ -1,5 +1,7 @@ from __future__ import annotations +import logging +import os import string from enum import Enum from inspect import stack @@ -7,6 +9,8 @@ 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 index a0f3d228b..aecf05bb5 100644 --- 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 index 4151748dd..941e293d3 100644 --- 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 new file mode 100644 index 000000000..a06eb797a --- /dev/null +++ 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 new file mode 100644 index 000000000..d6eb599d4 --- /dev/null +++ b/test/unit/test_spcs_token.py @@ -0,0 +1,144 @@ +from __future__ import annotations + +import json +from unittest import mock + +import pytest + +import snowflake.connector + + +@pytest.mark.skipolddriver +def test_spcs_token_included_in_login_request(monkeypatch): + """Verify that SPCS_TOKEN is injected into the login request body when present.""" + + # Use a custom SPCS token path and mock its existence and contents + 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") + monkeypatch.setattr("snowflake.connector._utils.open", mock_open, raising=False) + + captured_bodies: list[dict] = [] + + def mock_post_request(url, headers, body, **kwargs): + captured_bodies.append(json.loads(body)) + # Return a minimal successful login response + return { + "success": True, + "message": None, + "data": { + "token": "TOKEN", + "masterToken": "MASTER_TOKEN", + }, + } + + with mock.patch( + "snowflake.connector.network.SnowflakeRestful._post_request", + side_effect=mock_post_request, + ): + conn = snowflake.connector.connect( + account="testaccount", + user="testuser", + password="testpwd", + host="testaccount.snowflakecomputing.com", + ) + assert conn._rest.token == "TOKEN" + assert conn._rest.master_token == "MASTER_TOKEN" + + # Exactly one login-request should have been sent for this simple flow + assert len(captured_bodies) == 1 + body = captured_bodies[0] + assert body["data"]["SPCS_TOKEN"] == "TEST_SPCS_TOKEN" + + +@pytest.mark.skipolddriver +def test_spcs_token_not_included_when_file_missing(monkeypatch): + """Verify that SPCS_TOKEN is not added when the token file does not exist.""" + + # Ensure env var is unset so default path is used, but not created + monkeypatch.delenv("SF_SPCS_TOKEN_PATH", raising=False) + + captured_bodies: list[dict] = [] + + def mock_post_request(url, headers, body, **kwargs): + captured_bodies.append(json.loads(body)) + return { + "success": True, + "message": None, + "data": { + "token": "TOKEN", + "masterToken": "MASTER_TOKEN", + }, + } + + with mock.patch( + "snowflake.connector.network.SnowflakeRestful._post_request", + side_effect=mock_post_request, + ): + conn = snowflake.connector.connect( + account="testaccount", + user="testuser", + password="testpwd", + host="testaccount.snowflakecomputing.com", + ) + assert conn._rest.token == "TOKEN" + assert conn._rest.master_token == "MASTER_TOKEN" + + # Exactly one login-request should have been sent for this simple flow + assert len(captured_bodies) == 1 + body = captured_bodies[0] + assert "SPCS_TOKEN" not in body["data"] + + +@pytest.mark.skipolddriver +def test_spcs_token_default_path_used_when_env_unset(monkeypatch): + """When SF_SPCS_TOKEN_PATH is not set, default path should be used.""" + + # Ensure env var is unset so default path logic is exercised + monkeypatch.delenv("SF_SPCS_TOKEN_PATH", raising=False) + + # Default SPCS path inside SPCS container + 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") + monkeypatch.setattr("snowflake.connector._utils.open", mock_open, raising=False) + + captured_bodies: list[dict] = [] + + def mock_post_request(url, headers, body, **kwargs): + captured_bodies.append(json.loads(body)) + return { + "success": True, + "message": None, + "data": { + "token": "TOKEN", + "masterToken": "MASTER_TOKEN", + }, + } + + with mock.patch( + "snowflake.connector.network.SnowflakeRestful._post_request", + side_effect=mock_post_request, + ): + conn = snowflake.connector.connect( + account="testaccount", + user="testuser", + password="testpwd", + host="testaccount.snowflakecomputing.com", + ) + assert conn._rest.token == "TOKEN" + assert conn._rest.master_token == "MASTER_TOKEN" + + # Exactly one login-request should have been sent for this simple flow + assert len(captured_bodies) == 1 + body = captured_bodies[0] + assert body["data"]["SPCS_TOKEN"] == "DEFAULT_PATH_SPCS_TOKEN"