Prepare for Workload Identity Federation (WIF) GA (#2368)

Author: sfc-gh-pmansour

DESCRIPTION.md

@@ -8,6 +8,9 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
 
 # Release Notes
 - v3.16.1(TBD)
+  - Added new authentication methods support for Workload Identity Federation (WIF).
+    - Added the `WORKLOAD_IDENTITY` value for authenticator type.
+    - Added the `workload_identity_provider` and `workload_identity_entra_resource` parameters.
   - Added in-band OCSP exception telemetry.
   - Added `APPLICATION_PATH` within `CLIENT_ENVIRONMENT` to distinguish between multiple scripts using the PythonConnector in the same environment.
   - Disabled token caching for OAuth Client Credentials authentication

ci/container/test_authentication.sh

@@ -13,7 +13,6 @@ export SNOWFLAKE_AUTH_TEST_PRIVATE_KEY_PATH=./.github/workflows/parameters/priva
 export SNOWFLAKE_AUTH_TEST_INVALID_PRIVATE_KEY_PATH=./.github/workflows/parameters/private/rsa_keys/rsa_key_invalid.p8
 
 export SF_OCSP_TEST_MODE=true
-export SF_ENABLE_EXPERIMENTAL_AUTHENTICATION=true
 export RUN_AUTH_TESTS=true
 export AUTHENTICATION_TESTS_ENV="docker"
 export PYTHONPATH=$SOURCE_ROOT

ci/wif/test_wif.sh

@@ -3,7 +3,6 @@
 set -o pipefail
 
 export SF_OCSP_TEST_MODE=true
-export SF_ENABLE_EXPERIMENTAL_AUTHENTICATION=true
 export RUN_WIF_TESTS=true
 
 /opt/python/cp39-cp39/bin/python -m pip install --break-system-packages -e .

src/snowflake/connector/connection.py

@@ -59,7 +59,6 @@
     _CONNECTIVITY_ERR_MSG,
     _DOMAIN_NAME_MAP,
     _OAUTH_DEFAULT_SCOPE,
-    ENV_VAR_EXPERIMENTAL_AUTHENTICATION,
     ENV_VAR_PARTNER,
     PARAMETER_AUTOCOMMIT,
     PARAMETER_CLIENT_PREFETCH_THREADS,
@@ -88,7 +87,6 @@
 from .direct_file_operation_utils import FileOperationParser, StreamDownloader
 from .errorcode import (
     ER_CONNECTION_IS_CLOSED,
-    ER_EXPERIMENTAL_AUTHENTICATION_NOT_SUPPORTED,
     ER_FAILED_PROCESSING_PYFORMAT,
     ER_FAILED_PROCESSING_QMARK,
     ER_FAILED_TO_CONNECT_TO_DB,
@@ -1329,8 +1327,6 @@ def __open_connection(self):
                     self._token, self._external_session_id
                 )
             elif self._authenticator == WORKLOAD_IDENTITY_AUTHENTICATOR:
-                self._check_experimental_authentication_flag()
-
                 if isinstance(self._workload_identity_provider, str):
                     self._workload_identity_provider = AttestationProvider.from_string(
                         self._workload_identity_provider
@@ -2293,18 +2289,6 @@ def is_valid(self) -> bool:
             logger.debug("session could not be validated due to exception: %s", e)
             return False
 
-    def _check_experimental_authentication_flag(self) -> None:
-        if os.getenv(ENV_VAR_EXPERIMENTAL_AUTHENTICATION, "false").lower() != "true":
-            Error.errorhandler_wrapper(
-                self,
-                None,
-                ProgrammingError,
-                {
-                    "msg": f"Please set the '{ENV_VAR_EXPERIMENTAL_AUTHENTICATION}' environment variable true to use the '{self._authenticator}' authenticator.",
-                    "errno": ER_EXPERIMENTAL_AUTHENTICATION_NOT_SUPPORTED,
-                },
-            )
-
     @staticmethod
     def _detect_application() -> None | str:
         if ENV_VAR_PARTNER in os.environ.keys():

src/snowflake/connector/constants.py

@@ -428,7 +428,6 @@ class IterUnit(Enum):
 # TODO: all env variables definitions should be here
 ENV_VAR_PARTNER = "SF_PARTNER"
 ENV_VAR_TEST_MODE = "SNOWFLAKE_TEST_MODE"
-ENV_VAR_EXPERIMENTAL_AUTHENTICATION = "SF_ENABLE_EXPERIMENTAL_AUTHENTICATION"  # Needed to enable new strong auth features during the private preview.
 
 
 _DOMAIN_NAME_MAP = {_DEFAULT_HOSTNAME_TLD: "GLOBAL", _CHINA_HOSTNAME_TLD: "CHINA"}

src/snowflake/connector/errorcode.py

@@ -33,6 +33,7 @@
 ER_OAUTH_SERVER_TIMEOUT = 251016
 ER_INVALID_WIF_SETTINGS = 251017
 ER_WIF_CREDENTIALS_NOT_FOUND = 251018
+# not used but keep here to reserve errno
 ER_EXPERIMENTAL_AUTHENTICATION_NOT_SUPPORTED = 251019
 ER_NO_CLIENT_SECRET = 251020
 

test/unit/test_connection.py

@@ -639,17 +639,6 @@ def test_cannot_set_dependent_params_without_wlid_authenticator(
     )
 
 
-def test_cannot_set_wlid_authenticator_without_env_variable(mock_post_requests):
-    with pytest.raises(ProgrammingError) as excinfo:
-        snowflake.connector.connect(
-            account="account", authenticator="WORKLOAD_IDENTITY"
-        )
-    assert (
-        "Please set the 'SF_ENABLE_EXPERIMENTAL_AUTHENTICATION' environment variable true to use the 'WORKLOAD_IDENTITY' authenticator"
-        in str(excinfo.value)
-    )
-
-
 @pytest.mark.parametrize(
     "provider_param",
     [
@@ -665,7 +654,6 @@ def test_workload_identity_provider_is_required_for_wif_authenticator(
         m.setattr(
             "snowflake.connector.SnowflakeConnection._authenticate", lambda *_: None
         )
-        m.setenv("SF_ENABLE_EXPERIMENTAL_AUTHENTICATION", "true")
 
         with pytest.raises(ProgrammingError) as excinfo:
             snowflake.connector.connect(
@@ -701,7 +689,6 @@ def test_connection_params_are_plumbed_into_authbyworkloadidentity(
         m.setattr(
             "snowflake.connector.SnowflakeConnection._authenticate", lambda *_: None
         )
-        m.setenv("SF_ENABLE_EXPERIMENTAL_AUTHENTICATION", "true")
 
         conn = snowflake.connector.connect(
             account="my_account_1",
@@ -743,7 +730,6 @@ def test_toml_connection_params_are_plumbed_into_authbyworkloadidentity(
         m.setattr(
             "snowflake.connector.SnowflakeConnection._authenticate", lambda *_: None
         )
-        m.setenv("SF_ENABLE_EXPERIMENTAL_AUTHENTICATION", "true")
 
         conn = snowflake.connector.connect(connections_file_path=connections_file)
         assert conn.auth_class.provider == AttestationProvider.OIDC
@@ -762,7 +748,6 @@ def test_single_use_refresh_tokens_option_is_plumbed_into_authbyauthcode(
         m.setattr(
             "snowflake.connector.SnowflakeConnection._authenticate", lambda *_: None
         )
-        m.setenv("SF_ENABLE_EXPERIMENTAL_AUTHENTICATION", "true")
 
         conn = snowflake.connector.connect(
             account="my_account_1",