SNOW-3338458: Snowpark python support for WIF secret in EA (#4202)

Author: sfc-gh-tmehta

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.

docs/source/snowpark/secrets.rst

@@ -22,3 +22,4 @@ Snowpark Secrets
     get_secret_type
     get_username_password
     get_cloud_provider_token
+    get_wif_token

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  # pragma: no cover
+
 
 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)

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,14 @@ def set_up_external_access_integration_resources(
       ALLOWED_NETWORK_RULES = ({rule3})
       ALLOWED_AUTHENTICATION_SECRETS = ({key3})
       ENABLED = true;
+    """
+        ).collect()
+        session.sql(
+            f"""
+    CREATE EXTERNAL ACCESS INTEGRATION IF NOT EXISTS {integration4}
+      ALLOWED_NETWORK_RULES = ({rule1})
+      ALLOWED_AUTHENTICATION_SECRETS = ({key4})
+      ENABLED = true;
     """
         ).collect()
         CONNECTION_PARAMETERS["external_access_rule1"] = rule1
@@ -156,9 +173,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 +204,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 +338,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 +377,9 @@ def session(
             integration1,
             integration2,
             integration3,
+            key4,
+            integration4,
+            wif_audience,
         )
 
     if validate_ast:
@@ -387,9 +416,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 +441,9 @@ def profiler_session(
             integration1,
             integration2,
             integration3,
+            key4,
+            integration4,
+            wif_audience,
         )
     set_up_test_session_parameters(session, local_testing_mode)
     try:

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
@@ -152,6 +153,66 @@ 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):
+    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=StringType(),
+            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")
+        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.")
+
+
+@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):
+    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=StringType(),
+            packages=["snowflake-snowpark-python"],
+            external_access_integrations=[
+                db_parameters["external_access_integration4"]
+            ],
+            secrets={"cred": f"{db_parameters['external_access_key4']}"},
+            anonymous=True,
+        )
+        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.")
+
+
 @pytest.mark.skipif(
     IS_IN_STORED_PROC,
     reason="Run only outside Snowflake server to validate NotImplementedError",
@@ -169,3 +230,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")

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