more fixes

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

Author: vikrantpuppala

src/databricks/sql/auth/retry.py

@@ -294,7 +294,7 @@ def sleep_for_retry(self, response: BaseHTTPResponse) -> bool:
         else:
             proposed_wait = self.get_backoff_time()
 
-        proposed_wait = min(proposed_wait, self.delay_max)
+        proposed_wait = max(proposed_wait, self.delay_max)
         self.check_proposed_wait(proposed_wait)
         logger.debug(f"Retrying after {proposed_wait} seconds")
         time.sleep(proposed_wait)

src/databricks/sql/backend/thrift_backend.py

@@ -105,7 +105,7 @@ def __init__(
         http_headers,
         auth_provider: AuthProvider,
         ssl_options: SSLOptions,
-        http_client=None,
+        http_client,
         **kwargs,
     ):
         # Internal arguments in **kwargs:

src/databricks/sql/client.py

@@ -277,7 +277,7 @@ def read(self) -> Optional[OAuthToken]:
                 host_url=server_hostname,
                 http_path=http_path,
                 port=kwargs.get("_port", 443),
-                http_client=http_client,
+                client_context=client_context,
                 user_agent=self.session.useragent_header
                 if hasattr(self, "session")
                 else None,
@@ -299,7 +299,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.session.http_client,
+            client_context=client_context,
         )
 
         self._telemetry_client = TelemetryClientFactory.get_telemetry_client(

src/databricks/sql/cloudfetch/downloader.py

@@ -13,9 +13,6 @@
 
 logger = logging.getLogger(__name__)
 
-# TODO: Ideally, we should use a common retry policy (DatabricksRetryPolicy) for all the requests across the library.
-#       But DatabricksRetryPolicy should be updated first - currently it can work only with Thrift requests
-
 
 @dataclass
 class DownloadedFile:

src/databricks/sql/common/unified_http_client.py

@@ -193,38 +193,6 @@ def request(
             response._body = response.data
             return response
 
-    def upload_file(
-        self, url: str, file_path: str, headers: Optional[Dict[str, str]] = None
-    ) -> urllib3.HTTPResponse:
-        """
-        Upload a file using PUT method.
-
-        Args:
-            url: URL to upload to
-            file_path: Path to the file to upload
-            headers: Optional headers
-
-        Returns:
-            urllib3.HTTPResponse: The response from the server
-        """
-        with open(file_path, "rb") as file_obj:
-            return self.request("PUT", url, body=file_obj.read(), headers=headers)
-
-    def download_file(
-        self, url: str, file_path: str, headers: Optional[Dict[str, str]] = None
-    ) -> None:
-        """
-        Download a file using GET method.
-
-        Args:
-            url: URL to download from
-            file_path: Path where to save the downloaded file
-            headers: Optional headers
-        """
-        response = self.request("GET", url, headers=headers)
-        with open(file_path, "wb") as file_obj:
-            file_obj.write(response.data)
-
     def close(self):
         """Close the underlying connection pools."""
         if self._pool_manager:
@@ -236,14 +204,3 @@ def __enter__(self):
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         self.close()
-
-
-# Compatibility class to maintain requests-like interface for OAuth
-class IgnoreNetrcAuth:
-    """
-    Compatibility class for OAuth code that expects requests.auth.AuthBase interface.
-    This is a no-op auth handler since OAuth handles auth differently.
-    """
-
-    def __call__(self, request):
-        return request

src/databricks/sql/telemetry/telemetry_client.py

@@ -1,8 +1,11 @@
 import threading
 import time
 import logging
+import json
 from concurrent.futures import ThreadPoolExecutor
-from typing import Dict, Optional, TYPE_CHECKING
+from concurrent.futures import Future
+from datetime import datetime, timezone
+from typing import List, Dict, Any, Optional, TYPE_CHECKING
 from databricks.sql.telemetry.models.event import (
     TelemetryEvent,
     DriverSystemConfiguration,
@@ -36,8 +39,7 @@
 import locale
 from databricks.sql.telemetry.utils import BaseTelemetryClient
 from databricks.sql.common.feature_flag import FeatureFlagsContextFactory
-
-from src.databricks.sql.common.unified_http_client import UnifiedHttpClient
+from databricks.sql.common.unified_http_client import UnifiedHttpClient
 
 if TYPE_CHECKING:
     from databricks.sql.client import Connection
@@ -151,6 +153,44 @@ def _flush(self):
         pass
 
 
+class TelemetryHttpClientSingleton:
+    """
+    Singleton HTTP client for telemetry operations.
+
+    This ensures that telemetry has its own dedicated HTTP client that
+    is independent of individual connection lifecycles.
+    """
+
+    _instance = None
+    _lock = threading.RLock()
+
+    def __new__(cls):
+        if cls._instance is None:
+            with cls._lock:
+                if cls._instance is None:
+                    cls._instance = super().__new__(cls)
+                    cls._instance._http_client = None
+                    cls._instance._initialized = False
+        return cls._instance
+
+    def get_http_client(self, client_context):
+        """Get or create the singleton HTTP client."""
+        if not self._initialized and client_context:
+            with self._lock:
+                if not self._initialized:
+                    self._http_client = UnifiedHttpClient(client_context)
+                    self._initialized = True
+        return self._http_client
+
+    def close(self):
+        """Close the singleton HTTP client."""
+        with self._lock:
+            if self._http_client:
+                self._http_client.close()
+                self._http_client = None
+                self._initialized = False
+
+
 class TelemetryClient(BaseTelemetryClient):
     """
     Telemetry client class that handles sending telemetry events in batches to the server.
@@ -169,7 +209,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
@@ -182,7 +222,10 @@ def __init__(
         self._driver_connection_params = None
         self._host_url = host_url
         self._executor = executor
-        self._http_client = http_client
+
+        # Use singleton HTTP client for telemetry instead of connection-specific client
+        self._http_client_singleton = TelemetryHttpClientSingleton()
+        self._http_client = self._http_client_singleton.get_http_client(client_context)
 
     def _export_event(self, event):
         """Add an event to the batch queue and flush if batch is full"""
@@ -246,17 +289,24 @@ def _send_telemetry(self, events):
         except Exception as e:
             logger.debug("Failed to submit telemetry request: %s", e)
 
-    def _send_with_unified_client(self, url, data, headers):
+    def _send_with_unified_client(self, url, data, headers, timeout=900):
         """Helper method to send telemetry using the unified HTTP client."""
         try:
             response = self._http_client.request(
-                "POST", url, body=data, headers=headers, timeout=900
+                "POST", url, body=data, headers=headers, timeout=timeout
             )
             # Convert urllib3 response to requests-like response for compatibility
             response.status_code = response.status
+            response.ok = 200 <= response.status < 300
             response.json = (
                 lambda: json.loads(response.data.decode()) if response.data else {}
             )
+            # Add raise_for_status method
+            def raise_for_status():
+                if not response.ok:
+                    raise Exception(f"HTTP {response.status_code}")
+
+            response.raise_for_status = raise_for_status
             return response
         except Exception as e:
             logger.error("Failed to send telemetry with unified client: %s", e)
@@ -452,7 +502,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:
@@ -475,7 +525,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[
@@ -528,7 +578,7 @@ def connection_failure_log(
         host_url: str,
         http_path: str,
         port: int,
-        http_client: UnifiedHttpClient,
+        client_context,
         user_agent: Optional[str] = None,
     ):
         """Send error telemetry when connection creation fails, without requiring a session"""
@@ -541,7 +591,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

@@ -247,7 +247,6 @@ def test_retry_exponential_backoff(self, mock_send_telemetry, extra_params):
         """
         retry_policy = self._retry_policy.copy()
         retry_policy["_retry_delay_min"] = 1
-        retry_policy["_retry_delay_max"] = 10
 
         time_start = time.time()
         with mocked_server_response(
@@ -283,11 +282,9 @@ def test_retry_max_duration_not_exceeded(self, extra_params):
         WHEN the server sends a Retry-After header of 60 seconds
         THEN the connector raises a MaxRetryDurationError
         """
-        retry_policy = self._retry_policy.copy()
-        retry_policy["_retry_delay_max"] = 60
         with mocked_server_response(status=429, headers={"Retry-After": "60"}):
             with pytest.raises(RequestError) as cm:
-                extra_params = {**extra_params, **retry_policy}
+                extra_params = {**extra_params, **self._retry_policy}
                 with self.connection(extra_params=extra_params) as conn:
                     pass
             assert isinstance(cm.value.args[1], MaxRetryDurationError)

tests/e2e/common/staging_ingestion_tests.py

@@ -72,6 +72,7 @@ def test_staging_ingestion_life_cycle(self, ingestion_user):
         extra_params = {
             "staging_allowed_local_path": "/",
             "_retry_stop_after_attempts_count": 1,
+            "_retry_delay_max": 10,
         }
         with self.connection(extra_params=extra_params) as conn:
             cursor = conn.cursor()

tests/e2e/common/uc_volume_tests.py

@@ -72,6 +72,7 @@ def test_uc_volume_life_cycle(self, catalog, schema):
         extra_params = {
             "staging_allowed_local_path": "/",
             "_retry_stop_after_attempts_count": 1,
+            "_retry_delay_max": 10,
         }
         with self.connection(extra_params=extra_params) as conn:
             cursor = conn.cursor()