Merge remote-tracking branch 'target/main' into close-conn

Author: Varun0157

src/databricks/sql/auth/authenticators.py

@@ -199,7 +199,7 @@ 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

@@ -115,18 +115,14 @@ def get_azure_tenant_id_from_host(host: str, http_client) -> str:
     login_url = f"{host}/aad/auth"
     logger.debug("Loading tenant ID from %s", login_url)
 
-    with http_client.request_context(
-        HttpMethod.GET, login_url, allow_redirects=False
-    ) as resp:
-        if resp.status // 100 != 3:
+    with http_client.request_context(HttpMethod.GET, login_url) as resp:
+        entra_id_endpoint = resp.retries.history[-1].redirect_location
+        if entra_id_endpoint is None:
             raise ValueError(
-                f"Failed to get tenant ID from {login_url}: expected status code 3xx, got {resp.status}"
+                f"No Location header in response from {login_url}: {entra_id_endpoint}"
             )
-        entra_id_endpoint = dict(resp.headers).get("Location")
-        if entra_id_endpoint is None:
-            raise ValueError(f"No Location header in response from {login_url}")
 
-    # The Location header has the following form: https://login.microsoftonline.com/<tenant-id>/oauth2/authorize?...
+    # The final redirect URL 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).
     url = urlparse(entra_id_endpoint)
     path_segments = url.path.split("/")

src/databricks/sql/auth/oauth.py

@@ -337,17 +337,17 @@ def refresh(self) -> Token:
             }
         )
 
-        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())
-                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}"
-                )
+        response = self._http_client.request(
+            method=HttpMethod.POST, url=self.token_url, headers=headers, body=data
+        )
+        if response.status == 200:
+            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} {response.data.decode('utf-8')}"
+            )

src/databricks/sql/auth/retry.py

@@ -355,8 +355,14 @@ 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:
-            return False, "2xx codes are not retried"
+        if status_code // 100 <= 3:
+            return False, "2xx/3xx codes are not retried"
+
+        if status_code == 400:
+            return (
+                False,
+                "Received 400 - BAD_REQUEST. Please check the request parameters.",
+            )
 
         if status_code == 401:
             return (

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,9 @@ 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,8 @@ 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
+        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 +456,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 +479,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 +532,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 +545,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/common/retry_test_mixins.py

@@ -346,7 +346,7 @@ def test_retry_dangerous_codes(self, extra_params):
 
         # These http codes are not retried by default
         # For some applications, idempotency is not important so we give users a way to force retries anyway
-        DANGEROUS_CODES = [502, 504, 400]
+        DANGEROUS_CODES = [502, 504]
 
         additional_settings = {
             "_retry_dangerous_codes": DANGEROUS_CODES,

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_auth.py

@@ -263,15 +263,11 @@ def test_get_token_success(self, token_source, http_response):
         with patch.object(token_source, "_http_client", mock_http_client):
             # Create a mock response with the expected format
             mock_response = MagicMock()
-            mock_response.status_code = 200
-            mock_response.json.return_value = {
-                "access_token": "abc123",
-                "token_type": "Bearer",
-                "refresh_token": None,
-            }
-            # Mock the context manager (execute returns context manager)
-            mock_http_client.execute.return_value.__enter__.return_value = mock_response
-            mock_http_client.execute.return_value.__exit__.return_value = None
+            mock_response.status = 200
+            mock_response.data.decode.return_value = '{"access_token": "abc123", "token_type": "Bearer", "refresh_token": null}'
+            
+            # Mock the request method to return the response directly
+            mock_http_client.request.return_value = mock_response
 
             token = token_source.get_token()
 
@@ -287,12 +283,11 @@ def test_get_token_failure(self, token_source, http_response):
         with patch.object(token_source, "_http_client", mock_http_client):
             # Create a mock response with error
             mock_response = MagicMock()
-            mock_response.status_code = 400
-            mock_response.text = "Bad Request"
-            mock_response.json.return_value = {"error": "invalid_client"}
-            # Mock the context manager (execute returns context manager)
-            mock_http_client.execute.return_value.__enter__.return_value = mock_response
-            mock_http_client.execute.return_value.__exit__.return_value = None
+            mock_response.status = 400
+            mock_response.data.decode.return_value = "Bad Request"
+            
+            # Mock the request method to return the response directly
+            mock_http_client.request.return_value = mock_response
 
             with pytest.raises(Exception) as e:
                 token_source.get_token()