fixup! fixup! Make the changes backward compatible

Author: sfc-gh-turbaszek

src/snowflake/connector/auth/oauth_code.py

@@ -8,7 +8,6 @@
 import hashlib
 import json
 import logging
-import os
 import secrets
 import socket
 import time
@@ -17,11 +16,7 @@
 from typing import TYPE_CHECKING, Any
 
 from ..compat import parse_qs, urlparse, urlsplit
-from ..constants import (
-    ENV_VAR_OAUTH_SOCKET_ADDRESS,
-    ENV_VAR_OAUTH_SOCKET_PORT,
-    OAUTH_TYPE_AUTHORIZATION_CODE,
-)
+from ..constants import OAUTH_TYPE_AUTHORIZATION_CODE
 from ..errorcode import (
     ER_INVALID_VALUE,
     ER_OAUTH_CALLBACK_ERROR,
@@ -70,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(
@@ -97,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"
@@ -124,18 +121,11 @@ def _request_tokens(
         logger.debug("authenticating with OAuth authorization code flow")
         with AuthHttpServer(
             redirect_uri=self._redirect_uri,
-            uri=self._read_uri_from_env(),
+            uri=self._uri or self._redirect_uri,  # To preserve backward compatibility
         ) as callback_server:
             code = self._do_authorization_request(callback_server, conn)
             return self._do_token_request(code, callback_server, conn)
 
-    def _read_uri_from_env(self) -> str:
-        oauth_socket_address = os.getenv(
-            ENV_VAR_OAUTH_SOCKET_ADDRESS, "http://localhost"
-        )
-        oauth_socket_port = os.getenv(ENV_VAR_OAUTH_SOCKET_PORT, "0")
-        return f"{oauth_socket_address}:{oauth_socket_port}"
-
     def _check_post_requested(
         self, data: list[str]
     ) -> tuple[str, str] | tuple[None, None]:

src/snowflake/connector/connection.py

@@ -384,6 +384,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,
@@ -1456,6 +1461,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=(

src/snowflake/connector/constants.py

@@ -441,8 +441,6 @@ class IterUnit(Enum):
 # TODO: all env variables definitions should be here
 ENV_VAR_PARTNER = "SF_PARTNER"
 ENV_VAR_TEST_MODE = "SNOWFLAKE_TEST_MODE"
-ENV_VAR_OAUTH_SOCKET_ADDRESS = "SNOWFLAKE_OAUTH_SOCKET_ADDRESS"
-ENV_VAR_OAUTH_SOCKET_PORT = "SNOWFLAKE_OAUTH_SOCKET_PORT"
 
 _DOMAIN_NAME_MAP = {_DEFAULT_HOSTNAME_TLD: "GLOBAL", _CHINA_HOSTNAME_TLD: "CHINA"}
 

test/unit/test_auth_oauth_auth_code.py

@@ -394,3 +394,48 @@ def mock_request_tokens(self, **kwargs):
     assert isinstance(conn.auth_class, AuthByOauthCode)
 
     conn.close()
+
+
+@pytest.mark.parametrize(
+    "uri,redirect_uri",
+    [
+        ("https://example.com/server", "http://localhost:8080"),
+        ("http://localhost:8080", "https://example.com/redirect"),
+        ("http://127.0.0.1:9090", "https://server.com/oauth/callback"),
+        (None, "https://redirect.example.com"),
+    ],
+)
+@mock.patch(
+    "snowflake.connector.auth.oauth_code.AuthByOauthCode._do_authorization_request"
+)
+@mock.patch("snowflake.connector.auth.oauth_code.AuthByOauthCode._do_token_request")
+def test_auth_oauth_auth_code_passes_uri_to_http_server(
+    _, __, uri, redirect_uri, omit_oauth_urls_check
+):
+    """Test that uri and redirect_uri parameters are passed correctly to AuthHttpServer."""
+    auth = AuthByOauthCode(
+        "app",
+        "clientId",
+        "clientSecret",
+        "https://auth_url",
+        "tokenRequestUrl",
+        redirect_uri,
+        "scope",
+        "host",
+        uri=uri,
+    )
+
+    with patch(
+        "snowflake.connector.auth.oauth_code.AuthHttpServer",
+        # return_value=None,
+    ) as mock_http_server_init:
+        auth._request_tokens(
+            conn=mock.MagicMock(),
+            authenticator="authenticator",
+            service_name="service_name",
+            account="account",
+            user="user",
+        )
+        mock_http_server_init.assert_called_once_with(
+            uri=uri, redirect_uri=redirect_uri
+        )