SNOW-862388: fix okta retry bug (#1674)

Author: sfc-gh-aling

DESCRIPTION.md

@@ -10,6 +10,7 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
 
 - v3.1.1(TBD)
 
+  - Fixed a bug in retry logic for okta authentication to refresh token.
   - Support `RSAPublicKey` when constructing `AuthByKeyPair` in addition to raw bytes.
 
 - v3.1.0(July 31,2023)

src/snowflake/connector/auth/okta.py

@@ -277,13 +277,14 @@ def _step4(
         logger.debug("step 4: query IDP URL snowflake app to get SAML " "response")
         timeout_time = time.time() + conn.login_timeout if conn.login_timeout else None
         response_html = {}
+        origin_sso_url = sso_url
         while timeout_time is None or time.time() < timeout_time:
             try:
                 url_parameters = {
                     "RelayState": "/some/deep/link",
                     "onetimetoken": generate_one_time_token(),
                 }
-                sso_url = sso_url + "?" + urlencode(url_parameters)
+                sso_url = origin_sso_url + "?" + urlencode(url_parameters)
                 headers = {
                     HTTP_HEADER_ACCEPT: "*/*",
                 }

test/unit/test_auth_okta.py

@@ -209,13 +209,21 @@ def get_one_time_token():
         nonlocal second_token_generated
         if raise_token_refresh_error:
             assert not second_token_generated
+            return "1token1"
         else:
             second_token_generated = True
-        return "1token1"
+            return "2token2"
 
     # the first time, when step4 gets executed, we return 429
     # the second time when step4 gets retried, we return 200
     def mock_session_request(*args, **kwargs):
+        nonlocal second_token_generated
+        url = kwargs.get("url")
+        assert url == (
+            "https://testsso.snowflake.net/sso?RelayState=%2Fsome%2Fdeep%2Flink&onetimetoken=1token1"
+            if not second_token_generated
+            else "https://testsso.snowflake.net/sso?RelayState=%2Fsome%2Fdeep%2Flink&onetimetoken=2token2"
+        )
         nonlocal raise_token_refresh_error
         if raise_token_refresh_error:
             raise_token_refresh_error = False