Merge branch 'main' into zyao-SNOW-2483517-sproc-continuous-integration

Author: sfc-gh-turbaszek

DESCRIPTION.md

@@ -14,6 +14,7 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
   - Pin lower versions of dependencies to oldest version without vulnerabilities.
   - Added no_proxy parameter for proxy configuration without using environmental variables.
   - Added OAUTH_AUTHORIZATION_CODE and OAUTH_CLIENT_CREDENTIALS to list of authenticators that don't require user to be set
+  - Added `oauth_socket_uri` connection parameter allowing to separate server and redirect URIs for local OAuth server.
 
 - v4.0.0(October 09,2025)
   - Added support for checking certificates revocation using revocation lists (CRLs)

ci/wif/parameters/parameters_wif.json.gpg

@@ -514,6 +514,16 @@ def read_temporary_credentials(
         user: str,
         session_parameters: dict[str, Any],
     ) -> None:
+        """Attempt to load cached credentials to skip interactive authentication.
+
+        SSO (ID_TOKEN): If present, avoids opening browser for external authentication.
+            Controlled by client_store_temporary_credential parameter.
+
+        MFA (MFA_TOKEN): If present, skips MFA prompt on next connection.
+            Controlled by client_request_mfa_token parameter.
+
+        If cached tokens are expired/invalid, they're deleted and normal auth proceeds.
+        """
         if session_parameters.get(PARAMETER_CLIENT_STORE_TEMPORARY_CREDENTIAL, False):
             self._rest.id_token = self._read_temporary_credential(
                 host,
@@ -549,6 +559,13 @@ def write_temporary_credentials(
         session_parameters: dict[str, Any],
         response: dict[str, Any],
     ) -> None:
+        """Cache credentials received from successful authentication for future use.
+
+        Tokens are only cached if:
+        1. Server returned the token in response (server-side caching must be enabled)
+        2. Client has caching enabled via session parameters
+        3. User consented to caching (consent_cache_id_token for ID tokens)
+        """
         if (
             self._rest._connection.auth_class.consent_cache_id_token
             and session_parameters.get(

ci/wif/parameters/rsa_wif_aws_azure.gpg

@@ -70,8 +70,10 @@ def __init__(
         self,
         uri: str,
         buf_size: int = 16384,
+        redirect_uri: str | None = None,
     ) -> None:
         parsed_uri = urllib.parse.urlparse(uri)
+        parsed_redirect = urllib.parse.urlparse(redirect_uri) if redirect_uri else None
         self._socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
         self.buf_size = buf_size
         if os.getenv("SNOWFLAKE_AUTH_SOCKET_REUSE_PORT", "False").lower() == "true":
@@ -82,30 +84,34 @@ def __init__(
             else:
                 self._socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1)
 
-        port = parsed_uri.port or 0
+        if parsed_redirect and self._is_local_uri(parsed_redirect):
+            server_port = parsed_redirect.port or 0
+        else:
+            server_port = parsed_uri.port or 0
+
         for attempt in range(1, self.DEFAULT_MAX_ATTEMPTS + 1):
             try:
                 self._socket.bind(
                     (
                         parsed_uri.hostname,
-                        port,
+                        server_port,
                     )
                 )
                 break
             except socket.gaierror as ex:
                 logger.error(
-                    f"Failed to bind authorization callback server to port {port}: {ex}"
+                    f"Failed to bind authorization callback server to port {server_port}: {ex}"
                 )
                 raise
             except OSError as ex:
                 if attempt == self.DEFAULT_MAX_ATTEMPTS:
                     logger.error(
-                        f"Failed to bind authorization callback server to port {port}: {ex}"
+                        f"Failed to bind authorization callback server to port {server_port}: {ex}"
                     )
                     raise
                 logger.warning(
                     f"Attempt {attempt}/{self.DEFAULT_MAX_ATTEMPTS}. "
-                    f"Failed to bind authorization callback server to port {port}: {ex}"
+                    f"Failed to bind authorization callback server to port {server_port}: {ex}"
                 )
                 time.sleep(self.PORT_BIND_TIMEOUT / self.PORT_BIND_MAX_ATTEMPTS)
         try:
@@ -114,16 +120,47 @@ def __init__(
             logger.error(f"Failed to start listening for auth callback: {ex}")
             self.close()
             raise
-        port = self._socket.getsockname()[1]
+
+        server_port = self._socket.getsockname()[1]
         self._uri = urllib.parse.ParseResult(
             scheme=parsed_uri.scheme,
-            netloc=parsed_uri.hostname + ":" + str(port),
+            netloc=parsed_uri.hostname + ":" + str(server_port),
             path=parsed_uri.path,
             params=parsed_uri.params,
             query=parsed_uri.query,
             fragment=parsed_uri.fragment,
         )
 
+        if parsed_redirect:
+            if (
+                self._is_local_uri(parsed_redirect)
+                and server_port != parsed_redirect.port
+            ):
+                logger.debug(
+                    f"Updating redirect port {parsed_redirect.port} to match the server port {server_port}."
+                )
+                self._redirect_uri = urllib.parse.ParseResult(
+                    scheme=parsed_redirect.scheme,
+                    netloc=parsed_redirect.hostname + ":" + str(server_port),
+                    path=parsed_redirect.path,
+                    params=parsed_redirect.params,
+                    query=parsed_redirect.query,
+                    fragment=parsed_redirect.fragment,
+                )
+            else:
+                self._redirect_uri = parsed_redirect
+        else:
+            # For backwards compatibility
+            self._redirect_uri = self._uri
+
+    @staticmethod
+    def _is_local_uri(uri):
+        return uri.hostname in ("localhost", "127.0.0.1")
+
+    @property
+    def redirect_uri(self) -> str | None:
+        return self._redirect_uri.geturl()
+
     @property
     def url(self) -> str:
         return self._uri.geturl()

ci/wif/parameters/rsa_wif_gcp.gpg

@@ -35,6 +35,14 @@
 
 
 class _OAuthTokensMixin:
+    """Manages OAuth token caching to avoid repeated browser authentication flows.
+
+    Access tokens: Short-lived (typically 10 minutes), cached to avoid immediate re-auth.
+    Refresh tokens: Long-lived (hours/days), used to obtain new access tokens silently.
+
+    Tokens are cached per (user, IDP host) to support multiple OAuth providers/accounts.
+    """
+
     def __init__(
         self,
         token_cache: TokenCache | None,
@@ -77,12 +85,18 @@ def _pop_cached_token(self, key: TokenKey | None) -> str | None:
         return self._token_cache.retrieve(key)
 
     def _pop_cached_access_token(self) -> bool:
-        """Retrieves OAuth access token from the token cache if enabled"""
+        """Retrieves OAuth access token from the token cache if enabled, available and still valid.
+
+        Returns True if cached token found, allowing authentication to skip OAuth flow.
+        """
         self._access_token = self._pop_cached_token(self._get_access_token_cache_key())
         return self._access_token is not None
 
     def _pop_cached_refresh_token(self) -> bool:
-        """Retrieves OAuth refresh token from the token cache if enabled"""
+        """Retrieves OAuth refresh token from the token cache (if enabled) to silently obtain new access token.
+
+        Returns True if refresh token found, enabling automatic token renewal without user interaction.
+        """
         if self._refresh_token_enabled:
             self._refresh_token = self._pop_cached_token(
                 self._get_refresh_token_cache_key()

src/snowflake/connector/auth/_auth.py

@@ -65,6 +65,7 @@ def __init__(
         external_browser_timeout: int | None = None,
         enable_single_use_refresh_tokens: bool = False,
         connection: SnowflakeConnection | None = None,
+        uri: str | None = None,
         **kwargs,
     ) -> None:
         authentication_url, redirect_uri = self._validate_oauth_code_uris(
@@ -92,6 +93,7 @@ def __init__(
         self._origin: str | None = None
         self._authentication_url = authentication_url
         self._redirect_uri = redirect_uri
+        self._uri = uri
         self._state = secrets.token_urlsafe(43)
         logger.debug("chose oauth state: %s", "".join("*" for _ in self._state))
         self._protocol = "http"
@@ -117,7 +119,10 @@ def _request_tokens(
     ) -> (str | None, str | None):
         """Web Browser based Authentication."""
         logger.debug("authenticating with OAuth authorization code flow")
-        with AuthHttpServer(self._redirect_uri) as callback_server:
+        with AuthHttpServer(
+            redirect_uri=self._redirect_uri,
+            uri=self._uri or self._redirect_uri,  # for backward compatibility
+        ) as callback_server:
             code = self._do_authorization_request(callback_server, conn)
             return self._do_token_request(code, callback_server, conn)
 
@@ -260,7 +265,7 @@ def _do_authorization_request(
         connection: SnowflakeConnection,
     ) -> str | None:
         authorization_request = self._construct_authorization_request(
-            callback_server.url
+            callback_server.redirect_uri
         )
         logger.debug("step 1: going to open authorization URL")
         print(
@@ -315,7 +320,7 @@ def _do_token_request(
         fields = {
             "grant_type": "authorization_code",
             "code": code,
-            "redirect_uri": callback_server.url,
+            "redirect_uri": callback_server.redirect_uri,
         }
         if self._enable_single_use_refresh_tokens:
             fields["enable_single_use_refresh_tokens"] = "true"

src/snowflake/connector/auth/_http_server.py

@@ -280,8 +280,13 @@ def _get_private_bytes_from_file(
     "support_negative_year": (True, bool),  # snowflake
     "log_max_query_length": (LOG_MAX_QUERY_LENGTH, int),  # snowflake
     "disable_request_pooling": (False, bool),  # snowflake
-    # enable temporary credential file for Linux, default false. Mac/Win will overlook this
+    # Cache SSO ID tokens to avoid repeated browser popups. Must be enabled on the server-side.
+    # Storage: keyring (macOS/Windows), file (Linux). Auto-enabled on macOS/Windows.
+    # Sets session PARAMETER_CLIENT_STORE_TEMPORARY_CREDENTIAL as well
     "client_store_temporary_credential": (False, bool),
+    # Cache MFA tokens to skip MFA prompts on reconnect. Must be enabled on the server-side.
+    # Storage: keyring (macOS/Windows), file (Linux). Auto-enabled on macOS/Windows.
+    # In driver, we extract this from session using PARAMETER_CLIENT_REQUEST_MFA_TOKEN.
     "client_request_mfa_token": (False, bool),
     "use_openssl_only": (
         True,
@@ -384,6 +389,11 @@ def _get_private_bytes_from_file(
         # SNOW-1825621: OAUTH implementation
     ),
     "oauth_redirect_uri": ("http://127.0.0.1", str),
+    "oauth_socket_uri": (
+        "http://127.0.0.1",
+        str,
+        # SNOW-2194055: Separate server and redirect URIs in AuthHttpServer
+    ),
     "oauth_scope": (
         "",
         str,
@@ -1082,22 +1092,23 @@ def connect(self, **kwargs) -> None:
 
         self._crl_config: CRLConfig = CRLConfig.from_connection(self)
 
+        no_proxy_csv_str = (
+            ",".join(str(x) for x in self.no_proxy)
+            if (
+                self.no_proxy is not None
+                and isinstance(self.no_proxy, Iterable)
+                and not isinstance(self.no_proxy, (str, bytes))
+            )
+            else self.no_proxy
+        )
         self._http_config = HttpConfig(
             adapter_factory=ProxySupportAdapterFactory(),
             use_pooling=(not self.disable_request_pooling),
             proxy_host=self.proxy_host,
             proxy_port=self.proxy_port,
             proxy_user=self.proxy_user,
             proxy_password=self.proxy_password,
-            no_proxy=(
-                ",".join(str(x) for x in self.no_proxy)
-                if (
-                    self.no_proxy is not None
-                    and isinstance(self.no_proxy, Iterable)
-                    and not isinstance(self.no_proxy, (str, bytes))
-                )
-                else self.no_proxy
-            ),
+            no_proxy=no_proxy_csv_str,
         )
         self._session_manager = SessionManagerFactory.get_manager(self._http_config)
 
@@ -1157,7 +1168,8 @@ def close(self, retry: bool = True) -> None:
 
             # close telemetry first, since it needs rest to send remaining data
             logger.debug("closed")
-            self._telemetry.close(send_on_close=bool(retry and self.telemetry_enabled))
+            if self.telemetry_enabled:
+                self._telemetry.close(retry=retry)
             if (
                 self._all_async_queries_finished()
                 and not self._server_session_keep_alive
@@ -1390,9 +1402,11 @@ def __open_connection(self):
                     backoff_generator=self._backoff_generator,
                 )
             elif self._authenticator == EXTERNAL_BROWSER_AUTHENTICATOR:
+                # Enable SSO credential caching
                 self._session_parameters[
                     PARAMETER_CLIENT_STORE_TEMPORARY_CREDENTIAL
                 ] = (self._client_store_temporary_credential if IS_LINUX else True)
+                # Try to load cached ID token to avoid browser popup
                 auth.read_temporary_credentials(
                     self.host,
                     self.user,
@@ -1456,6 +1470,7 @@ def __open_connection(self):
                         host=self.host, port=self.port
                     ),
                     redirect_uri=self._oauth_redirect_uri,
+                    uri=self._oauth_socket_uri,
                     scope=self._oauth_scope,
                     pkce_enabled=not self._oauth_disable_pkce,
                     token_cache=(
@@ -1483,9 +1498,11 @@ def __open_connection(self):
                     connection=self,
                 )
             elif self._authenticator == USR_PWD_MFA_AUTHENTICATOR:
+                # Enable MFA token caching
                 self._session_parameters[PARAMETER_CLIENT_REQUEST_MFA_TOKEN] = (
                     self._client_request_mfa_token if IS_LINUX else True
                 )
+                # Try to load cached MFA token to skip MFA prompt
                 if self._session_parameters[PARAMETER_CLIENT_REQUEST_MFA_TOKEN]:
                     auth.read_temporary_credentials(
                         self.host,

src/snowflake/connector/auth/_oauth_base.py

@@ -442,7 +442,6 @@ class IterUnit(Enum):
 ENV_VAR_PARTNER = "SF_PARTNER"
 ENV_VAR_TEST_MODE = "SNOWFLAKE_TEST_MODE"
 
-
 _DOMAIN_NAME_MAP = {_DEFAULT_HOSTNAME_TLD: "GLOBAL", _CHINA_HOSTNAME_TLD: "CHINA"}
 
 _CONNECTIVITY_ERR_MSG = (