SNOW-2872400 Add MFA authentication support to Python wrapper

.gitignore

@@ -15,6 +15,9 @@ parameters*.json
 .DS_Store
 /.tmp
 
+# Local directory for AI prompts / inputs
+.ai/plans
+
 # Python/Cython build artifacts
 python/src/snowflake/connector/_internal/nanoarrow_cpp/ArrowIterator/arrow_stream_iterator.cpp
 python/src/snowflake/connector/_internal/*.so

python/src/snowflake/connector/connection.py

@@ -69,6 +69,18 @@ def __init__(
             port: Port number
             private_key: Private key in bytes, str (base64), or RSAPrivateKey format
             session_parameters: Optional dict of session parameters to set at connection time
+            authenticator: Authentication method. Use ``"USERNAME_PASSWORD_MFA"`` for MFA authentication.
+            passcode: MFA passcode (TOTP one-time code from an authenticator app). Used when
+                ``authenticator="USERNAME_PASSWORD_MFA"`` and ``ext_authn_duo_method="passcode"``.
+            passcode_in_password: If ``True``, the MFA passcode is appended to the password field
+                rather than sent separately. Default ``False``.
+            client_store_temporary_credential: If ``True``, a successfully obtained MFA token is
+                cached in the OS keyring and reused for subsequent connections, avoiding repeated
+                MFA prompts. Default ``False``. The server must have
+                ``ALLOW_CLIENT_MFA_CACHING`` enabled. This also implicitly requests an MFA token
+                from the server (``CLIENT_REQUEST_MFA_TOKEN``).
+            ext_authn_duo_method: DUO Security authentication method. Either ``"push"`` (send a
+                push notification to the registered device) or ``"passcode"`` (use a TOTP code).
             **kwargs: Additional connection parameters
         """
         # paramstyle
@@ -77,6 +89,7 @@ def __init__(
         self._paramstyle = ParamStyle.from_string(paramstyle or default_paramstyle)
 
         kwargs = self._rewrite_private_key_password(kwargs)
+        kwargs = self._rewrite_mfa_params(kwargs)
 
         self.db_api = database_driver_client()
         self.db_handle = self.db_api.database_new(DatabaseNewRequest()).db_handle
@@ -133,7 +146,7 @@ def __init__(
             )
 
         self.db_api.connection_init(ConnectionInitRequest(conn_handle=self.conn_handle, db_handle=self.db_handle))
-        _sensitive_keys = {"password", "private_key"}
+        _sensitive_keys = {"password", "private_key", "passcode"}
         self.kwargs = {k: ("***" if k in _sensitive_keys else v) for k, v in kwargs.items()}
         self._closed = False
         self._messages: list[tuple[type[Exception], dict[str, str | bool]]] = []
@@ -357,6 +370,13 @@ def _rewrite_private_key_password(self, kwargs: ConnectionParameters) -> Connect
             kwargs = {**kwargs, "private_key_password": private_key_file_pwd}
         return kwargs
 
+    def _rewrite_mfa_params(self, kwargs: ConnectionParameters) -> ConnectionParameters:
+        """Translate Python-style MFA parameter names to the keys expected by the Rust core."""
+        passcode_in_password = kwargs.pop("passcode_in_password", None)
+        if passcode_in_password is not None:
+            kwargs = {**kwargs, "passcodeInPassword": passcode_in_password}
+        return kwargs
+
     @property
     def role(self) -> str | None:
         """The current role in use for the session."""

python/tests/config.py

@@ -39,5 +39,8 @@ def get_test_parameters() -> dict[str, Any]:
         "SNOWFLAKE_TEST_PRIVATE_KEY_FILE",
         "SNOWFLAKE_TEST_PRIVATE_KEY_CONTENTS",
         "SNOWFLAKE_TEST_PRIVATE_KEY_PASSWORD",
+        "SNOWFLAKE_TEST_MFA_USER",
+        "SNOWFLAKE_TEST_MFA_PASSWORD",
+        "SNOWFLAKE_TEST_MFA_PASSCODE",
     ]
     return {k: os.environ.get(k) for k in env_vars}

python/tests/e2e/authentication/test_user_password_mfa.py

@@ -0,0 +1,181 @@
+import pytest
+
+from ...compatibility import NEW_DRIVER_ONLY
+from ...config import get_test_parameters
+from .auth_helpers import verify_simple_query_execution
+
+
+# ---------------------------------------------------------------------------
+# Fixtures
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture(scope="module")
+def mfa_params():
+    """
+    Read MFA-specific test parameters from the environment / parameters.json.
+
+    Required keys:
+        SNOWFLAKE_TEST_MFA_USER      – account user that has MFA enabled
+        SNOWFLAKE_TEST_MFA_PASSWORD  – password for that user
+
+    Optional keys:
+        SNOWFLAKE_TEST_MFA_PASSCODE  – current TOTP code (needed for passcode-based tests)
+    """
+    params = get_test_parameters()
+    user = params.get("SNOWFLAKE_TEST_MFA_USER")
+    password = params.get("SNOWFLAKE_TEST_MFA_PASSWORD")
+
+    if not user or not password:
+        pytest.skip("MFA test credentials not configured. Set SNOWFLAKE_TEST_MFA_USER and SNOWFLAKE_TEST_MFA_PASSWORD.")
+
+    return {
+        "user": user,
+        "password": password,
+        "passcode": params.get("SNOWFLAKE_TEST_MFA_PASSCODE"),
+    }
+
+
+@pytest.fixture(scope="module")
+def mfa_passcode(mfa_params):
+    """Skip the test when no TOTP passcode is available."""
+    passcode = mfa_params["passcode"]
+    if not passcode:
+        pytest.skip("No MFA passcode configured. Set SNOWFLAKE_TEST_MFA_PASSCODE.")
+    return passcode
+
+
+# ---------------------------------------------------------------------------
+# Tests
+# ---------------------------------------------------------------------------
+
+
+class TestUserPasswordMfaAuthentication:
+    # ------------------------------------------------------------------
+    # Passcode flow
+    # ------------------------------------------------------------------
+
+    def test_should_authenticate_with_explicit_passcode(self, connection_factory, mfa_params, mfa_passcode):
+        # Given Authentication is set to USERNAME_PASSWORD_MFA and a TOTP passcode is provided
+        connection = connection_factory(
+            authenticator="USERNAME_PASSWORD_MFA",
+            user=mfa_params["user"],
+            password=mfa_params["password"],
+            passcode=mfa_passcode,
+        )
+
+        # Then Login is successful and a simple query can be executed
+        with connection:
+            verify_simple_query_execution(connection)
+
+    def test_should_authenticate_with_passcode_in_password(self, connection_factory, mfa_params, mfa_passcode):
+        # Given Authentication is set to USERNAME_PASSWORD_MFA and the passcode is appended to the password
+        combined_password = mfa_params["password"] + mfa_passcode
+
+        connection = connection_factory(
+            authenticator="USERNAME_PASSWORD_MFA",
+            user=mfa_params["user"],
+            password=combined_password,
+            passcode_in_password=True,
+        )
+
+        # Then Login is successful and a simple query can be executed
+        with connection:
+            verify_simple_query_execution(connection)
+
+    # ------------------------------------------------------------------
+    # Token caching flow
+    # ------------------------------------------------------------------
+
+    @pytest.mark.skipif(not NEW_DRIVER_ONLY("BD#MFA1"), reason="Token caching only supported in new driver")
+    def test_should_cache_mfa_token_on_first_connection(self, connection_factory, mfa_params, mfa_passcode):
+        # Given caching flags are enabled and a valid passcode is provided
+        connection = connection_factory(
+            authenticator="USERNAME_PASSWORD_MFA",
+            user=mfa_params["user"],
+            password=mfa_params["password"],
+            passcode=mfa_passcode,
+            client_store_temporary_credential=True,
+        )
+
+        # Then Login is successful – the server issues an MFA token that the driver caches
+        with connection:
+            verify_simple_query_execution(connection)
+
+    @pytest.mark.skipif(not NEW_DRIVER_ONLY("BD#MFA2"), reason="Token caching only supported in new driver")
+    def test_should_reuse_cached_mfa_token_without_passcode(self, connection_factory, mfa_params, mfa_passcode):
+        # Given the first connection has already cached an MFA token (see test above)
+        # Establish first connection to warm the cache
+        first = connection_factory(
+            authenticator="USERNAME_PASSWORD_MFA",
+            user=mfa_params["user"],
+            password=mfa_params["password"],
+            passcode=mfa_passcode,
+            client_store_temporary_credential=True,
+        )
+        with first:
+            verify_simple_query_execution(first)
+
+        # When a second connection is made without a passcode
+        second = connection_factory(
+            authenticator="USERNAME_PASSWORD_MFA",
+            user=mfa_params["user"],
+            password=mfa_params["password"],
+            client_store_temporary_credential=True,
+        )
+
+        # Then the cached token is used and login succeeds without prompting for MFA
+        with second:
+            verify_simple_query_execution(second)
+
+    # ------------------------------------------------------------------
+    # DUO push flow
+    # ------------------------------------------------------------------
+
+    @pytest.mark.skip(reason="DUO push requires interactive device approval – run manually")
+    def test_should_authenticate_with_duo_push(self, connection_factory, mfa_params):
+        # Given DUO push is configured as the MFA method
+        connection = connection_factory(
+            authenticator="USERNAME_PASSWORD_MFA",
+            user=mfa_params["user"],
+            password=mfa_params["password"],
+            ext_authn_duo_method="push",
+        )
+
+        # Then (after approving the push on the registered device) login succeeds
+        with connection:
+            verify_simple_query_execution(connection)
+
+    # ------------------------------------------------------------------
+    # Error cases
+    # ------------------------------------------------------------------
+
+    def test_should_fail_with_invalid_passcode(self, connection_factory, mfa_params):
+        # Given an incorrect TOTP passcode is provided
+        with pytest.raises(Exception) as exc_info:
+            connection_factory(
+                authenticator="USERNAME_PASSWORD_MFA",
+                user=mfa_params["user"],
+                password=mfa_params["password"],
+                passcode="000000",
+            )
+
+        # Then a DatabaseError is raised
+        from snowflake.connector.errors import DatabaseError
+
+        assert isinstance(exc_info.value, DatabaseError), f"Expected DatabaseError, got: {type(exc_info.value)}"
+
+    def test_should_fail_with_invalid_password(self, connection_factory, mfa_params, mfa_passcode):
+        # Given the password is wrong
+        with pytest.raises(Exception) as exc_info:
+            connection_factory(
+                authenticator="USERNAME_PASSWORD_MFA",
+                user=mfa_params["user"],
+                password="wrong_password",
+                passcode=mfa_passcode,
+            )
+
+        # Then a DatabaseError is raised
+        from snowflake.connector.errors import DatabaseError
+
+        assert isinstance(exc_info.value, DatabaseError), f"Expected DatabaseError, got: {type(exc_info.value)}"

sf_core/src/auth.rs

@@ -34,6 +34,12 @@ pub enum Credentials {
         username: String,
         token: SensitiveString,
     },
+    UserPasswordMfa {
+        username: String,
+        password: SensitiveString,
+        passcode_in_password: bool,
+        passcode: Option<SensitiveString>,
+    },
 }
 
 #[derive(Debug, Serialize)]
@@ -148,6 +154,18 @@ pub fn create_credentials(login_parameters: &LoginParameters) -> Result<Credenti
             username: username.clone(),
             token: token.clone(),
         }),
+        LoginMethod::UserPasswordMfa {
+            username,
+            password,
+            passcode_in_password,
+            passcode,
+            ..
+        } => Ok(Credentials::UserPasswordMfa {
+            username: username.clone(),
+            password: password.clone(),
+            passcode: passcode.clone(),
+            passcode_in_password: *passcode_in_password,
+        }),
     }
 }
 

sf_core/src/config/rest_parameters.rs

@@ -157,6 +157,13 @@ pub enum LoginMethod {
         username: String,
         token: SensitiveString,
     },
+    UserPasswordMfa {
+        username: String,
+        password: SensitiveString,
+        passcode_in_password: bool,
+        passcode: Option<SensitiveString>,
+        client_store_temporary_credential: bool,
+    },
 }
 
 impl LoginMethod {
@@ -345,10 +352,42 @@ impl LoginMethod {
                     authentication_timeout_secs,
                 }))
             }
+            "USERNAME_PASSWORD_MFA" => Ok(Self::UserPasswordMfa {
+                username: settings
+                    .get_string("user")
+                    .context(MissingParameterSnafu { parameter: "user" })?,
+                password: settings
+                    .get_string("password")
+                    .context(MissingParameterSnafu { parameter: "password" })?
+                    .into(),
+                passcode_in_password: settings
+                    .get_bool("passcodeInPassword")
+                    .or_else(|| {
+                        settings
+                            .get_string("passcodeInPassword")
+                            .map(|v| v.eq_ignore_ascii_case("true") || v == "1")
+                    })
+                    .or_else(|| settings.get_int("passcodeInPassword").map(|v| v != 0))
+                    .unwrap_or(false),
+                passcode: settings.get_string("passcode").map(SensitiveString::from),
+                client_store_temporary_credential: settings
+                    .get_bool("client_store_temporary_credential")
+                    .or_else(|| {
+                        settings
+                            .get_string("client_store_temporary_credential")
+                            .map(|v| v.eq_ignore_ascii_case("true") || v == "1")
+                    })
+                    .or_else(|| {
+                        settings
+                            .get_int("client_store_temporary_credential")
+                            .map(|v| v != 0)
+                    })
+                    .unwrap_or(false),
+            }),
             _ => InvalidParameterValueSnafu {
                 parameter: "authenticator",
                 value: authenticator,
-                explanation: "Allowed values are SNOWFLAKE_JWT, SNOWFLAKE_PASSWORD, PROGRAMMATIC_ACCESS_TOKEN, or an https:// URL for native Okta SSO",
+                explanation: "Allowed values are SNOWFLAKE_JWT, SNOWFLAKE_PASSWORD, PROGRAMMATIC_ACCESS_TOKEN, USERNAME_PASSWORD_MFA, or an https:// URL for native Okta SSO",
             }
             .fail()?,
         }

sf_core/src/rest/snowflake/auth.rs

@@ -56,6 +56,11 @@ pub struct AuthRequestData {
         skip_serializing_if = "Option::is_none"
     )]
     pub ext_authn_duo_method: Option<String>,
+    #[serde(
+        rename = "CLIENT_REQUEST_MFA_TOKEN",
+        skip_serializing_if = "Option::is_none"
+    )]
+    pub client_request_mfa_token: Option<bool>,
     #[serde(rename = "PASSCODE", skip_serializing_if = "Option::is_none")]
     pub passcode: Option<SensitiveString>,
     #[serde(rename = "AUTHENTICATOR", skip_serializing_if = "Option::is_none")]
@@ -126,9 +131,9 @@ pub struct AuthResponseMain {
     )]
     pub master_validity: Option<Duration>,
     #[serde(rename = "mfaToken")]
-    pub _mfa_token: Option<String>,
+    pub mfa_token: Option<String>,
     #[serde(rename = "mfaTokenValidityInSeconds")]
-    pub _mfa_token_validity: Option<u64>,
+    pub mfa_token_validity: Option<u64>,
     #[serde(rename = "idToken")]
     pub _id_token: Option<String>,
     #[serde(rename = "idTokenValidityInSeconds")]