changes done in auth util as well as sea http

Signed-off-by: Nikhil Suri <[email protected]>

Author: nikhilsuri-db

src/databricks/sql/auth/auth_utils.py

@@ -7,23 +7,6 @@
 logger = logging.getLogger(__name__)
 
 
-def parse_hostname(hostname: str) -> str:
-    """
-    Normalize the hostname to include scheme and trailing slash.
-
-    Args:
-        hostname: The hostname to normalize
-
-    Returns:
-        Normalized hostname with scheme and trailing slash
-    """
-    if not hostname.startswith("http://") and not hostname.startswith("https://"):
-        hostname = f"https://{hostname}"
-    if not hostname.endswith("/"):
-        hostname = f"{hostname}/"
-    return hostname
-
-
 def decode_token(access_token: str) -> Optional[Dict]:
     """
     Decode a JWT token without verification to extract claims.

src/databricks/sql/auth/token_federation.py

@@ -6,10 +6,10 @@
 
 from databricks.sql.auth.authenticators import AuthProvider
 from databricks.sql.auth.auth_utils import (
-    parse_hostname,
     decode_token,
     is_same_host,
 )
+from databricks.sql.common.url_utils import normalize_host_with_protocol
 from databricks.sql.common.http import HttpMethod
 
 logger = logging.getLogger(__name__)
@@ -99,7 +99,7 @@ def __init__(
         if not http_client:
             raise ValueError("http_client is required for TokenFederationProvider")
 
-        self.hostname = parse_hostname(hostname)
+        self.hostname = normalize_host_with_protocol(hostname)
         self.external_provider = external_provider
         self.http_client = http_client
         self.identity_federation_client_id = identity_federation_client_id
@@ -164,7 +164,7 @@ def _should_exchange_token(self, access_token: str) -> bool:
 
     def _exchange_token(self, access_token: str) -> Token:
         """Exchange the external token for a Databricks token."""
-        token_url = f"{self.hostname.rstrip('/')}{self.TOKEN_EXCHANGE_ENDPOINT}"
+        token_url = f"{self.hostname}{self.TOKEN_EXCHANGE_ENDPOINT}"
 
         data = {
             "grant_type": self.TOKEN_EXCHANGE_GRANT_TYPE,

src/databricks/sql/backend/sea/utils/http_client.py

@@ -18,6 +18,7 @@
 from databricks.sql.common.http_utils import (
     detect_and_parse_proxy,
 )
+from databricks.sql.common.url_utils import normalize_host_with_protocol
 
 logger = logging.getLogger(__name__)
 
@@ -66,8 +67,9 @@ def __init__(
         self.auth_provider = auth_provider
         self.ssl_options = ssl_options
 
-        # Build base URL
-        self.base_url = f"https://{server_hostname}:{self.port}"
+        # Build base URL using url_utils for consistent normalization
+        normalized_host = normalize_host_with_protocol(server_hostname)
+        self.base_url = f"{normalized_host}:{self.port}"
 
         # Parse URL for proxy handling
         parsed_url = urllib.parse.urlparse(self.base_url)

src/databricks/sql/common/url_utils.py

@@ -5,7 +5,7 @@
 
 def normalize_host_with_protocol(host: str) -> str:
     """
-    Normalize a connection hostname by ensuring it has a protocol and removing trailing slashes.
+    Normalize a connection hostname by ensuring it has a protocol.
 
     This is useful for handling cases where users may provide hostnames with or without protocols
     (common with dbt-databricks users copying URLs from their browser).
@@ -15,12 +15,13 @@ def normalize_host_with_protocol(host: str) -> str:
               and may or may not have a trailing slash
 
     Returns:
-        Normalized hostname with protocol prefix and no trailing slash
+        Normalized hostname with protocol prefix and no trailing slashes
 
     Examples:
         normalize_host_with_protocol("myserver.com") -> "https://myserver.com"
         normalize_host_with_protocol("https://myserver.com") -> "https://myserver.com"
-        normalize_host_with_protocol("HTTPS://myserver.com") -> "https://myserver.com"
+        normalize_host_with_protocol("HTTPS://myserver.com/") -> "https://myserver.com"
+        normalize_host_with_protocol("http://localhost:8080/") -> "http://localhost:8080"
 
     Raises:
         ValueError: If host is None or empty string
