From df7672cf164119bc99a827d50d11def014fe1d47 Mon Sep 17 00:00:00 2001 From: Tanmay Mehta Date: Mon, 27 Apr 2026 16:57:40 +0000 Subject: [PATCH 1/5] initial changes --- CHANGELOG.md | 1 + docs/source/snowpark/secrets.rst | 1 + src/snowflake/snowpark/secrets.py | 39 +++++++++++++++++++++++++++++++ tests/integ/test_secrets.py | 3 +++ tests/unit/test_secrets.py | 12 ++++++++++ 5 files changed, 56 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index fba8710a2f..2ed988dcf2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -38,6 +38,7 @@ #### New Features - Added `artifact_repository` support to `udtf_configs` in `session.read.dbapi()`, enabling users to specify a custom artifact repository (e.g. PyPI) for packages used by the internal UDTF during distributed ingestion. +- Added `get_wif_token` to `snowflake.snowpark.secrets` for workload identity federation tokens on the Snowflake server (not available in SPCS file-based secret environments). #### Bug Fixes diff --git a/docs/source/snowpark/secrets.rst b/docs/source/snowpark/secrets.rst index 0c650171fe..2edcd41e19 100644 --- a/docs/source/snowpark/secrets.rst +++ b/docs/source/snowpark/secrets.rst @@ -22,3 +22,4 @@ Snowpark Secrets get_secret_type get_username_password get_cloud_provider_token + get_wif_token diff --git a/src/snowflake/snowpark/secrets.py b/src/snowflake/snowpark/secrets.py index cce9a15d7b..cae3521ba2 100644 --- a/src/snowflake/snowpark/secrets.py +++ b/src/snowflake/snowpark/secrets.py @@ -20,6 +20,7 @@ "get_secret_type", "get_username_password", "get_cloud_provider_token", + "get_wif_token", "UsernamePassword", "CloudProviderToken", ] @@ -61,6 +62,10 @@ def get_username_password(self, secret_name: str) -> UsernamePassword: def get_cloud_provider_token(self, secret_name: str) -> CloudProviderToken: pass + @abstractmethod + def get_wif_token(self, secret_name: str, audience: str) -> str: + pass + class _SnowflakeSecretsServer(_SnowflakeSecrets): """Secret instance for Snowflake server environment (using _snowflake module).""" @@ -89,6 +94,9 @@ def get_cloud_provider_token(self, secret_name: str) -> CloudProviderToken: secret_object.token, ) + def get_wif_token(self, secret_name: str, audience: str) -> str: + return self._snowflake.get_wif_token(secret_name, audience) + class _SnowflakeSecretsSPCS(_SnowflakeSecrets): """Secret instance for SPCS container environment (file-based secrets).""" @@ -173,6 +181,11 @@ def get_cloud_provider_token(self, secret_name: str) -> CloudProviderToken: "Cloud provider token secrets are not supported in SPCS container environments." ) + def get_wif_token(self, secret_name: str, audience: str) -> str: + raise NotImplementedError( + "WIF token secrets are not supported in SPCS container environments." + ) + def _is_spcs_environment() -> bool: return os.getenv(_SCLS_SPCS_SECRET_ENV_NAME, None) is not None @@ -259,3 +272,29 @@ def get_cloud_provider_token(secret_name: str) -> CloudProviderToken: NotImplementedError: If running outside Snowflake server environment. """ return _get_secrets_instance().get_cloud_provider_token(secret_name) + + +def get_wif_token(secret_name: str, audience: str) -> str: + """Get a workload identity federation (WIF) token from Snowflake. + + Note: + Requires a Snowflake environment with a WIF secret configured and an + external access integration that allows the UDF or stored procedure to + use that secret. The ``audience`` must match the token audience expected + by the external system (for example, an OAuth token endpoint URL). + + Args: + secret_name: The secret reference name bound to the WIF secret. + audience: The intended audience (``aud``) for the issued token. + + Returns: + The issued token as a string (typically a JWT). + + Raises: + NotImplementedError: If running outside the Snowflake server environment + (including SPCS file-based secret environments, where WIF tokens cannot + be minted). + ValueError: If the secret does not exist or is not authorized (when + applicable in supported environments). + """ + return _get_secrets_instance().get_wif_token(secret_name, audience) diff --git a/tests/integ/test_secrets.py b/tests/integ/test_secrets.py index f67b9b5701..1a7623bfe5 100644 --- a/tests/integ/test_secrets.py +++ b/tests/integ/test_secrets.py @@ -9,6 +9,7 @@ get_secret_type, get_cloud_provider_token, get_oauth_access_token, + get_wif_token, ) from snowflake.snowpark.types import BooleanType, StringType from tests.utils import IS_NOT_ON_GITHUB, RUNNING_ON_JENKINS, IS_IN_STORED_PROC, Utils @@ -169,3 +170,5 @@ def test_secrets_import_error(): get_cloud_provider_token("c1") with pytest.raises(NotImplementedError): get_oauth_access_token("o1") + with pytest.raises(NotImplementedError): + get_wif_token("w1", "https://audience") diff --git a/tests/unit/test_secrets.py b/tests/unit/test_secrets.py index 0964f212da..b5e393b860 100644 --- a/tests/unit/test_secrets.py +++ b/tests/unit/test_secrets.py @@ -12,6 +12,7 @@ get_secret_type, get_username_password, get_cloud_provider_token, + get_wif_token, UsernamePassword, CloudProviderToken, _SCLS_SPCS_SECRET_ENV_NAME, @@ -31,6 +32,7 @@ def _build_fake_snowflake_module() -> object: get_secret_type=lambda secret_name: "PASSWORD", get_username_password=lambda secret_name: fake_username_password, get_cloud_provider_token=lambda secret_name: fake_cloud_token, + get_wif_token=lambda secret_name, audience: f"wif:{secret_name}:{audience}", ) @@ -52,6 +54,11 @@ def test_secrets_mock_server_paths(): assert cloud.secret_access_key == "SECRET_TEST" assert cloud.token == "STS_TOKEN_TEST" + assert ( + get_wif_token("w1", "https://example.com/aud") + == "wif:w1:https://example.com/aud" + ) + @pytest.fixture def scls_spcs_mock_env(tmp_path): @@ -135,6 +142,9 @@ def test_secrets_mock_scls_spcs_error_cases(scls_spcs_mock_env): with pytest.raises(NotImplementedError): get_cloud_provider_token("any_secret") + with pytest.raises(NotImplementedError): + get_wif_token("any_secret", "https://audience") + with pytest.raises(ValueError, match="Unknown secret type"): get_secret_type("unknown_secret") @@ -159,6 +169,8 @@ def test_secrets_import_error_paths(): get_username_password("p1") with pytest.raises(NotImplementedError): get_cloud_provider_token("c1") + with pytest.raises(NotImplementedError): + get_wif_token("w1", "https://audience") finally: if original_env is not None: os.environ[_SCLS_SPCS_SECRET_ENV_NAME] = original_env From 8ff6f5686fdf465eb0dd6c5a9f3bcc675625fe8d Mon Sep 17 00:00:00 2001 From: Tanmay Mehta Date: Wed, 6 May 2026 22:36:12 +0000 Subject: [PATCH 2/5] added test in tests/integ/test_secrets.py --- tests/integ/conftest.py | 34 +++++++++++++++++++++++++ tests/integ/test_secrets.py | 50 +++++++++++++++++++++++++++++++++++++ 2 files changed, 84 insertions(+) diff --git a/tests/integ/conftest.py b/tests/integ/conftest.py index 0ac231a396..9c845b8e14 100644 --- a/tests/integ/conftest.py +++ b/tests/integ/conftest.py @@ -74,6 +74,9 @@ def set_up_external_access_integration_resources( integration1, integration2, integration3, + key4, + integration4, + wif_audience, ): try: # IMPORTANT SETUP NOTES: the test role needs to be granted the creation privilege @@ -128,6 +131,12 @@ def set_up_external_access_integration_resources( ).collect() session.sql( f""" + CREATE SECRET IF NOT EXISTS {key4} + TYPE = WORKLOAD_IDENTITY_FEDERATION; + """ + ).collect() + session.sql( + f""" CREATE IF NOT EXISTS EXTERNAL ACCESS INTEGRATION {integration1} ALLOWED_NETWORK_RULES = ({rule1}) ALLOWED_AUTHENTICATION_SECRETS = ({key1}) @@ -148,6 +157,13 @@ def set_up_external_access_integration_resources( ALLOWED_NETWORK_RULES = ({rule3}) ALLOWED_AUTHENTICATION_SECRETS = ({key3}) ENABLED = true; + """ + ).collect() + session.sql( + f""" + CREATE IF NOT EXISTS EXTERNAL ACCESS INTEGRATION {integration4} + ALLOWED_AUTHENTICATION_SECRETS = ({key4}) + ENABLED = true; """ ).collect() CONNECTION_PARAMETERS["external_access_rule1"] = rule1 @@ -156,9 +172,12 @@ def set_up_external_access_integration_resources( CONNECTION_PARAMETERS["external_access_key1"] = key1 CONNECTION_PARAMETERS["external_access_key2"] = key2 CONNECTION_PARAMETERS["external_access_key3"] = key3 + CONNECTION_PARAMETERS["external_access_key4"] = key4 CONNECTION_PARAMETERS["external_access_integration1"] = integration1 CONNECTION_PARAMETERS["external_access_integration2"] = integration2 CONNECTION_PARAMETERS["external_access_integration3"] = integration3 + CONNECTION_PARAMETERS["external_access_integration4"] = integration4 + CONNECTION_PARAMETERS["wif_audience"] = wif_audience except SnowparkSQLException: # GCP currently does not support external access integration # we can remove the exception once the integration is available on GCP @@ -184,9 +203,12 @@ def clean_up_external_access_integration_resources(): CONNECTION_PARAMETERS.pop("external_access_key1", None) CONNECTION_PARAMETERS.pop("external_access_key2", None) CONNECTION_PARAMETERS.pop("external_access_key3", None) + CONNECTION_PARAMETERS.pop("external_access_key4", None) CONNECTION_PARAMETERS.pop("external_access_integration1", None) CONNECTION_PARAMETERS.pop("external_access_integration2", None) CONNECTION_PARAMETERS.pop("external_access_integration3", None) + CONNECTION_PARAMETERS.pop("external_access_integration4", None) + CONNECTION_PARAMETERS.pop("wif_audience", None) def set_up_dataframe_processor_parameters( @@ -315,9 +337,12 @@ def session( key1 = "snowpark_python_test_key1" key2 = "snowpark_python_test_key2" key3 = "snowpark_python_test_key3" + key4 = "snowpark_python_test_key4" integration1 = "snowpark_python_test_integration1" integration2 = "snowpark_python_test_integration2" integration3 = "snowpark_python_test_integration3" + integration4 = "snowpark_python_test_integration4" + wif_audience = "https://replace-with-your-wif-audience" session = ( Session.builder.configs(db_parameters) @@ -351,6 +376,9 @@ def session( integration1, integration2, integration3, + key4, + integration4, + wif_audience, ) if validate_ast: @@ -387,9 +415,12 @@ def profiler_session( key1 = "snowpark_python_profiler_test_key1" key2 = "snowpark_python_profiler_test_key2" key3 = "snowpark_python_profiler_test_key3" + key4 = "snowpark_python_profiler_test_key4" integration1 = "snowpark_python_profiler_test_integration1" integration2 = "snowpark_python_profiler_test_integration2" integration3 = "snowpark_python_profiler_test_integration3" + integration4 = "snowpark_python_profiler_test_integration4" + wif_audience = "https://replace-with-your-wif-audience" session = ( Session.builder.configs(db_parameters) .config("local_testing", local_testing_mode) @@ -409,6 +440,9 @@ def profiler_session( integration1, integration2, integration3, + key4, + integration4, + wif_audience, ) set_up_test_session_parameters(session, local_testing_mode) try: diff --git a/tests/integ/test_secrets.py b/tests/integ/test_secrets.py index 1a7623bfe5..c8183c7c50 100644 --- a/tests/integ/test_secrets.py +++ b/tests/integ/test_secrets.py @@ -153,6 +153,56 @@ def get_secret(): ) +@pytest.mark.skipif( + IS_NOT_ON_GITHUB or not RUNNING_ON_JENKINS, + reason="Secret API is only supported on Snowflake server environment", +) +def test_get_wif_token_udf(session, db_parameters): + def get_wif(): + token = get_wif_token("cred", db_parameters["wif_audience"]) + return len(token) > 0 + + try: + get_wif_udf = session.udf.register( + get_wif, + return_type=BooleanType(), + packages=["snowflake-snowpark-python"], + external_access_integrations=[ + db_parameters["external_access_integration4"] + ], + secrets={"cred": f"{db_parameters['external_access_key4']}"}, + ) + df = session.create_dataframe([[1], [2]]).to_df("x") + Utils.check_answer(df.select(get_wif_udf()), [Row(True), Row(True)]) + except KeyError: + pytest.skip("External Access Integration is not supported on the deployment.") + + +@pytest.mark.skipif( + IS_NOT_ON_GITHUB or not RUNNING_ON_JENKINS, + reason="Secret API is only supported on Snowflake server environment", +) +def test_get_wif_token_sproc(session, db_parameters): + def get_wif_in_sproc(session_): + token = get_wif_token("cred", db_parameters["wif_audience"]) + return len(token) > 0 + + try: + get_wif_sp = session.sproc.register( + get_wif_in_sproc, + return_type=BooleanType(), + packages=["snowflake-snowpark-python"], + external_access_integrations=[ + db_parameters["external_access_integration4"] + ], + secrets={"cred": f"{db_parameters['external_access_key4']}"}, + anonymous=True, + ) + assert get_wif_sp() + except KeyError: + pytest.skip("External Access Integration is not supported on the deployment.") + + @pytest.mark.skipif( IS_IN_STORED_PROC, reason="Run only outside Snowflake server to validate NotImplementedError", From edbcd6d4c9736a9adceb7b548abb4dad8f2c4cdf Mon Sep 17 00:00:00 2001 From: Tanmay Mehta Date: Thu, 21 May 2026 06:47:23 +0000 Subject: [PATCH 3/5] fix for precommit coverage Co-authored-by: Cursor --- src/snowflake/snowpark/secrets.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/snowflake/snowpark/secrets.py b/src/snowflake/snowpark/secrets.py index cae3521ba2..ffad8500f8 100644 --- a/src/snowflake/snowpark/secrets.py +++ b/src/snowflake/snowpark/secrets.py @@ -64,7 +64,7 @@ def get_cloud_provider_token(self, secret_name: str) -> CloudProviderToken: @abstractmethod def get_wif_token(self, secret_name: str, audience: str) -> str: - pass + pass # pragma: no cover class _SnowflakeSecretsServer(_SnowflakeSecrets): From febd04c70d0cee69572e4cb535bd8a19edf39dc8 Mon Sep 17 00:00:00 2001 From: Tanmay Mehta Date: Fri, 22 May 2026 22:48:45 +0000 Subject: [PATCH 4/5] addressed comments --- tests/integ/conftest.py | 3 ++- tests/integ/test_secrets.py | 34 ++++++++++++++++++++++------------ 2 files changed, 24 insertions(+), 13 deletions(-) diff --git a/tests/integ/conftest.py b/tests/integ/conftest.py index 9c845b8e14..a113b2a3d7 100644 --- a/tests/integ/conftest.py +++ b/tests/integ/conftest.py @@ -161,7 +161,8 @@ def set_up_external_access_integration_resources( ).collect() session.sql( f""" - CREATE IF NOT EXISTS EXTERNAL ACCESS INTEGRATION {integration4} + CREATE EXTERNAL ACCESS INTEGRATION IF NOT EXISTS {integration4} + ALLOWED_NETWORK_RULES = ({rule1}) ALLOWED_AUTHENTICATION_SECRETS = ({key4}) ENABLED = true; """ diff --git a/tests/integ/test_secrets.py b/tests/integ/test_secrets.py index c8183c7c50..a4779cd073 100644 --- a/tests/integ/test_secrets.py +++ b/tests/integ/test_secrets.py @@ -158,14 +158,15 @@ def get_secret(): reason="Secret API is only supported on Snowflake server environment", ) def test_get_wif_token_udf(session, db_parameters): - def get_wif(): - token = get_wif_token("cred", db_parameters["wif_audience"]) - return len(token) > 0 - try: + wif_audience = db_parameters["wif_audience"] + + def get_wif(): + return get_wif_token("cred", wif_audience) + get_wif_udf = session.udf.register( get_wif, - return_type=BooleanType(), + return_type=StringType(), packages=["snowflake-snowpark-python"], external_access_integrations=[ db_parameters["external_access_integration4"] @@ -173,7 +174,12 @@ def get_wif(): secrets={"cred": f"{db_parameters['external_access_key4']}"}, ) df = session.create_dataframe([[1], [2]]).to_df("x") - Utils.check_answer(df.select(get_wif_udf()), [Row(True), Row(True)]) + rows = df.select(get_wif_udf()).collect() + for row in rows: + token = row[0] + assert ( + isinstance(token, str) and len(token.split(".")) == 3 + ), f"expected JWT-shaped token (header.payload.signature), got {token!r}" except KeyError: pytest.skip("External Access Integration is not supported on the deployment.") @@ -183,14 +189,15 @@ def get_wif(): reason="Secret API is only supported on Snowflake server environment", ) def test_get_wif_token_sproc(session, db_parameters): - def get_wif_in_sproc(session_): - token = get_wif_token("cred", db_parameters["wif_audience"]) - return len(token) > 0 - try: + wif_audience = db_parameters["wif_audience"] + + def get_wif_in_sproc(session_): + return get_wif_token("cred", wif_audience) + get_wif_sp = session.sproc.register( get_wif_in_sproc, - return_type=BooleanType(), + return_type=StringType(), packages=["snowflake-snowpark-python"], external_access_integrations=[ db_parameters["external_access_integration4"] @@ -198,7 +205,10 @@ def get_wif_in_sproc(session_): secrets={"cred": f"{db_parameters['external_access_key4']}"}, anonymous=True, ) - assert get_wif_sp() + token = get_wif_sp() + assert ( + isinstance(token, str) and len(token.split(".")) == 3 + ), f"expected JWT-shaped token (header.payload.signature), got {token!r}" except KeyError: pytest.skip("External Access Integration is not supported on the deployment.") From c8b222b634dcd8a20aa7cee1e2289e67a2c31ffc Mon Sep 17 00:00:00 2001 From: Tanmay Mehta Date: Wed, 27 May 2026 18:59:20 +0000 Subject: [PATCH 5/5] address comment for changelog --- CHANGELOG.md | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2ed988dcf2..6dd3c2fe06 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,10 @@ ### Snowpark Python API Updates +#### New Features + +- Added `get_wif_token` to `snowflake.snowpark.secrets` for workload identity federation tokens on the Snowflake server (not available in SPCS file-based secret environments). + #### Documentation - Clarified that the JDBC driver JAR referenced via `udtf_configs.imports` in `DataFrameReader.jdbc()` must be downloaded from the database vendor and uploaded to a Snowflake stage. @@ -38,7 +42,6 @@ #### New Features - Added `artifact_repository` support to `udtf_configs` in `session.read.dbapi()`, enabling users to specify a custom artifact repository (e.g. PyPI) for packages used by the internal UDTF during distributed ingestion. -- Added `get_wif_token` to `snowflake.snowpark.secrets` for workload identity federation tokens on the Snowflake server (not available in SPCS file-based secret environments). #### Bug Fixes