more fixes

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

Author: vikrantpuppala

src/databricks/sql/auth/authenticators.py

@@ -199,8 +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,
-            http_client
+            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
-    ) 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}"
-        #     )
+    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"No Location header in response from {login_url}: {entra_id_endpoint}")
+            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 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

@@ -336,15 +336,15 @@ def refresh(self) -> Token:
                 **self.extra_params,
             }
         )
-        
 
         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')))
+                oauth_response = OAuthResponse(
+                    **json.loads(response.data.decode("utf-8"))
+                )
                 return Token(
                     oauth_response.access_token,
                     oauth_response.token_type,

src/databricks/sql/auth/retry.py

@@ -356,7 +356,13 @@ def should_retry(self, method: str, status_code: int) -> Tuple[bool, str]:
 
         # Request succeeded. Don't retry.
         if status_code // 100 <= 3:
-            return False, "2xx codes are not retried"
+            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/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, urllib.parse.urlparse(url).netloc)
+        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

@@ -187,8 +187,6 @@ def __init__(
         self._executor = executor
 
         # 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):

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/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()