@@ -29,7 +30,7 @@ def normalize_host_with_protocol(host: str) -> str:
     if not host or not host.strip():
         raise ValueError("Host cannot be None or empty")
 
-    # Remove trailing slash
+    # Remove trailing slashes
     host = host.rstrip("/")
 
     # Add protocol if not present (case-insensitive check)

tests/unit/test_sea_http_client.py

@@ -44,6 +44,40 @@ def sea_http_client(self, mock_auth_provider, ssl_options):
             client._pool = Mock()
             return client
 
+    @pytest.mark.parametrize(
+        "server_hostname,port,expected_base_url",
+        [
+            # Basic hostname without protocol
+            ("myserver.com", 443, "https://myserver.com:443"),
+            # Hostname with trailing slash
+            ("myserver.com/", 443, "https://myserver.com:443"),
+            # Hostname with https:// protocol
+            ("https://myserver.com", 443, "https://myserver.com:443"),
+            # Hostname with http:// protocol (preserved as-is)
+            ("http://myserver.com", 443, "http://myserver.com:443"),
+            # Hostname with protocol and trailing slash
+            ("https://myserver.com/", 443, "https://myserver.com:443"),
+            # Custom port
+            ("myserver.com", 8080, "https://myserver.com:8080"),
+            # Protocol with custom port
+            ("https://myserver.com", 8080, "https://myserver.com:8080"),
+        ],
+    )
+    def test_base_url_construction(
+        self, server_hostname, port, expected_base_url, mock_auth_provider, ssl_options
+    ):
+        """Test that base_url is constructed correctly from various hostname inputs."""
+        with patch("databricks.sql.backend.sea.utils.http_client.HTTPSConnectionPool"):
+            client = SeaHttpClient(
+                server_hostname=server_hostname,
+                port=port,
+                http_path="/sql/1.0/warehouses/test",
+                http_headers=[],
+                auth_provider=mock_auth_provider,
+                ssl_options=ssl_options,
+            )
+            assert client.base_url == expected_base_url
+
     def test_get_command_type_from_path(self, sea_http_client):
         """Test the _get_command_type_from_path method with various paths and methods."""
         # Test statement execution

tests/unit/test_token_federation.py

@@ -6,10 +6,10 @@
 
 from databricks.sql.auth.token_federation import TokenFederationProvider, Token
 from databricks.sql.auth.auth_utils import (
-    parse_hostname,
     decode_token,
     is_same_host,
 )
+from databricks.sql.common.url_utils import normalize_host_with_protocol
 from databricks.sql.common.http import HttpMethod
 
 
