Snow 2324796 secrets api migration 1 (#3781)

Author: sfc-gh-mayliu

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

docs/source/snowpark/index.rst

@@ -20,6 +20,7 @@ Snowpark APIs
    udf
    udaf
    udtf
+   secrets
    observability
    files
    catalog

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

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",

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."
+        )

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