SNOW-721779 fix MFA token cache logic (#1391)

Author: sfc-gh-mkeller

src/snowflake/connector/auth/_auth.py

@@ -450,7 +450,7 @@ def post_request_wrapper(self, url, headers, body):
             self._rest._connection._update_parameters(session_parameters)
             return session_parameters
 
-    def _read_temporary_credential(self, host, user, cred_type):
+    def _read_temporary_credential(self, host, user, cred_type) -> str | None:
         cred = None
         if IS_MACOS or IS_WINDOWS:
             if not installed_keyring:

src/snowflake/connector/connection.py

@@ -808,10 +808,19 @@ def __open_connection(self):
         elif self._authenticator == OAUTH_AUTHENTICATOR:
             self.auth_class = AuthByOAuth(oauth_token=self._token)
         elif self._authenticator == USR_PWD_MFA_AUTHENTICATOR:
-            self.auth_class = AuthByUsrPwdMfa(password=self._password)
             self._session_parameters[PARAMETER_CLIENT_REQUEST_MFA_TOKEN] = (
                 self._client_request_mfa_token if IS_LINUX else True
             )
+            if self._session_parameters[PARAMETER_CLIENT_REQUEST_MFA_TOKEN]:
+                auth.read_temporary_credentials(
+                    self.host,
+                    self.user,
+                    self._session_parameters,
+                )
+            self.auth_class = AuthByUsrPwdMfa(
+                password=self._password,
+                mfa_token=self.rest.mfa_token,
+            )
         else:
             # okta URL, e.g., https://<account>.okta.com/
             self.auth_class = AuthByOkta(application=self.application)

src/snowflake/connector/network.py

@@ -403,11 +403,11 @@ def id_token(self, value):
         self._id_token = value
 
     @property
-    def mfa_token(self):
+    def mfa_token(self) -> str | None:
         return getattr(self, "_mfa_token", None)
 
     @mfa_token.setter
-    def mfa_token(self, value):
+    def mfa_token(self, value: str) -> None:
         self._mfa_token = value
 
     def close(self):

test/unit/test_auth_mfa.py

@@ -0,0 +1,51 @@
+#
+# Copyright (c) 2012-2021 Snowflake Computing Inc. All rights reserved.
+#
+
+from unittest import mock
+
+from snowflake.connector import connect
+
+
+def test_mfa_token_cache():
+    with mock.patch(
+        "snowflake.connector.network.SnowflakeRestful.fetch",
+    ):
+        with mock.patch(
+            "snowflake.connector.auth._auth.Auth._write_temporary_credential",
+        ) as save_mock:
+            with connect(
+                account="account",
+                user="user",
+                password="password",
+                authenticator="username_password_mfa",
+                client_store_temporary_credential=True,
+                client_request_mfa_token=True,
+            ):
+                assert save_mock.called
+    with mock.patch(
+        "snowflake.connector.network.SnowflakeRestful.fetch",
+        return_value={
+            "data": {
+                "token": "abcd",
+                "masterToken": "defg",
+            },
+            "success": True,
+        },
+    ):
+        with mock.patch(
+            "snowflake.connector.cursor.SnowflakeCursor._init_result_and_meta",
+        ):
+            with mock.patch(
+                "snowflake.connector.auth._auth.Auth._read_temporary_credential",
+                return_value=None,
+            ) as load_mock:
+                with connect(
+                    account="account",
+                    user="user",
+                    password="password",
+                    authenticator="username_password_mfa",
+                    client_store_temporary_credential=True,
+                    client_request_mfa_token=True,
+                ):
+                    assert load_mock.called