Support client-side opt-in of Refresh Token Rotation in Snowflake OAuth (#2294)

Author: sfc-gh-pmansour

src/snowflake/connector/auth/oauth_code.py

@@ -58,6 +58,7 @@ def __init__(
         token_cache: TokenCache | None = None,
         refresh_token_enabled: bool = False,
         external_browser_timeout: int | None = None,
+        enable_single_use_refresh_tokens: bool = False,
         **kwargs,
     ) -> None:
         super().__init__(
@@ -81,6 +82,7 @@ def __init__(
             logger.debug("oauth pkce is going to be used")
         self._verifier: str | None = None
         self._external_browser_timeout = external_browser_timeout
+        self._enable_single_use_refresh_tokens = enable_single_use_refresh_tokens
 
     def _get_oauth_type_id(self) -> str:
         return OAUTH_TYPE_AUTHORIZATION_CODE
@@ -296,6 +298,8 @@ def _do_token_request(
             "code": code,
             "redirect_uri": callback_server.url,
         }
+        if self._enable_single_use_refresh_tokens:
+            fields["enable_single_use_refresh_tokens"] = "true"
         if self._pkce_enabled:
             assert self._verifier is not None
             fields["code_verifier"] = self._verifier

src/snowflake/connector/connection.py

@@ -355,6 +355,11 @@ def _get_private_bytes_from_file(
         True,
         bool,
     ),  # SNOW-XXXXX: remove the check_arrow_conversion_error_on_every_column flag
+    # Client-side opt-in to single-use refresh tokens.
+    "oauth_enable_single_use_refresh_tokens": (
+        False,
+        bool,
+    ),
 }
 
 APPLICATION_RE = re.compile(r"[\w\d_]+")
@@ -1236,6 +1241,7 @@ def __open_connection(self):
                     ),
                     refresh_token_enabled=features.refresh_token_enabled,
                     external_browser_timeout=self._external_browser_timeout,
+                    enable_single_use_refresh_tokens=self._oauth_enable_single_use_refresh_tokens,
                 )
             elif self._authenticator == OAUTH_CLIENT_CREDENTIALS:
                 self._check_experimental_authentication_flag()

test/unit/test_auth_oauth_auth_code.py

@@ -3,7 +3,12 @@
 # Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
 #
 
+from unittest.mock import patch
+
+import pytest
+
 from snowflake.connector.auth import AuthByOauthCode
+from snowflake.connector.network import OAUTH_AUTHORIZATION_CODE
 
 
 def test_auth_oauth_auth_code_oauth_type():
@@ -20,3 +25,42 @@ def test_auth_oauth_auth_code_oauth_type():
     body = {"data": {}}
     auth.update_body(body)
     assert body["data"]["OAUTH_TYPE"] == "authorization_code"
+
+
+@pytest.mark.parametrize("rtr_enabled", [True, False])
+def test_auth_oauth_auth_code_single_use_refresh_tokens(rtr_enabled: bool):
+    """Verifies that the enable_single_use_refresh_tokens option is plumbed into the authz code request."""
+    auth = AuthByOauthCode(
+        "app",
+        "clientId",
+        "clientSecret",
+        "auth_url",
+        "tokenRequestUrl",
+        "http://127.0.0.1:8080",
+        "scope",
+        pkce_enabled=False,
+        enable_single_use_refresh_tokens=rtr_enabled,
+    )
+
+    def fake_get_request_token_response(_, fields: dict[str, str]):
+        if rtr_enabled:
+            assert fields.get("enable_single_use_refresh_tokens") == "true"
+        else:
+            assert "enable_single_use_refresh_tokens" not in fields
+        return ("access_token", "refresh_token")
+
+    with patch(
+        "snowflake.connector.auth.AuthByOauthCode._do_authorization_request",
+        return_value="abc",
+    ):
+        with patch(
+            "snowflake.connector.auth.AuthByOauthCode._get_request_token_response",
+            side_effect=fake_get_request_token_response,
+        ):
+            auth.prepare(
+                conn=None,
+                authenticator=OAUTH_AUTHORIZATION_CODE,
+                service_name=None,
+                account="acc",
+                user="user",
+            )

test/unit/test_connection.py

@@ -687,3 +687,24 @@ def test_toml_connection_params_are_plumbed_into_authbyworkloadidentity(
             == "api://0b2f151f-09a2-46eb-ad5a-39d5ebef917b"
         )
         assert conn.auth_class.token == "my_token"
+
+
+@pytest.mark.parametrize("rtr_enabled", [True, False])
+def test_single_use_refresh_tokens_option_is_plumbed_into_authbyauthcode(
+    monkeypatch, rtr_enabled: bool
+):
+    with monkeypatch.context() as m:
+        m.setattr(
+            "snowflake.connector.SnowflakeConnection._authenticate", lambda *_: None
+        )
+        m.setenv("SF_ENABLE_EXPERIMENTAL_AUTHENTICATION", "true")
+
+        conn = snowflake.connector.connect(
+            account="my_account_1",
+            user="user",
+            oauth_client_id="client_id",
+            oauth_client_secret="client_secret",
+            authenticator="OAUTH_AUTHORIZATION_CODE",
+            oauth_enable_single_use_refresh_tokens=rtr_enabled,
+        )
+        assert conn.auth_class._enable_single_use_refresh_tokens == rtr_enabled