more fixes

Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>

Author: vikrantpuppala

src/databricks/sql/auth/authenticators.py

@@ -199,7 +199,8 @@ def __init__(
         self.azure_client_secret = azure_client_secret
         self.azure_workspace_resource_id = azure_workspace_resource_id
         self.azure_tenant_id = azure_tenant_id or get_azure_tenant_id_from_host(
-            hostname
+            hostname,
+            http_client
         )
         self._http_client = http_client
 

src/databricks/sql/auth/common.py

@@ -116,15 +116,15 @@ def get_azure_tenant_id_from_host(host: str, http_client) -> str:
     logger.debug("Loading tenant ID from %s", login_url)
 
     with http_client.request_context(
-        HttpMethod.GET, login_url, allow_redirects=False
+        HttpMethod.GET, login_url
     ) as resp:
-        if resp.status // 100 != 3:
-            raise ValueError(
-                f"Failed to get tenant ID from {login_url}: expected status code 3xx, got {resp.status}"
-            )
-        entra_id_endpoint = dict(resp.headers).get("Location")
+        # if resp.status // 100 != 3:
+        #     raise ValueError(
+        #         f"Failed to get tenant ID from {login_url}: expected status code 3xx, got {resp.status}"
+        #     )
+        entra_id_endpoint = resp.retries.history[-1].redirect_location
         if entra_id_endpoint is None:
-            raise ValueError(f"No Location header in response from {login_url}")
+            raise ValueError(f"No Location header in response from {login_url}: {entra_id_endpoint}")
 
     # The Location header has the following form: https://login.microsoftonline.com/<tenant-id>/oauth2/authorize?...
     # The domain may change depending on the Azure cloud (e.g. login.microsoftonline.us for US Government cloud).

src/databricks/sql/auth/oauth.py

@@ -336,18 +336,23 @@ def refresh(self) -> Token:
                 **self.extra_params,
             }
         )
+        
 
-        with self._http_client.execute(
-            method=HttpMethod.POST, url=self.token_url, headers=headers, data=data
-        ) as response:
-            if response.status_code == 200:
-                oauth_response = OAuthResponse(**response.json())
+        response = self._http_client.request(
+            method=HttpMethod.POST, url=self.token_url, headers=headers, body=data
+        )
+        try:
+            if response.status == 200:
+                import json
+                oauth_response = OAuthResponse(**json.loads(response.data.decode('utf-8')))
                 return Token(
                     oauth_response.access_token,
                     oauth_response.token_type,
                     oauth_response.refresh_token,
                 )
             else:
                 raise Exception(
-                    f"Failed to get token: {response.status_code} {response.text}"
+                    f"Failed to get token: {response.status} {response.data.decode('utf-8')}"
                 )
+        finally:
+            response.close()

src/databricks/sql/auth/retry.py

@@ -355,7 +355,7 @@ def should_retry(self, method: str, status_code: int) -> Tuple[bool, str]:
         logger.info(f"Received status code {status_code} for {method} request")
 
         # Request succeeded. Don't retry.
-        if status_code // 100 == 2:
+        if status_code // 100 <= 3:
             return False, "2xx codes are not retried"
 
         if status_code == 401:

src/databricks/sql/client.py

@@ -279,7 +279,7 @@ def read(self) -> Optional[OAuthToken]:
                 host_url=server_hostname,
                 http_path=http_path,
                 port=kwargs.get("_port", 443),
-                http_client=self.http_client,
+                client_context=client_context,
                 user_agent=self.session.useragent_header
                 if hasattr(self, "session")
                 else None,
@@ -301,7 +301,7 @@ def read(self) -> Optional[OAuthToken]:
             auth_provider=self.session.auth_provider,
             host_url=self.session.host,
             batch_size=self.telemetry_batch_size,
