diff --git a/CHANGELOG.md b/CHANGELOG.md index eeb4140527..d54d009514 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,13 @@ #### New Features +- Added a new module `snowflake.snowpark.secrets` that provides Python wrappers for accessing Snowflake Secrets within Python UDFs and stored procedures that execute inside Snowflake. + - `get_generic_secret_string` + - `get_oauth_access_token` + - `get_secret_type` + - `get_username_password` + - `get_cloud_provider_token` + ### Snowpark pandas API Updates #### New Features diff --git a/docs/source/snowpark/index.rst b/docs/source/snowpark/index.rst index 9b0b354b1e..e7de9022d4 100644 --- a/docs/source/snowpark/index.rst +++ b/docs/source/snowpark/index.rst @@ -20,6 +20,7 @@ Snowpark APIs udf udaf udtf + secrets observability files catalog diff --git a/docs/source/snowpark/secrets.rst b/docs/source/snowpark/secrets.rst new file mode 100644 index 0000000000..0c650171fe --- /dev/null +++ b/docs/source/snowpark/secrets.rst @@ -0,0 +1,24 @@ +============================= +Snowpark Secrets +============================= + +.. currentmodule:: snowflake.snowpark.secrets + +.. rubric:: Classes + +.. autosummary:: + :toctree: api/ + + UsernamePassword + CloudProviderToken + +.. rubric:: Functions + +.. autosummary:: + :toctree: api/ + + get_generic_secret_string + get_oauth_access_token + get_secret_type + get_username_password + get_cloud_provider_token diff --git a/src/snowflake/snowpark/__init__.py b/src/snowflake/snowpark/__init__.py index bf13c6700a..292dc8f01a 100644 --- a/src/snowflake/snowpark/__init__.py +++ b/src/snowflake/snowpark/__init__.py @@ -7,7 +7,7 @@ Contains core classes of Snowpark. """ -# types, udf, functions, exceptions still use its own modules +# types, udf, functions, exceptions, secrets still use its own modules __all__ = [ "Column", diff --git a/src/snowflake/snowpark/secrets.py b/src/snowflake/snowpark/secrets.py new file mode 100644 index 0000000000..56f168f5ce --- /dev/null +++ b/src/snowflake/snowpark/secrets.py @@ -0,0 +1,127 @@ +# +# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved. +# +from snowflake.snowpark._internal.utils import publicapi + +# Reference for Python API for Secret Access: +# https://docs.snowflake.com/en/developer-guide/external-network-access/secret-api-reference#python-api-for-secret-access + + +class UsernamePassword: + def __init__(self, username, password) -> None: + self.username = username + self.password = password + + +class CloudProviderToken: + def __init__(self, id, key, token) -> None: + self.access_key_id = id + self.secret_access_key = key + self.token = token + + +@publicapi +def get_generic_secret_string(secret_name: str) -> str: + """Get a generic token string from Snowflake. + Note: + Require a Snowflake environment with generic secret strings configured + Returns: + The secret value as a string. + Raises: + NotImplementedError: If the _snowflake module cannot be imported. + """ + try: + import _snowflake + + return _snowflake.get_generic_secret_string(secret_name) + except ImportError: + raise NotImplementedError( + "Cannot import _snowflake module. Secret API is only supported on Snowflake server environment." + ) + + +@publicapi +def get_oauth_access_token(secret_name: str) -> str: + """Get an OAuth2 access token from Snowflake. + Note: + Require a Snowflake environment with OAuth secrets configured + Returns: + The OAuth2 access token as a string. + Raises: + NotImplementedError: If the _snowflake module cannot be imported. + """ + try: + import _snowflake + + return _snowflake.get_oauth_access_token(secret_name) + except ImportError: + raise NotImplementedError( + "Cannot import _snowflake module. Secret API is only supported on Snowflake server environment." + ) + + +@publicapi +def get_secret_type(secret_name: str) -> str: + """Get the type of a secret from Snowflake. + Note: + Require a Snowflake environment with secrets configured + Returns: + The type of the secret as a string. + Raises: + NotImplementedError: If the _snowflake module cannot be imported. + """ + try: + import _snowflake + + return str(_snowflake.get_secret_type(secret_name)) + except ImportError: + raise NotImplementedError( + "Cannot import _snowflake module. Secret API is only supported on Snowflake server environment." + ) + + +@publicapi +def get_username_password(secret_name: str) -> UsernamePassword: + """Get a username and password secret from Snowflake. + Note: + Require a Snowflake environment with username/password secrets configured + Returns: + UsernamePassword: An object with attributes ``username`` and ``password``. + Raises: + NotImplementedError: If the _snowflake module cannot be imported. + """ + try: + import _snowflake + + secret_object = _snowflake.get_username_password(secret_name) + return UsernamePassword(secret_object.username, secret_object.password) + except ImportError: + raise NotImplementedError( + "Cannot import _snowflake module. Secret API is only supported on Snowflake server environment." + ) + + +@publicapi +def get_cloud_provider_token(secret_name: str) -> CloudProviderToken: + """Get a cloud provider token secret from Snowflake. + Note: + Require a Snowflake environment with cloud provider secrets configured + Returns: + CloudProviderToken: An object with attributes ``access_key_id``, + ``secret_access_key``, and ``token``. + Raises: + NotImplementedError: If the _snowflake module cannot be imported. + """ + try: + import _snowflake + + secret_object = _snowflake.get_cloud_provider_token(secret_name) + return CloudProviderToken( + secret_object.access_key_id, + secret_object.secret_access_key, + secret_object.token, + ) + except ImportError: + raise NotImplementedError( + "Cannot import _snowflake module. Secret API is only supported on Snowflake server environment." + ) diff --git a/tests/integ/conftest.py b/tests/integ/conftest.py index 2ff90e2494..dd20b9c261 100644 --- a/tests/integ/conftest.py +++ b/tests/integ/conftest.py @@ -62,7 +62,16 @@ def print_help() -> None: def set_up_external_access_integration_resources( - session, rule1, rule2, key1, key2, integration1, integration2 + session, + rule1, + rule2, + rule3, + key1, + key2, + key3, + integration1, + integration2, + integration3, ): try: # IMPORTANT SETUP NOTES: the test role needs to be granted the creation privilege @@ -87,6 +96,14 @@ def set_up_external_access_integration_resources( ).collect() session.sql( f""" + CREATE IF NOT EXISTS NETWORK RULE {rule3} + MODE = EGRESS + TYPE = HOST_PORT + VALUE_LIST = ('www.amazon.com'); + """ + ).collect() + session.sql( + f""" CREATE IF NOT EXISTS SECRET {key1} TYPE = GENERIC_STRING SECRET_STRING = 'replace-with-your-api-key'; @@ -101,6 +118,14 @@ def set_up_external_access_integration_resources( ).collect() session.sql( f""" + CREATE IF NOT EXISTS SECRET {key3} + TYPE = PASSWORD + USERNAME = 'replace-with-your-username'; + PASSWORD = 'replace-with-your-password'; + """ + ).collect() + session.sql( + f""" CREATE IF NOT EXISTS EXTERNAL ACCESS INTEGRATION {integration1} ALLOWED_NETWORK_RULES = ({rule1}) ALLOWED_AUTHENTICATION_SECRETS = ({key1}) @@ -113,14 +138,25 @@ def set_up_external_access_integration_resources( ALLOWED_NETWORK_RULES = ({rule2}) ALLOWED_AUTHENTICATION_SECRETS = ({key2}) ENABLED = true; + """ + ).collect() + session.sql( + f""" + CREATE IF NOT EXISTS EXTERNAL ACCESS INTEGRATION {integration3} + ALLOWED_NETWORK_RULES = ({rule3}) + ALLOWED_AUTHENTICATION_SECRETS = ({key3}) + ENABLED = true; """ ).collect() CONNECTION_PARAMETERS["external_access_rule1"] = rule1 CONNECTION_PARAMETERS["external_access_rule2"] = rule2 + CONNECTION_PARAMETERS["external_access_rule3"] = rule3 CONNECTION_PARAMETERS["external_access_key1"] = key1 CONNECTION_PARAMETERS["external_access_key2"] = key2 + CONNECTION_PARAMETERS["external_access_key3"] = key3 CONNECTION_PARAMETERS["external_access_integration1"] = integration1 CONNECTION_PARAMETERS["external_access_integration2"] = integration2 + CONNECTION_PARAMETERS["external_access_integration3"] = integration3 except SnowparkSQLException: # GCP currently does not support external access integration # we can remove the exception once the integration is available on GCP @@ -142,10 +178,13 @@ def set_up_external_access_integration_resources( def clean_up_external_access_integration_resources(): CONNECTION_PARAMETERS.pop("external_access_rule1", None) CONNECTION_PARAMETERS.pop("external_access_rule2", None) + CONNECTION_PARAMETERS.pop("external_access_rule3", None) 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_integration1", None) CONNECTION_PARAMETERS.pop("external_access_integration2", None) + CONNECTION_PARAMETERS.pop("external_access_integration3", None) def set_up_dataframe_processor_parameters( @@ -264,10 +303,13 @@ def session( set_ast_state(AstFlagSource.TEST, ast_enabled) rule1 = "snowpark_python_test_rule1" rule2 = "snowpark_python_test_rule2" + rule3 = "snowpark_python_test_rule3" key1 = "snowpark_python_test_key1" key2 = "snowpark_python_test_key2" + key3 = "snowpark_python_test_key3" integration1 = "snowpark_python_test_integration1" integration2 = "snowpark_python_test_integration2" + integration3 = "snowpark_python_test_integration3" session = ( Session.builder.configs(db_parameters) @@ -291,7 +333,16 @@ def session( if (RUNNING_ON_GH or RUNNING_ON_JENKINS) and not local_testing_mode: set_up_external_access_integration_resources( - session, rule1, rule2, key1, key2, integration1, integration2 + session, + rule1, + rule2, + rule3, + key1, + key2, + key3, + integration1, + integration2, + integration3, ) if validate_ast: @@ -321,10 +372,13 @@ def profiler_session( ): rule1 = "snowpark_python_profiler_test_rule1" rule2 = "snowpark_python_profiler_test_rule2" + rule3 = "snowpark_python_profiler_test_rule3" key1 = "snowpark_python_profiler_test_key1" key2 = "snowpark_python_profiler_test_key2" + key3 = "snowpark_python_profiler_test_key3" integration1 = "snowpark_python_profiler_test_integration1" integration2 = "snowpark_python_profiler_test_integration2" + integration3 = "snowpark_python_profiler_test_integration3" session = ( Session.builder.configs(db_parameters) .config("local_testing", local_testing_mode) @@ -334,7 +388,16 @@ def profiler_session( session._cte_optimization_enabled = cte_optimization_enabled if RUNNING_ON_GH and not local_testing_mode: set_up_external_access_integration_resources( - session, rule1, rule2, key1, key2, integration1, integration2 + session, + rule1, + rule2, + rule3, + key1, + key2, + key3, + integration1, + integration2, + integration3, ) try: yield session diff --git a/tests/integ/test_secrets.py b/tests/integ/test_secrets.py new file mode 100644 index 0000000000..f67b9b5701 --- /dev/null +++ b/tests/integ/test_secrets.py @@ -0,0 +1,171 @@ +# +# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved. +# + +import pytest +from snowflake.snowpark.secrets import ( + get_generic_secret_string, + get_username_password, + get_secret_type, + get_cloud_provider_token, + get_oauth_access_token, +) +from snowflake.snowpark.types import BooleanType, StringType +from tests.utils import IS_NOT_ON_GITHUB, RUNNING_ON_JENKINS, IS_IN_STORED_PROC, Utils +from snowflake.snowpark import Row +from snowflake.snowpark.exceptions import SnowparkSQLException + + +@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_generic_secret_string_sproc(session, db_parameters): + def get_secret_string(session_): + if get_generic_secret_string("cred") == "replace-with-your-api-key": + return True + return False + + try: + get_secret_string_sp = session.sproc.register( + get_secret_string, + return_type=BooleanType(), + packages=["snowflake-snowpark-python"], + external_access_integrations=[ + db_parameters["external_access_integration1"] + ], + secrets={"cred": f"{db_parameters['external_access_key1']}"}, + ) + result = get_secret_string_sp() + assert result + 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_generic_secret_string_udf(session, db_parameters): + def get_secret_string(): + if get_generic_secret_string("cred") == "replace-with-your-api-key_2": + return True + return False + + try: + get_secret_string_udf = session.udf.register( + get_secret_string, + return_type=BooleanType(), + packages=["snowflake-snowpark-python"], + external_access_integrations=[ + db_parameters["external_access_integration2"] + ], + secrets={"cred": f"{db_parameters['external_access_key2']}"}, + ) + df = session.create_dataframe([[1], [2]]).to_df("x") + Utils.check_answer(df.select(get_secret_string_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_username_password_udf(session, db_parameters): + def get_secret_username_password(): + creds = get_username_password("cred") + if ( + creds.username == "replace-with-your-username" + and creds.password == "replace-with-your-password" + ): + return True + return False + + try: + get_secret_udf = session.udf.register( + get_secret_username_password, + return_type=BooleanType(), + packages=["snowflake-snowpark-python"], + external_access_integrations=[ + db_parameters["external_access_integration2"] + ], + secrets={"cred": f"{db_parameters['external_access_key2']}"}, + ) + df = session.create_dataframe([[1], [2]]).to_df("x") + Utils.check_answer(df.select(get_secret_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_secret_type_sproc(session, db_parameters): + def get_type(session_): + t_str = get_secret_type("cred_str") + t_pwd = get_secret_type("cred_pwd") + return f"{t_str},{t_pwd}" + + try: + get_type_sp = session.sproc.register( + get_type, + return_type=StringType(), + packages=["snowflake-snowpark-python"], + external_access_integrations=[ + db_parameters["external_access_integration1"], + db_parameters["external_access_integration3"], + ], + secrets={ + "cred_str": f"{db_parameters['external_access_key1']}", + "cred_pwd": f"{db_parameters['external_access_key3']}", + }, + anonymous=True, + ) + result = get_type_sp() + parts = result.split(",") + assert parts[0] == "GENERIC_STRING" + assert parts[1] == "PASSWORD" + 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_nonexistent_secret(session, db_parameters): + def get_secret(): + get_generic_secret_string("cred") + return False + + with pytest.raises(SnowparkSQLException): + session.udf.register( + get_secret, + return_type=BooleanType(), + packages=["snowflake-snowpark-python"], + external_access_integrations=[ + db_parameters["external_access_integration2"] + ], + secrets={"cred": "nonexistent_secret"}, + ) + + +@pytest.mark.skipif( + IS_IN_STORED_PROC, + reason="Run only outside Snowflake server to validate NotImplementedError", +) +def test_secrets_import_error(): + # [SNOW-2324796] Phase 1 relies on Snowflake server environment + # Remove this test once secrets local development support is complete + with pytest.raises(NotImplementedError): + get_generic_secret_string("s1") + with pytest.raises(NotImplementedError): + get_username_password("p1") + with pytest.raises(NotImplementedError): + get_secret_type("t1") + with pytest.raises(NotImplementedError): + get_cloud_provider_token("c1") + with pytest.raises(NotImplementedError): + get_oauth_access_token("o1") diff --git a/tests/unit/test_secrets.py b/tests/unit/test_secrets.py new file mode 100644 index 0000000000..1f1e75d073 --- /dev/null +++ b/tests/unit/test_secrets.py @@ -0,0 +1,68 @@ +# +# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved. +# + +import pytest +from types import SimpleNamespace +from unittest.mock import patch +from snowflake.snowpark.secrets import ( + get_generic_secret_string, + get_oauth_access_token, + get_secret_type, + get_username_password, + get_cloud_provider_token, + UsernamePassword, + CloudProviderToken, +) + + +def _build_fake_snowflake_module() -> object: + fake_username_password = SimpleNamespace(username="user1", password="pass1") + fake_cloud_token = SimpleNamespace( + access_key_id="AKIA_TEST", + secret_access_key="SECRET_TEST", + token="STS_TOKEN_TEST", + ) + return SimpleNamespace( + get_generic_secret_string=lambda secret_name: f"generic:{secret_name}", + get_oauth_access_token=lambda secret_name: f"oauth:{secret_name}", + 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, + ) + + +def test_secrets_mock_server_paths(): + fake_module = _build_fake_snowflake_module() + with patch.dict("sys.modules", {"_snowflake": fake_module}): + assert get_generic_secret_string("s1") == "generic:s1" + assert get_oauth_access_token("o1") == "oauth:o1" + assert get_secret_type("t1") == "PASSWORD" + + creds = get_username_password("p1") + assert isinstance(creds, UsernamePassword) + assert creds.username == "user1" + assert creds.password == "pass1" + + cloud = get_cloud_provider_token("c1") + assert isinstance(cloud, CloudProviderToken) + assert cloud.access_key_id == "AKIA_TEST" + assert cloud.secret_access_key == "SECRET_TEST" + assert cloud.token == "STS_TOKEN_TEST" + + +def test_secrets_import_error_paths(): + with patch.dict("sys.modules", {"_snowflake": None}): + if "_snowflake" in __import__("sys").modules: + del __import__("sys").modules["_snowflake"] + + with pytest.raises(NotImplementedError): + get_generic_secret_string("s1") + with pytest.raises(NotImplementedError): + get_oauth_access_token("o1") + with pytest.raises(NotImplementedError): + get_secret_type("t1") + with pytest.raises(NotImplementedError): + get_username_password("p1") + with pytest.raises(NotImplementedError): + get_cloud_provider_token("c1")