more fixes

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

Author: vikrantpuppala

src/databricks/sql/auth/auth.py

@@ -19,6 +19,7 @@ def get_auth_provider(cfg: ClientContext, http_client):
                 cfg.hostname,
                 cfg.azure_client_id,
                 cfg.azure_client_secret,
+                http_client,
                 cfg.azure_tenant_id,
                 cfg.azure_workspace_resource_id,
             )
@@ -34,8 +35,8 @@ def get_auth_provider(cfg: ClientContext, http_client):
             cfg.oauth_redirect_port_range,
             cfg.oauth_client_id,
             cfg.oauth_scopes,
+            http_client,
             cfg.auth_type,
-            http_client=http_client,
         )
     elif cfg.access_token is not None:
         return AccessTokenAuthProvider(cfg.access_token)
@@ -54,7 +55,8 @@ def get_auth_provider(cfg: ClientContext, http_client):
                 cfg.oauth_redirect_port_range,
                 cfg.oauth_client_id,
                 cfg.oauth_scopes,
-                http_client=http_client,
+                http_client,
+                cfg.auth_type or "databricks-oauth",
             )
         else:
             raise RuntimeError("No valid authentication settings!")

src/databricks/sql/auth/authenticators.py

@@ -190,6 +190,7 @@ def __init__(
         hostname,
         azure_client_id,
         azure_client_secret,
+        http_client,
         azure_tenant_id=None,
         azure_workspace_resource_id=None,
     ):
@@ -200,6 +201,7 @@ def __init__(
         self.azure_tenant_id = azure_tenant_id or get_azure_tenant_id_from_host(
             hostname
         )
+        self._http_client = http_client
 
     def auth_type(self) -> str:
         return AuthType.AZURE_SP_M2M.value
@@ -209,6 +211,7 @@ def get_token_source(self, resource: str) -> RefreshableTokenSource:
             token_url=f"{self.AZURE_AAD_ENDPOINT}/{self.azure_tenant_id}/{self.AZURE_TOKEN_ENDPOINT}",
             client_id=self.azure_client_id,
             client_secret=self.azure_client_secret,
+            http_client=self._http_client,
             extra_params={"resource": resource},
         )
 

src/databricks/sql/auth/oauth.py

