Clarify error messages detected during WIF training (#2469)

Author: sfc-gh-pczajka

src/snowflake/connector/connection.py

@@ -1214,6 +1214,7 @@ def __open_connection(self):
                     raise TypeError("auth_class must be a child class of AuthByKeyPair")
                     # TODO: add telemetry for custom auth
                 self.auth_class = self.auth_class
+            # match authentivator - validation happens in __config
             elif self._authenticator == DEFAULT_AUTHENTICATOR:
                 self.auth_class = AuthByDefault(
                     password=self._password,
@@ -1468,20 +1469,30 @@ def __config(self, **kwargs):
         # type to be the same as the custom auth class
         if self._auth_class:
             self._authenticator = self._auth_class.type_.value
-
-        if self._authenticator:
-            # Only upper self._authenticator if it is a non-okta link
+        elif self._authenticator:
+            # Validate authenticator and convert it to uppercase if it is a non-okta link
             auth_tmp = self._authenticator.upper()
-            if auth_tmp in [  # Non-okta authenticators
+            if auth_tmp in [
                 DEFAULT_AUTHENTICATOR,
                 EXTERNAL_BROWSER_AUTHENTICATOR,
                 KEY_PAIR_AUTHENTICATOR,
                 OAUTH_AUTHENTICATOR,
+                OAUTH_AUTHORIZATION_CODE,
+                OAUTH_CLIENT_CREDENTIALS,
                 USR_PWD_MFA_AUTHENTICATOR,
                 WORKLOAD_IDENTITY_AUTHENTICATOR,
+                PROGRAMMATIC_ACCESS_TOKEN,
                 PAT_WITH_EXTERNAL_SESSION,
             ]:
                 self._authenticator = auth_tmp
+            elif auth_tmp.startswith("HTTPS://"):
+                # okta authenticator link
+                pass
+            else:
+                raise ProgrammingError(
+                    msg=f"Unknown authenticator: {self._authenticator}",
+                    errno=ER_INVALID_VALUE,
+                )
 
         # read OAuth token from
         token_file_path = kwargs.get("token_file_path")

src/snowflake/connector/wif_util.py

@@ -13,7 +13,7 @@
 from botocore.awsrequest import AWSRequest
 from botocore.utils import InstanceMetadataRegionFetcher
 
-from .errorcode import ER_WIF_CREDENTIALS_NOT_FOUND
+from .errorcode import ER_INVALID_WIF_SETTINGS, ER_WIF_CREDENTIALS_NOT_FOUND
 from .errors import ProgrammingError
 from .session_manager import SessionManager
 
@@ -38,7 +38,13 @@ class AttestationProvider(Enum):
     @staticmethod
     def from_string(provider: str) -> AttestationProvider:
         """Converts a string to a strongly-typed enum value of AttestationProvider."""
-        return AttestationProvider[provider.upper()]
+        try:
+            return AttestationProvider[provider.upper()]
+        except KeyError:
+            raise ProgrammingError(
+                msg=f"Unknown workload_identity_provider: '{provider}'. Expected one of: {', '.join(AttestationProvider.all_string_values())}",
+                errno=ER_INVALID_WIF_SETTINGS,
+            )
 
     @staticmethod
     def all_string_values() -> list[str]:
@@ -65,7 +71,13 @@ def extract_iss_and_sub_without_signature_verification(jwt_str: str) -> tuple[st
 
     Any errors during token parsing will be bubbled up. Missing 'iss' or 'sub' claims will also raise an error.
     """
-    claims = jwt.decode(jwt_str, options={"verify_signature": False})
+    try:
+        claims = jwt.decode(jwt_str, options={"verify_signature": False})
+    except jwt.InvalidTokenError as e:
+        raise ProgrammingError(
+            msg=f"Invalid JWT token: {e}",
+            errno=ER_INVALID_WIF_SETTINGS,
+        )
 
     if not ("iss" in claims and "sub" in claims):
         raise ProgrammingError(
@@ -179,14 +191,20 @@ def create_gcp_attestation(
 
     If the application isn't running on GCP or no credentials were found, raises an error.
     """
-    res = session_manager.request(
-        method="GET",
-        url=f"http://169.254.169.254/computeMetadata/v1/instance/service-accounts/default/identity?audience={SNOWFLAKE_AUDIENCE}",
-        headers={
-            "Metadata-Flavor": "Google",
-        },
-    )
-    res.raise_for_status()
+    try:
+        res = session_manager.request(
+            method="GET",
+            url=f"http://169.254.169.254/computeMetadata/v1/instance/service-accounts/default/identity?audience={SNOWFLAKE_AUDIENCE}",
+            headers={
+                "Metadata-Flavor": "Google",
+            },
+        )
+        res.raise_for_status()
+    except Exception as e:
+        raise ProgrammingError(
+            msg=f"Error fetching GCP metadata: {e}. Ensure the application is running on GCP.",
+            errno=ER_WIF_CREDENTIALS_NOT_FOUND,
+        )
 
     jwt_str = res.content.decode("utf-8")
     _, subject = extract_iss_and_sub_without_signature_verification(jwt_str)
@@ -230,12 +248,18 @@ def create_azure_attestation(
     if managed_identity_client_id:
         query_params += f"&client_id={managed_identity_client_id}"
 
-    res = session_manager.request(
-        method="GET",
-        url=f"{url_without_query_string}?{query_params}",
-        headers=headers,
-    )
-    res.raise_for_status()
+    try:
+        res = session_manager.request(
+            method="GET",
+            url=f"{url_without_query_string}?{query_params}",
+            headers=headers,
+        )
+        res.raise_for_status()
+    except Exception as e:
+        raise ProgrammingError(
+            msg=f"Error fetching Azure metadata: {e}. Ensure the application is running on Azure.",
+            errno=ER_WIF_CREDENTIALS_NOT_FOUND,
+        )
 
     jwt_str = res.json().get("access_token")
     if not jwt_str:

test/auth/authorization_parameters.py

@@ -79,7 +79,7 @@ def get_base_connection_parameters(self) -> dict[str, Union[str, bool, int]]:
 
     def get_key_pair_connection_parameters(self):
         config = self.basic_config.copy()
-        config["authenticator"] = "KEY_PAIR_AUTHENTICATOR"
+        config["authenticator"] = "SNOWFLAKE_JWT"
         config["user"] = _get_env_variable("SNOWFLAKE_AUTH_TEST_BROWSER_USER")
 
         return config

test/unit/test_auth_keypair.py

@@ -3,6 +3,7 @@
 
 from unittest.mock import Mock, PropertyMock, patch
 
+import pytest
 from cryptography.hazmat.backends import default_backend
 from cryptography.hazmat.primitives import serialization
 from cryptography.hazmat.primitives.asymmetric import rsa
@@ -36,7 +37,8 @@ def _mock_auth_key_pair_rest_response(url, headers, body, **kwargs):
     return _mock_auth_key_pair_rest_response
 
 
-def test_auth_keypair():
+@pytest.mark.parametrize("authenticator", ["SNOWFLAKE_JWT", "snowflake_jwt"])
+def test_auth_keypair(authenticator):
     """Simple Key Pair test."""
     private_key_der, public_key_der_encoded = generate_key_pair(2048)
     application = "testapplication"
@@ -45,7 +47,7 @@ def test_auth_keypair():
     auth_instance = AuthByKeyPair(private_key=private_key_der)
     auth_instance._retry_ctx.set_start_time()
     auth_instance.handle_timeout(
-        authenticator="SNOWFLAKE_JWT",
+        authenticator=authenticator,
         service_name=None,
         account=account,
         user=user,

test/unit/test_auth_mfa.py

@@ -1,9 +1,14 @@
 from unittest import mock
 
+import pytest
+
 from snowflake.connector import connect
 
 
-def test_mfa_token_cache():
+@pytest.mark.parametrize(
+    "authenticator", ["USERNAME_PASSWORD_MFA", "username_password_mfa"]
+)
+def test_mfa_token_cache(authenticator):
     with mock.patch(
         "snowflake.connector.network.SnowflakeRestful.fetch",
     ):
@@ -14,7 +19,7 @@ def test_mfa_token_cache():
                 account="account",
                 user="user",
                 password="password",
-                authenticator="username_password_mfa",
+                authenticator=authenticator,
                 client_store_temporary_credential=True,
                 client_request_mfa_token=True,
             ):
@@ -40,7 +45,7 @@ def test_mfa_token_cache():
                     account="account",
                     user="user",
                     password="password",
-                    authenticator="username_password_mfa",
+                    authenticator=authenticator,
                     client_store_temporary_credential=True,
                     client_request_mfa_token=True,
                 ):

test/unit/test_auth_oauth.py

@@ -5,6 +5,7 @@
     from snowflake.connector.auth import AuthByOAuth
 except ImportError:
     from snowflake.connector.auth_oauth import AuthByOAuth
+import pytest
 
 
 def test_auth_oauth():
@@ -15,3 +16,38 @@ def test_auth_oauth():
     auth.update_body(body)
     assert body["data"]["TOKEN"] == token, body
     assert body["data"]["AUTHENTICATOR"] == "OAUTH", body
+
+
+@pytest.mark.parametrize("authenticator", ["oauth", "OAUTH"])
+def test_oauth_authenticator_is_case_insensitive(monkeypatch, authenticator):
+    """Test that oauth authenticator is case insensitive."""
+    import snowflake.connector
+
+    def mock_post_request(self, url, headers, json_body, **kwargs):
+        return {
+            "success": True,
+            "message": None,
+            "data": {
+                "token": "TOKEN",
+                "masterToken": "MASTER_TOKEN",
+                "idToken": None,
+                "parameters": [{"name": "SERVICE_NAME", "value": "FAKE_SERVICE_NAME"}],
+            },
+        }
+
+    monkeypatch.setattr(
+        snowflake.connector.network.SnowflakeRestful, "_post_request", mock_post_request
+    )
+
+    # Create connection with oauth authenticator - OAuth requires a token parameter
+    conn = snowflake.connector.connect(
+        user="testuser",
+        account="testaccount",
+        authenticator=authenticator,
+        token="test_oauth_token",  # OAuth authentication requires a token
+    )
+
+    # Verify that the auth_class is an instance of AuthByOAuth
+    assert isinstance(conn.auth_class, AuthByOAuth)
+
+    conn.close()

test/unit/test_auth_oauth_auth_code.py

@@ -211,3 +211,50 @@ def assert_initialized_correctly() -> None:
             assert_initialized_correctly()
     else:
         assert_initialized_correctly()
+
+
+@pytest.mark.parametrize(
+    "authenticator", ["OAUTH_AUTHORIZATION_CODE", "oauth_authorization_code"]
+)
+def test_oauth_authorization_code_authenticator_is_case_insensitive(
+    monkeypatch, authenticator
+):
+    """Test that OAuth authorization code authenticator is case insensitive."""
+    import snowflake.connector
+
+    def mock_post_request(self, url, headers, json_body, **kwargs):
+        return {
+            "success": True,
+            "message": None,
+            "data": {
+                "token": "TOKEN",
+                "masterToken": "MASTER_TOKEN",
+                "idToken": None,
+                "parameters": [{"name": "SERVICE_NAME", "value": "FAKE_SERVICE_NAME"}],
+            },
+        }
+
+    monkeypatch.setattr(
+        snowflake.connector.network.SnowflakeRestful, "_post_request", mock_post_request
+    )
+
+    # Mock the OAuth authorization flow to avoid opening browser and starting HTTP server
+    def mock_request_tokens(self, **kwargs):
+        # Simulate successful token retrieval
+        return ("mock_access_token", "mock_refresh_token")
+
+    monkeypatch.setattr(AuthByOauthCode, "_request_tokens", mock_request_tokens)
+
+    # Create connection with OAuth authorization code authenticator
+    conn = snowflake.connector.connect(
+        user="testuser",
+        account="testaccount",
+        authenticator=authenticator,
+        oauth_client_id="test_client_id",
+        oauth_client_secret="test_client_secret",
+    )
+
+    # Verify that the auth_class is an instance of AuthByOauthCode
+    assert isinstance(conn.auth_class, AuthByOauthCode)
+
+    conn.close()