-            http_client=self.http_client,
+            client_context=client_context,
         )
 
         self._telemetry_client = TelemetryClientFactory.get_telemetry_client(

src/databricks/sql/common/unified_http_client.py

@@ -154,7 +154,7 @@ def request_context(
         Yields:
             urllib3.HTTPResponse: The HTTP response object
         """
-        logger.debug("Making %s request to %s", method, url)
+        logger.debug("Making %s request to %s", method, urllib.parse.urlparse(url).netloc)
 
         request_headers = self._prepare_headers(headers)
 

src/databricks/sql/telemetry/telemetry_client.py

@@ -172,7 +172,7 @@ def __init__(
         host_url,
         executor,
         batch_size,
-        http_client,
+        client_context,
     ):
         logger.debug("Initializing TelemetryClient for connection: %s", session_id_hex)
         self._telemetry_enabled = telemetry_enabled
@@ -186,8 +186,10 @@ def __init__(
         self._host_url = host_url
         self._executor = executor
 
-        # Use the provided HTTP client directly
-        self._http_client = http_client
+        # Create own HTTP client from client context
+        from databricks.sql.common.unified_http_client import UnifiedHttpClient
+
+        self._http_client = UnifiedHttpClient(client_context)
 
     def _export_event(self, event):
         """Add an event to the batch queue and flush if batch is full"""
@@ -456,7 +458,7 @@ def initialize_telemetry_client(
         auth_provider,
         host_url,
         batch_size,
-        http_client,
+        client_context,
     ):
         """Initialize a telemetry client for a specific connection if telemetry is enabled"""
         try:
@@ -479,7 +481,7 @@ def initialize_telemetry_client(
                             host_url=host_url,
                             executor=TelemetryClientFactory._executor,
                             batch_size=batch_size,
-                            http_client=http_client,
+                            client_context=client_context,
                         )
                     else:
                         TelemetryClientFactory._clients[
@@ -532,10 +534,10 @@ def connection_failure_log(
         host_url: str,
         http_path: str,
         port: int,
-        http_client,
+        client_context,
         user_agent: Optional[str] = None,
     ):
-        """Send error telemetry when connection creation fails, using existing HTTP client"""
+        """Send error telemetry when connection creation fails, using provided client context"""
 
         UNAUTH_DUMMY_SESSION_ID = "unauth_session_id"
 
@@ -545,7 +547,7 @@ def connection_failure_log(
             auth_provider=None,
             host_url=host_url,
             batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE,
-            http_client=http_client,
+            client_context=client_context,
         )
 
         telemetry_client = TelemetryClientFactory.get_telemetry_client(

tests/e2e/test_concurrent_telemetry.py

@@ -4,6 +4,7 @@
 import time
 from unittest.mock import patch
 import pytest
+import json
 
 from databricks.sql.telemetry.models.enums import StatementType
 from databricks.sql.telemetry.telemetry_client import (
@@ -119,8 +120,12 @@ def execute_query_worker(thread_id):
             for future in done:
                 try:
                     response = future.result()
-                    response.raise_for_status()
-                    captured_responses.append(response.json())
+                    # Check status using urllib3 method (response.status instead of response.raise_for_status())
+                    if response.status >= 400:
+                        raise Exception(f"HTTP {response.status}: {getattr(response, 'reason', 'Unknown')}")
+                    # Parse JSON using urllib3 method (response.data.decode() instead of response.json())
+                    response_data = json.loads(response.data.decode()) if response.data else {}
+                    captured_responses.append(response_data)
                 except Exception as e:
                     captured_exceptions.append(e)
 

tests/unit/test_telemetry.py

@@ -24,17 +24,19 @@ def mock_telemetry_client():
     session_id = str(uuid.uuid4())
     auth_provider = AccessTokenAuthProvider("test-token")
     executor = MagicMock()
-    mock_http_client = MagicMock()
-
-    return TelemetryClient(
-        telemetry_enabled=True,
-        session_id_hex=session_id,
-        auth_provider=auth_provider,
-        host_url="test-host.com",
-        executor=executor,
-        batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE,
-        http_client=mock_http_client,
-    )
+    client_context = MagicMock()
+
+    # Patch the _setup_pool_manager method to avoid SSL file loading
+    with patch('databricks.sql.common.unified_http_client.UnifiedHttpClient._setup_pool_manager'):
+        return TelemetryClient(
+            telemetry_enabled=True,
+            session_id_hex=session_id,
+            auth_provider=auth_provider,
+            host_url="test-host.com",
+            executor=executor,
+            batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE,
+            client_context=client_context,
+        )
 
 
 class TestNoopTelemetryClient:
@@ -216,41 +218,42 @@ def test_client_lifecycle_flow(self):
         """Test complete client lifecycle: initialize -> use -> close."""
         session_id_hex = "test-session"
         auth_provider = AccessTokenAuthProvider("token")
-        mock_http_client = MagicMock()
+        client_context = MagicMock()
 
         # Initialize enabled client
-        TelemetryClientFactory.initialize_telemetry_client(
-            telemetry_enabled=True,
-            session_id_hex=session_id_hex,
-            auth_provider=auth_provider,
-            host_url="test-host.com",
-            batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE,
-            http_client=mock_http_client,
-        )
+        with patch('databricks.sql.common.unified_http_client.UnifiedHttpClient._setup_pool_manager'):
+            TelemetryClientFactory.initialize_telemetry_client(
+                telemetry_enabled=True,
+                session_id_hex=session_id_hex,
+                auth_provider=auth_provider,
+                host_url="test-host.com",
+                batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE,
+                client_context=client_context,
+            )
 
-        client = TelemetryClientFactory.get_telemetry_client(session_id_hex)
-        assert isinstance(client, TelemetryClient)
-        assert client._session_id_hex == session_id_hex
+            client = TelemetryClientFactory.get_telemetry_client(session_id_hex)
+            assert isinstance(client, TelemetryClient)
+            assert client._session_id_hex == session_id_hex
 
-        # Close client
-        with patch.object(client, "close") as mock_close:
-            TelemetryClientFactory.close(session_id_hex)
-            mock_close.assert_called_once()
+            # Close client
+            with patch.object(client, "close") as mock_close:
+                TelemetryClientFactory.close(session_id_hex)
+                mock_close.assert_called_once()
 
-        # Should get NoopTelemetryClient after close
+            # Should get NoopTelemetryClient after close
 
     def test_disabled_telemetry_creates_noop_client(self):
         """Test that disabled telemetry creates NoopTelemetryClient."""
         session_id_hex = "test-session"
-        mock_http_client = MagicMock()
+        client_context = MagicMock()
 
         TelemetryClientFactory.initialize_telemetry_client(
             telemetry_enabled=False,
             session_id_hex=session_id_hex,
             auth_provider=None,
             host_url="test-host.com",
             batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE,
-            http_client=mock_http_client,
+            client_context=client_context,
         )
 
         client = TelemetryClientFactory.get_telemetry_client(session_id_hex)
@@ -259,7 +262,7 @@ def test_disabled_telemetry_creates_noop_client(self):
     def test_factory_error_handling(self):
         """Test that factory errors fall back to NoopTelemetryClient."""
         session_id = "test-session"
-        mock_http_client = MagicMock()
+        client_context = MagicMock()
 
         # Simulate initialization error
         with patch(
@@ -272,7 +275,7 @@ def test_factory_error_handling(self):
                 auth_provider=AccessTokenAuthProvider("token"),
                 host_url="test-host.com",
                 batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE,
-                http_client=mock_http_client,
+                client_context=client_context,
             )
 
         # Should fall back to NoopTelemetryClient
@@ -283,31 +286,32 @@ def test_factory_shutdown_flow(self):
         """Test factory shutdown when last client is removed."""
         session1 = "session-1"
         session2 = "session-2"
-        mock_http_client = MagicMock()
+        client_context = MagicMock()
 
         # Initialize multiple clients
-        for session in [session1, session2]:
-            TelemetryClientFactory.initialize_telemetry_client(
-                telemetry_enabled=True,
-                session_id_hex=session,
-                auth_provider=AccessTokenAuthProvider("token"),
-                host_url="test-host.com",
-                batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE,
-                http_client=mock_http_client,
-            )
-
-        # Factory should be initialized
-        assert TelemetryClientFactory._initialized is True
-        assert TelemetryClientFactory._executor is not None
-
-        # Close first client - factory should stay initialized
-        TelemetryClientFactory.close(session1)
-        assert TelemetryClientFactory._initialized is True
-
-        # Close second client - factory should shut down
-        TelemetryClientFactory.close(session2)
-        assert TelemetryClientFactory._initialized is False
-        assert TelemetryClientFactory._executor is None
+        with patch('databricks.sql.common.unified_http_client.UnifiedHttpClient._setup_pool_manager'):
+            for session in [session1, session2]:
+                TelemetryClientFactory.initialize_telemetry_client(
+                    telemetry_enabled=True,
+                    session_id_hex=session,
+                    auth_provider=AccessTokenAuthProvider("token"),
+                    host_url="test-host.com",
+                    batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE,
+                    client_context=client_context,
+                )
+
+            # Factory should be initialized
+            assert TelemetryClientFactory._initialized is True
+            assert TelemetryClientFactory._executor is not None
+
+            # Close first client - factory should stay initialized
+            TelemetryClientFactory.close(session1)
+            assert TelemetryClientFactory._initialized is True
+
+            # Close second client - factory should shut down
+            TelemetryClientFactory.close(session2)
+            assert TelemetryClientFactory._initialized is False
+            assert TelemetryClientFactory._executor is None
 
     @patch(
         "databricks.sql.telemetry.telemetry_client.TelemetryClient.export_failure_log"