@@ -190,8 +190,7 @@ def __send_auth_code_token_request(
         data = f"{token_request_body}&code_verifier={verifier}"
         return self.__send_token_request(token_request_url, data)
 
-    @staticmethod
-    def __send_token_request(token_request_url, data):
+    def __send_token_request(self, token_request_url, data):
         headers = {
             "Accept": "application/json",
             "Content-Type": "application/x-www-form-urlencoded",
@@ -210,7 +209,7 @@ def __send_refresh_token_request(self, hostname, refresh_token):
         token_request_body = client.prepare_refresh_body(
             refresh_token=refresh_token, client_id=client.client_id
         )
-        return OAuthManager.__send_token_request(token_request_url, token_request_body)
+        return self.__send_token_request(token_request_url, token_request_body)
 
     @staticmethod
     def __get_tokens_from_response(oauth_response):

src/databricks/sql/backend/thrift_backend.py

@@ -8,6 +8,7 @@
 from typing import List, Optional, Union, Any, TYPE_CHECKING
 from uuid import UUID
 
+from databricks.sql.common.unified_http_client import UnifiedHttpClient
 from databricks.sql.result_set import ThriftResultSet
 from databricks.sql.telemetry.models.event import StatementType
 
@@ -105,7 +106,7 @@ def __init__(
         http_headers,
         auth_provider: AuthProvider,
         ssl_options: SSLOptions,
-        http_client,
+        http_client: UnifiedHttpClient,
         **kwargs,
     ):
         # Internal arguments in **kwargs:

src/databricks/sql/client.py

@@ -799,12 +799,6 @@ def _handle_staging_put(
             r = self.connection.session.http_client.request(
                 "PUT", presigned_url, body=fh.read(), headers=headers
             )
-            # Add compatibility attributes for urllib3 response
-            r.status_code = r.status
-            if hasattr(r, "data"):
-                r.content = r.data
-                r.ok = r.status < 400
-                r.text = r.data.decode() if r.data else ""
 
         # fmt: off
         # HTTP status codes
@@ -814,13 +808,15 @@ def _handle_staging_put(
         NO_CONTENT = 204
         # fmt: on
 
-        if r.status_code not in [OK, CREATED, NO_CONTENT, ACCEPTED]:
+        if r.status not in [OK, CREATED, NO_CONTENT, ACCEPTED]:
+            # Decode response data for error message
+            error_text = r.data.decode() if r.data else ""
             raise OperationalError(
-                f"Staging operation over HTTP was unsuccessful: {r.status_code}-{r.text}",
+                f"Staging operation over HTTP was unsuccessful: {r.status}-{error_text}",
                 session_id_hex=self.connection.get_session_id_hex(),
             )
 
-        if r.status_code == ACCEPTED:
+        if r.status == ACCEPTED:
             logger.debug(
                 f"Response code {ACCEPTED} from server indicates ingestion command was accepted "
                 + "but not yet applied on the server. It's possible this command may fail later."
@@ -844,23 +840,19 @@ def _handle_staging_get(
         r = self.connection.session.http_client.request(
             "GET", presigned_url, headers=headers
         )
-        # Add compatibility attributes for urllib3 response
-        r.status_code = r.status
-        if hasattr(r, "data"):
-            r.content = r.data
-            r.ok = r.status < 400
-            r.text = r.data.decode() if r.data else ""
 
         # response.ok verifies the status code is not between 400-600.
         # Any 2xx or 3xx will evaluate r.ok == True
-        if not r.ok:
+        if r.status >= 400:
+            # Decode response data for error message
+            error_text = r.data.decode() if r.data else ""
             raise OperationalError(
-                f"Staging operation over HTTP was unsuccessful: {r.status_code}-{r.text}",
+                f"Staging operation over HTTP was unsuccessful: {r.status}-{error_text}",
                 session_id_hex=self.connection.get_session_id_hex(),
             )
 
         with open(local_file, "wb") as fp:
-            fp.write(r.content)
+            fp.write(r.data)
 
     @log_latency(StatementType.SQL)
     def _handle_staging_remove(
@@ -871,16 +863,12 @@ def _handle_staging_remove(
         r = self.connection.session.http_client.request(
             "DELETE", presigned_url, headers=headers
         )
-        # Add compatibility attributes for urllib3 response
-        r.status_code = r.status
-        if hasattr(r, "data"):
-            r.content = r.data
-            r.ok = r.status < 400
-            r.text = r.data.decode() if r.data else ""
-
-        if not r.ok:
+
+        if r.status >= 400:
+            # Decode response data for error message
+            error_text = r.data.decode() if r.data else ""
             raise OperationalError(
-                f"Staging operation over HTTP was unsuccessful: {r.status_code}-{r.text}",
+                f"Staging operation over HTTP was unsuccessful: {r.status}-{error_text}",
                 session_id_hex=self.connection.get_session_id_hex(),
             )
 

src/databricks/sql/common/feature_flag.py

@@ -113,12 +113,11 @@ def _refresh_flags(self):
             response = self._http_client.request(
                 "GET", self._feature_flag_endpoint, headers=headers, timeout=30
             )
-            # Add compatibility attributes for urllib3 response
-            response.status_code = response.status
-            response.json = lambda: json.loads(response.data.decode())
 
-            if response.status_code == 200:
-                ff_response = FeatureFlagsResponse.from_dict(response.json())
+            if response.status == 200:
+                # Parse JSON response from urllib3 response data
+                response_data = json.loads(response.data.decode())
+                ff_response = FeatureFlagsResponse.from_dict(response_data)
                 self._update_cache_from_response(ff_response)
             else:
                 # On failure, initialize with an empty dictionary to prevent re-blocking.

tests/unit/test_auth.py

@@ -306,6 +306,7 @@ def credential_provider(self):
             hostname="hostname",
             azure_client_id="client_id",
             azure_client_secret="client_secret",
+            http_client=MagicMock(),
             azure_tenant_id="tenant_id",
         )
 

tests/unit/test_retry.py

@@ -34,7 +34,7 @@ def test_sleep__no_retry_after(self, t_mock, retry_policy, error_history):
         retry_policy.history = [error_history, error_history]
         retry_policy.sleep(HTTPResponse(status=503))
 
-        expected_backoff_time = min(
+        expected_backoff_time = max(
             self.calculate_backoff_time(
                 0, retry_policy.delay_min, retry_policy.delay_max
             ),
@@ -57,7 +57,7 @@ def test_sleep__no_retry_after_header__multiple_retries(self, t_mock, retry_poli
         expected_backoff_times = []
         for attempt in range(num_attempts):
             expected_backoff_times.append(
-                min(
+                max(
                     self.calculate_backoff_time(
                         attempt, retry_policy.delay_min, retry_policy.delay_max
                     ),

tests/unit/test_telemetry.py

@@ -19,12 +19,17 @@
 
 
 @pytest.fixture
-def mock_telemetry_client():
+@patch("databricks.sql.telemetry.telemetry_client.TelemetryHttpClientSingleton")
+def mock_telemetry_client(mock_singleton_class):
     """Create a mock telemetry client for testing."""
     session_id = str(uuid.uuid4())
     auth_provider = AccessTokenAuthProvider("test-token")
     executor = MagicMock()
-    mock_http_client = MagicMock()
+    mock_client_context = MagicMock()
+    
+    # Mock the singleton to return a mock HTTP client
+    mock_singleton = mock_singleton_class.return_value
+    mock_singleton.get_http_client.return_value = MagicMock()
 
     return TelemetryClient(
         telemetry_enabled=True,
@@ -33,7 +38,7 @@ def mock_telemetry_client():
         host_url="test-host.com",
         executor=executor,
         batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE,
-        http_client=mock_http_client,
+        client_context=mock_client_context,
     )
 
 
@@ -212,11 +217,16 @@ def telemetry_system_reset(self):
         TelemetryClientFactory._executor = None
         TelemetryClientFactory._initialized = False
 
-    def test_client_lifecycle_flow(self):
+    @patch("databricks.sql.telemetry.telemetry_client.TelemetryHttpClientSingleton")
+    def test_client_lifecycle_flow(self, mock_singleton_class):
         """Test complete client lifecycle: initialize -> use -> close."""
         session_id_hex = "test-session"
         auth_provider = AccessTokenAuthProvider("token")
-        mock_http_client = MagicMock()
+        mock_client_context = MagicMock()
+        
+        # Mock the singleton to return a mock HTTP client
+        mock_singleton = mock_singleton_class.return_value
+        mock_singleton.get_http_client.return_value = MagicMock()
 
         # Initialize enabled client
         TelemetryClientFactory.initialize_telemetry_client(
@@ -225,7 +235,7 @@ def test_client_lifecycle_flow(self):
             auth_provider=auth_provider,
             host_url="test-host.com",
             batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE,
-            http_client=mock_http_client,
+            client_context=mock_client_context,
         )
 
         client = TelemetryClientFactory.get_telemetry_client(session_id_hex)
@@ -241,27 +251,37 @@ def test_client_lifecycle_flow(self):
         client = TelemetryClientFactory.get_telemetry_client(session_id_hex)
         assert isinstance(client, NoopTelemetryClient)
 
-    def test_disabled_telemetry_flow(self):
+    @patch("databricks.sql.telemetry.telemetry_client.TelemetryHttpClientSingleton")
+    def test_disabled_telemetry_flow(self, mock_singleton_class):
         """Test that disabled telemetry uses NoopTelemetryClient."""
         session_id_hex = "test-session"
-        mock_http_client = MagicMock()
+        mock_client_context = MagicMock()
+        
+        # Mock the singleton to return a mock HTTP client
+        mock_singleton = mock_singleton_class.return_value
+        mock_singleton.get_http_client.return_value = 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=mock_client_context,
         )
 
         client = TelemetryClientFactory.get_telemetry_client(session_id_hex)
         assert isinstance(client, NoopTelemetryClient)
 
-    def test_factory_error_handling(self):
+    @patch("databricks.sql.telemetry.telemetry_client.TelemetryHttpClientSingleton")
+    def test_factory_error_handling(self, mock_singleton_class):
         """Test that factory errors fall back to NoopTelemetryClient."""
         session_id = "test-session"
-        mock_http_client = MagicMock()
+        mock_client_context = MagicMock()
+        
+        # Mock the singleton to return a mock HTTP client
+        mock_singleton = mock_singleton_class.return_value
+        mock_singleton.get_http_client.return_value = MagicMock()
 
         # Simulate initialization error
         with patch(
@@ -274,18 +294,23 @@ 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=mock_client_context,
             )
 
         # Should fall back to NoopTelemetryClient
         client = TelemetryClientFactory.get_telemetry_client(session_id)
         assert isinstance(client, NoopTelemetryClient)
 
-    def test_factory_shutdown_flow(self):
+    @patch("databricks.sql.telemetry.telemetry_client.TelemetryHttpClientSingleton")
+    def test_factory_shutdown_flow(self, mock_singleton_class):
         """Test factory shutdown when last client is removed."""
         session1 = "session-1"
         session2 = "session-2"
-        mock_http_client = MagicMock()
+        mock_client_context = MagicMock()
+        
+        # Mock the singleton to return a mock HTTP client
+        mock_singleton = mock_singleton_class.return_value
+        mock_singleton.get_http_client.return_value = MagicMock()
 
         # Initialize multiple clients
         for session in [session1, session2]:
@@ -295,7 +320,7 @@ def test_factory_shutdown_flow(self):
                 auth_provider=AccessTokenAuthProvider("token"),
                 host_url="test-host.com",
                 batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE,
-                http_client=mock_http_client,
+                client_context=mock_client_context,
             )
 
         # Factory should be initialized