SNOW-733835 fix AuthByIdToken expired token auth (#1444)

Author: sfc-gh-mkeller

DESCRIPTION.md

@@ -12,6 +12,7 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
 
   - Improved the robustness of OCSP response caching to handle errors in cases of serialization and deserialization.
   - Fixed a bug where `AuthByKeyPair.handle_timeout` should pass keyword arguments instead of positional arguments when calling `AuthByKeyPair.prepare`. PR #1440 (@emilhe)
+  - Fixed a bug where MFA token caching would refuse to work until restarted instead of reauthenticating
 
 - v3.0.0(January 26, 2023)
 

src/snowflake/connector/auth/__init__.py

@@ -7,6 +7,7 @@
 from ._auth import Auth, get_public_key_fingerprint, get_token_from_private_key
 from .by_plugin import AuthByPlugin, AuthType
 from .default import AuthByDefault
+from .idtoken import AuthByIdToken
 from .keypair import AuthByKeyPair
 from .oauth import AuthByOAuth
 from .okta import AuthByOkta
@@ -21,6 +22,7 @@
         AuthByOkta,
         AuthByUsrPwdMfa,
         AuthByWebBrowser,
+        AuthByIdToken,
     )
 )
 

src/snowflake/connector/auth/idtoken.py

@@ -5,10 +5,14 @@
 
 from __future__ import annotations
 
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from ..network import ID_TOKEN_AUTHENTICATOR
 from .by_plugin import AuthByPlugin, AuthType
+from .webbrowser import AuthByWebBrowser
+
+if TYPE_CHECKING:
+    from ..connection import SnowflakeConnection
 
 
 class AuthByIdToken(AuthByPlugin):
@@ -25,19 +29,43 @@ def type_(self) -> AuthType:
     def assertion_content(self) -> str:
         return self._id_token
 
-    def __init__(self, id_token: str) -> None:
+    def __init__(
+        self,
+        id_token: str,
+        application: str,
+        protocol: str | None,
+        host: str | None,
+        port: str | None,
+    ) -> None:
         """Initialized an instance with an IdToken."""
         super().__init__()
         self._id_token: str | None = id_token
+        self._application = application
+        self._protocol = protocol
+        self._host = host
+        self._port = port
 
     def reset_secrets(self) -> None:
         self._id_token = None
 
     def prepare(self, **kwargs: Any) -> None:
         pass
 
-    def reauthenticate(self, **kwargs: Any) -> dict[str, bool]:
-        return {"success": False}
+    def reauthenticate(
+        self,
+        *,
+        conn: SnowflakeConnection,
+        **kwargs: Any,
+    ) -> dict[str, bool]:
+        conn.auth_class = AuthByWebBrowser(
+            application=self._application,
+            protocol=self._protocol,
+            host=self._host,
+            port=self._port,
+        )
+        conn._authenticate(conn.auth_class)
+        conn._auth_class.reset_secrets()
+        return {"success": True}
 
     def update_body(self, body: dict[Any, Any]) -> None:
         """Idtoken needs the authenticator and token attributes set."""

src/snowflake/connector/connection.py

@@ -351,6 +351,8 @@ def region(self) -> str | None:
         warnings.warn(
             "Region has been deprecated and will be removed in the near future",
             PendingDeprecationWarning,
+            # Raise warning from where this property was called from
+            stacklevel=2,
         )
         return self._region
 
@@ -804,7 +806,13 @@ def __open_connection(self):
                     port=self.port,
                 )
             else:
-                self.auth_class = AuthByIdToken(id_token=self._rest.id_token)
+                self.auth_class = AuthByIdToken(
+                    id_token=self._rest.id_token,
+                    application=self.application,
+                    protocol=self._protocol,
+                    host=self.host,
+                    port=self.port,
+                )
 
         elif self._authenticator == KEY_PAIR_AUTHENTICATOR:
             self.auth_class = AuthByKeyPair(private_key=self._private_key)
@@ -866,7 +874,9 @@ def __config(self, **kwargs):
                     warnings.warn(
                         "'{}' is an unknown connection parameter{}".format(
                             name, f", did you mean '{guess}'?" if guess else ""
-                        )
+                        ),
+                        # Raise warning from where class was initiated
+                        stacklevel=4,
                     )
                 elif not isinstance(value, DEFAULT_CONFIGURATION[name][1]):
                     accepted_types = DEFAULT_CONFIGURATION[name][1]
@@ -879,7 +889,9 @@ def __config(self, **kwargs):
                             if isinstance(accepted_types, tuple)
                             else accepted_types.__name__,
                             type(value).__name__,
-                        )
+                        ),
+                        # Raise warning from where class was initiated
+                        stacklevel=4,
                     )
             setattr(self, "_" + name, value)
 
@@ -1088,7 +1100,12 @@ def authenticate_with_retry(self, auth_instance) -> None:
         except ReauthenticationRequest as ex:
             # cached id_token expiration error, we have cleaned id_token and try to authenticate again
             logger.debug("ID token expired. Reauthenticating...: %s", ex)
-            self._authenticate(auth_instance)
+            if isinstance(auth_instance, AuthByIdToken):
+                # Note: SNOW-733835 IDToken auth needs to authenticate through
+                #  SSO if it has expired
+                self._reauthenticate()
+            else:
+                self._authenticate(auth_instance)
 
     def _authenticate(self, auth_instance: AuthByPlugin):
         auth_instance.prepare(

test/unit/test_auth_webbrowser.py

@@ -5,13 +5,19 @@
 
 from __future__ import annotations
 
+from unittest import mock
 from unittest.mock import MagicMock, Mock, PropertyMock, patch
 
 import pytest
 
+from snowflake.connector import SnowflakeConnection
 from snowflake.connector.constants import OCSPMode
 from snowflake.connector.description import CLIENT_NAME, CLIENT_VERSION
-from snowflake.connector.network import EXTERNAL_BROWSER_AUTHENTICATOR, SnowflakeRestful
+from snowflake.connector.network import (
+    EXTERNAL_BROWSER_AUTHENTICATOR,
+    ReauthenticationRequest,
+    SnowflakeRestful,
+)
 
 try:  # pragma: no cover
     from snowflake.connector.auth import AuthByWebBrowser
@@ -275,3 +281,40 @@ def post_request(url, headers, body, **kwargs):
     rest._post_request = post_request
     connection._rest = rest
     return rest
+
+
+def test_idtoken_reauth():
+    """This test makes sure that AuthByIdToken reverts to AuthByWebBrowser.
+
+    This happens when the initial connection fails. Such as when the saved ID
+    token has expired.
+    """
+    from snowflake.connector.auth.idtoken import AuthByIdToken
+
+    auth_inst = AuthByIdToken(
+        id_token="token",
+        application="application",
+        protocol="protocol",
+        host="host",
+        port="port",
+    )
+
+    # We'll use this Exception to make sure AuthByWebBrowser authentication
+    #  flow is called as expected
+    class StopExecuting(Exception):
+        pass
+
+    with mock.patch(
+        "snowflake.connector.auth.idtoken.AuthByIdToken.prepare",
+        side_effect=ReauthenticationRequest(Exception()),
+    ):
+        with mock.patch(
+            "snowflake.connector.auth.webbrowser.AuthByWebBrowser.prepare",
+            side_effect=StopExecuting(),
+        ):
+            with pytest.raises(StopExecuting):
+                SnowflakeConnection(
+                    user="user",
+                    account="account",
+                    auth_class=auth_inst,
+                )