Make the changes backward compatible

Author: sfc-gh-turbaszek

src/snowflake/connector/auth/_http_server.py

@@ -69,11 +69,11 @@ class AuthHttpServer:
     def __init__(
         self,
         uri: str,
-        redirect_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)
+        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":
@@ -84,10 +84,11 @@ def __init__(
             else:
                 self._socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1)
 
-        if parsed_redirect.hostname in ("localhost", "127.0.0.1"):
+        if parsed_redirect and self._is_local_uri(parsed_redirect):
             port = parsed_redirect.port or 0
         else:
-            port = parsed_uri.port or 0
+            port = parsed_uri.port if parsed_uri and parsed_uri.port else 0
+
         for attempt in range(1, self.DEFAULT_MAX_ATTEMPTS + 1):
             try:
                 self._socket.bind(
@@ -128,27 +129,30 @@ def __init__(
             query=parsed_uri.query,
             fragment=parsed_uri.fragment,
         )
-        if (
-            parsed_redirect.hostname in ("localhost", "127.0.0.1")
-            and port != parsed_redirect.port
-        ):
-            logger.debug(
-                f"Updating redirect port {parsed_redirect.port} to match the server port {port}."
-            )
-            self._redirect_uri = urllib.parse.ParseResult(
-                scheme=parsed_redirect.scheme,
-                netloc=parsed_redirect.hostname + ":" + str(port),
-                path=parsed_redirect.path,
-                params=parsed_redirect.params,
-                query=parsed_redirect.query,
-                fragment=parsed_redirect.fragment,
-            )
-        else:
-            self._redirect_uri = parsed_redirect
+        if parsed_redirect:
+            if self._is_local_uri(parsed_redirect) and port != parsed_redirect.port:
+                logger.debug(
+                    f"Updating redirect port {parsed_redirect.port} to match the server port {port}."
+                )
+                self._redirect_uri = urllib.parse.ParseResult(
+                    scheme=parsed_redirect.scheme,
+                    netloc=parsed_redirect.hostname + ":" + str(port),
+                    path=parsed_redirect.path,
+                    params=parsed_redirect.params,
+                    query=parsed_redirect.query,
+                    fragment=parsed_redirect.fragment,
+                )
+            else:
+                self._redirect_uri = parsed_redirect
+
+    def _is_local_uri(self, parsed_redirect):
+        return parsed_redirect.hostname in ("localhost", "127.0.0.1")
 
     @property
-    def redirect_uri(self) -> str:
-        return self._redirect_uri.geturl()
+    def redirect_uri(self) -> str | None:
+        if self._redirect_uri:
+            return self._redirect_uri.geturl()
+        return self.url
 
     @property
     def url(self) -> str:

src/snowflake/connector/auth/oauth_code.py

@@ -17,7 +17,11 @@
 from typing import TYPE_CHECKING, Any
 
 from ..compat import parse_qs, urlparse, urlsplit
-from ..constants import OAUTH_TYPE_AUTHORIZATION_CODE
+from ..constants import (
+    ENV_VAR_OAUTH_SOCKET_ADDRESS,
+    ENV_VAR_OAUTH_SOCKET_PORT,
+    OAUTH_TYPE_AUTHORIZATION_CODE,
+)
 from ..errorcode import (
     ER_INVALID_VALUE,
     ER_OAUTH_CALLBACK_ERROR,
@@ -119,14 +123,19 @@ def _request_tokens(
         """Web Browser based Authentication."""
         logger.debug("authenticating with OAuth authorization code flow")
         with AuthHttpServer(
-            uri=os.environ.get("SNOWFLAKE_OAUTH_SOCKET_ADDRESS", "http://localhost")
-            + ":"
-            + os.environ.get("SNOWFLAKE_OAUTH_SOCKET_PORT", "0"),
             redirect_uri=self._redirect_uri,
+            uri=self._read_uri_from_env(),
         ) 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/constants.py

@@ -441,7 +441,8 @@ 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_callback_server.py

@@ -25,8 +25,7 @@ def test_auth_callback_success(monkeypatch, dontwait, timeout, reuse_port) -> No
     monkeypatch.setenv("SNOWFLAKE_AUTH_SOCKET_MSG_DONTWAIT", dontwait)
     test_response: requests.Response | None = None
     with AuthHttpServer(
-        uri="http://127.0.0.1/test_request",
-        redirect_uri="http://127.0.0.1/test_request",
+        "http://127.0.0.1/test_request",
     ) as callback_server:
 
         def request_callback():
@@ -61,8 +60,7 @@ def test_auth_callback_timeout(monkeypatch, dontwait, timeout, reuse_port) -> No
     monkeypatch.setenv("SNOWFLAKE_AUTH_SOCKET_REUSE_PORT", reuse_port)
     monkeypatch.setenv("SNOWFLAKE_AUTH_SOCKET_MSG_DONTWAIT", dontwait)
     with AuthHttpServer(
-        uri="http://127.0.0.1/test_request",
-        redirect_uri="http://127.0.0.1/test_request",
+        "http://127.0.0.1/test_request",
     ) as callback_server:
         block, client_socket = callback_server.receive_block(timeout=timeout)
         assert block is None