@@ -78,10 +78,10 @@ def test_init_requires_http_client(self, mock_external_provider):
     @pytest.mark.parametrize(
         "input_hostname,expected",
         [
-            ("test.databricks.com", "https://test.databricks.com/"),
-            ("https://test.databricks.com", "https://test.databricks.com/"),
-            ("https://test.databricks.com/", "https://test.databricks.com/"),
-            ("test.databricks.com/", "https://test.databricks.com/"),
+            ("test.databricks.com", "https://test.databricks.com"),
+            ("https://test.databricks.com", "https://test.databricks.com"),
+            ("https://test.databricks.com/", "https://test.databricks.com"),
+            ("test.databricks.com/", "https://test.databricks.com"),
         ],
     )
     def test_hostname_normalization(
@@ -305,15 +305,15 @@ class TestUtilityFunctions:
     @pytest.mark.parametrize(
         "input_hostname,expected",
         [
-            ("test.databricks.com", "https://test.databricks.com/"),
-            ("https://test.databricks.com", "https://test.databricks.com/"),
-            ("https://test.databricks.com/", "https://test.databricks.com/"),
-            ("test.databricks.com/", "https://test.databricks.com/"),
+            ("test.databricks.com", "https://test.databricks.com"),
+            ("https://test.databricks.com", "https://test.databricks.com"),
+            ("https://test.databricks.com/", "https://test.databricks.com"),
+            ("test.databricks.com/", "https://test.databricks.com"),
         ],
     )
-    def test_parse_hostname(self, input_hostname, expected):
-        """Test hostname parsing."""
-        assert parse_hostname(input_hostname) == expected
+    def test_normalize_hostname(self, input_hostname, expected):
+        """Test hostname normalization."""
+        assert normalize_host_with_protocol(input_hostname) == expected
 
     @pytest.mark.parametrize(
         "url1,url2,expected",

tests/unit/test_url_utils.py

@@ -6,65 +6,36 @@
 class TestNormalizeHostWithProtocol:
     """Tests for normalize_host_with_protocol function."""
 
-    @pytest.mark.parametrize("input_host,expected_output", [
-        # Hostname without protocol - should add https://
-        ("myserver.com", "https://myserver.com"),
-        ("workspace.databricks.com", "https://workspace.databricks.com"),
-        
-        # Hostname with https:// - should not duplicate
-        ("https://myserver.com", "https://myserver.com"),
-        ("https://workspace.databricks.com", "https://workspace.databricks.com"),
-        
-        # Hostname with http:// - should preserve
-        ("http://localhost", "http://localhost"),
-        ("http://myserver.com:8080", "http://myserver.com:8080"),
-        
-        # Hostname with port numbers
-        ("myserver.com:443", "https://myserver.com:443"),
-        ("https://myserver.com:443", "https://myserver.com:443"),
-        ("http://localhost:8080", "http://localhost:8080"),
-        
-        # Trailing slash - should be removed
-        ("myserver.com/", "https://myserver.com"),
-        ("https://myserver.com/", "https://myserver.com"),
-        ("http://localhost/", "http://localhost"),
-        
-        # Case-insensitive protocol handling - should normalize to lowercase
-        ("HTTPS://myserver.com", "https://myserver.com"),
-        ("HTTP://myserver.com", "http://myserver.com"),
-        ("HttPs://workspace.databricks.com", "https://workspace.databricks.com"),
-        ("HtTp://localhost:8080", "http://localhost:8080"),
-        ("HTTPS://MYSERVER.COM", "https://MYSERVER.COM"),  # Only protocol lowercased
-        
-        # Case-insensitive with trailing slashes
-        ("HTTPS://myserver.com/", "https://myserver.com"),
-        ("HTTP://localhost:8080/", "http://localhost:8080"),
-        ("HttPs://workspace.databricks.com//", "https://workspace.databricks.com"),
-        
-        # Mixed case protocols with ports
-        ("HTTPS://myserver.com:443", "https://myserver.com:443"),
-        ("HtTp://myserver.com:8080", "http://myserver.com:8080"),
-        
-        # Case preservation - only protocol lowercased, hostname case preserved
-        ("HTTPS://MyServer.DataBricks.COM", "https://MyServer.DataBricks.COM"),
-        ("HttPs://CamelCase.Server.com", "https://CamelCase.Server.com"),
-        ("HTTP://UPPERCASE.COM:8080", "http://UPPERCASE.COM:8080"),
-    ])
-    def test_normalize_host_with_protocol(self, input_host, expected_output):
+    @pytest.mark.parametrize(
+        "input_url,expected_output",
+        [
+            ("myserver.com", "https://myserver.com"),  # Add https://
+            ("https://myserver.com", "https://myserver.com"),  # No duplicate
+            ("http://localhost:8080", "http://localhost:8080"),  # Preserve http://
+            ("myserver.com:443", "https://myserver.com:443"),  # With port
+            ("myserver.com/", "https://myserver.com"),  # Remove trailing slash
+            ("https://myserver.com///", "https://myserver.com"),  # Multiple slashes
+            ("HTTPS://MyServer.COM", "https://MyServer.COM"),  # Case handling
+        ],
+    )
+    def test_normalize_host_with_protocol(self, input_url, expected_output):
         """Test host normalization with various input formats."""
-        result = normalize_host_with_protocol(input_host)
+        result = normalize_host_with_protocol(input_url)
         assert result == expected_output
-        
-        # Additional assertion: verify protocol is always lowercase
+
+        # Additional assertions
         assert result.startswith("https://") or result.startswith("http://")
+        assert not result.endswith("/")
 
-    @pytest.mark.parametrize("invalid_host", [
-        None,
-        "",
-        "   ",  # Whitespace only
-    ])
+    @pytest.mark.parametrize(
+        "invalid_host",
+        [
+            None,
+            "",
+            "   ",  # Whitespace only
+        ],
+    )
     def test_normalize_host_with_protocol_raises_on_invalid_input(self, invalid_host):
         """Test that function raises ValueError for None or empty host."""
         with pytest.raises(ValueError, match="Host cannot be None or empty"):
             normalize_host_with_protocol(invalid_host)
-