Merge branch 'main' into normalise-code

Author: varun-edachali-dbx

src/databricks/sql/backend/thrift_backend.py

@@ -1040,7 +1040,6 @@ def execute_command(
                 max_download_threads=self.max_download_threads,
                 ssl_options=self._ssl_options,
                 has_more_rows=has_more_rows,
-                session_id_hex=self._session_id_hex,
             )
 
     def get_catalogs(
@@ -1079,7 +1078,6 @@ def get_catalogs(
             max_download_threads=self.max_download_threads,
             ssl_options=self._ssl_options,
             has_more_rows=has_more_rows,
-            session_id_hex=self._session_id_hex,
         )
 
     def get_schemas(
@@ -1124,7 +1122,6 @@ def get_schemas(
             max_download_threads=self.max_download_threads,
             ssl_options=self._ssl_options,
             has_more_rows=has_more_rows,
-            session_id_hex=self._session_id_hex,
         )
 
     def get_tables(
@@ -1173,7 +1170,6 @@ def get_tables(
             max_download_threads=self.max_download_threads,
             ssl_options=self._ssl_options,
             has_more_rows=has_more_rows,
-            session_id_hex=self._session_id_hex,
         )
 
     def get_columns(
@@ -1222,7 +1218,6 @@ def get_columns(
             max_download_threads=self.max_download_threads,
             ssl_options=self._ssl_options,
             has_more_rows=has_more_rows,
-            session_id_hex=self._session_id_hex,
         )
 
     def _handle_execute_response(self, resp, cursor):

src/databricks/sql/cloudfetch/downloader.py

@@ -54,12 +54,14 @@ class DownloadableResultSettings:
         link_expiry_buffer_secs (int): Time in seconds to prevent download of a link before it expires. Default 0 secs.
         download_timeout (int): Timeout for download requests. Default 60 secs.
         max_consecutive_file_download_retries (int): Number of consecutive download retries before shutting down.
+        min_cloudfetch_download_speed (float): Threshold in MB/s below which to log warning. Default 0.1 MB/s.
     """
 
     is_lz4_compressed: bool
     link_expiry_buffer_secs: int = 0
     download_timeout: int = 60
     max_consecutive_file_download_retries: int = 0
+    min_cloudfetch_download_speed: float = 0.1
 
 
 class ResultSetDownloadHandler:
@@ -100,6 +102,8 @@ def run(self) -> DownloadedFile:
             self.link, self.settings.link_expiry_buffer_secs
         )
 
+        start_time = time.time()
+
         with self._http_client.execute(
             method=HttpMethod.GET,
             url=self.link.fileLink,
@@ -112,6 +116,13 @@ def run(self) -> DownloadedFile:
 
             # Save (and decompress if needed) the downloaded file
             compressed_data = response.content
+
+            # Log download metrics
+            download_duration = time.time() - start_time
+            self._log_download_metrics(
+                self.link.fileLink, len(compressed_data), download_duration
+            )
+
             decompressed_data = (
                 ResultSetDownloadHandler._decompress_data(compressed_data)
                 if self.settings.is_lz4_compressed
@@ -138,6 +149,32 @@ def run(self) -> DownloadedFile:
                 self.link.rowCount,
             )
 
+    def _log_download_metrics(
+        self, url: str, bytes_downloaded: int, duration_seconds: float
+    ):
+        """Log download speed metrics at INFO/WARN levels."""
+        # Calculate speed in MB/s (ensure float division for precision)
+        speed_mbps = (float(bytes_downloaded) / (1024 * 1024)) / duration_seconds
+
+        urlEndpoint = url.split("?")[0]
+        # INFO level logging
+        logger.info(
+            "CloudFetch download completed: %.4f MB/s, %d bytes in %.3fs from %s",
+            speed_mbps,
+            bytes_downloaded,
+            duration_seconds,
+            urlEndpoint,
+        )
+
+        # WARN level logging if below threshold
+        if speed_mbps < self.settings.min_cloudfetch_download_speed:
+            logger.warning(
+                "CloudFetch download slower than threshold: %.4f MB/s (threshold: %.1f MB/s) from %s",
+                speed_mbps,
+                self.settings.min_cloudfetch_download_speed,
+                url,
+            )
+
     @staticmethod
     def _validate_link(link: TSparkArrowResultLink, expiry_buffer_secs: int):
         """

src/databricks/sql/common/http.py

@@ -5,8 +5,10 @@
 import threading
 from dataclasses import dataclass
 from contextlib import contextmanager
-from typing import Generator
+from typing import Generator, Optional
 import logging
+from requests.adapters import HTTPAdapter
+from databricks.sql.auth.retry import DatabricksRetryPolicy, CommandType
 
 logger = logging.getLogger(__name__)
 
@@ -81,3 +83,70 @@ def execute(
 
     def close(self):
         self.session.close()
+
+
+class TelemetryHTTPAdapter(HTTPAdapter):
+    """
+    Custom HTTP adapter to prepare our DatabricksRetryPolicy before each request.
+    This ensures the retry timer is started and the command type is set correctly,
+    allowing the policy to manage its state for the duration of the request retries.
+    """
+
+    def send(self, request, **kwargs):
+        self.max_retries.command_type = CommandType.OTHER
+        self.max_retries.start_retry_timer()
+        return super().send(request, **kwargs)
+
+
+class TelemetryHttpClient:  # TODO: Unify all the http clients in the PySQL Connector
+    """Singleton HTTP client for sending telemetry data."""
+
+    _instance: Optional["TelemetryHttpClient"] = None
+    _lock = threading.Lock()
+
+    TELEMETRY_RETRY_STOP_AFTER_ATTEMPTS_COUNT = 3
+    TELEMETRY_RETRY_DELAY_MIN = 1.0
+    TELEMETRY_RETRY_DELAY_MAX = 10.0
+    TELEMETRY_RETRY_STOP_AFTER_ATTEMPTS_DURATION = 30.0
+
+    def __init__(self):
+        """Initializes the session and mounts the custom retry adapter."""
+        retry_policy = DatabricksRetryPolicy(
+            delay_min=self.TELEMETRY_RETRY_DELAY_MIN,
+            delay_max=self.TELEMETRY_RETRY_DELAY_MAX,
+            stop_after_attempts_count=self.TELEMETRY_RETRY_STOP_AFTER_ATTEMPTS_COUNT,
+            stop_after_attempts_duration=self.TELEMETRY_RETRY_STOP_AFTER_ATTEMPTS_DURATION,
+            delay_default=1.0,
+            force_dangerous_codes=[],
+        )
+        adapter = TelemetryHTTPAdapter(max_retries=retry_policy)
+        self.session = requests.Session()
+        self.session.mount("https://", adapter)
+        self.session.mount("http://", adapter)
+
+    @classmethod
+    def get_instance(cls) -> "TelemetryHttpClient":
+        """Get the singleton instance of the TelemetryHttpClient."""
+        if cls._instance is None:
+            with cls._lock:
+                if cls._instance is None:
+                    logger.debug("Initializing singleton TelemetryHttpClient")
+                    cls._instance = TelemetryHttpClient()
+        return cls._instance
+
+    def post(self, url: str, **kwargs) -> requests.Response:
+        """
+        Executes a POST request using the configured session.
+
+        This is a blocking call intended to be run in a background thread.
+        """
+        logger.debug("Executing telemetry POST request to: %s", url)
+        return self.session.post(url, **kwargs)
+
+    def close(self):
+        """Closes the underlying requests.Session."""
+        logger.debug("Closing TelemetryHttpClient session.")
+        self.session.close()
+        # Clear the instance to allow for re-initialization if needed
+        with TelemetryHttpClient._lock:
+            TelemetryHttpClient._instance = None

src/databricks/sql/exc.py

@@ -2,8 +2,6 @@
 import logging
 
 logger = logging.getLogger(__name__)
-from databricks.sql.telemetry.telemetry_client import TelemetryClientFactory
-
 
 ### PEP-249 Mandated ###
 # https://peps.python.org/pep-0249/#exceptions
@@ -22,6 +20,8 @@ def __init__(
 
         error_name = self.__class__.__name__
         if session_id_hex:
+            from databricks.sql.telemetry.telemetry_client import TelemetryClientFactory
+
             telemetry_client = TelemetryClientFactory.get_telemetry_client(
                 session_id_hex
             )

src/databricks/sql/result_set.py

@@ -170,7 +170,11 @@ def close(self) -> None:
         been closed on the server for some other reason, issue a request to the server to close it.
         """
         try:
-            self.results.close()
+            if self.results is not None:
+                self.results.close()
+            else:
+                logger.warning("result set close: queue not initialized")
+
             if (
                 self.status != CommandState.CLOSED
                 and not self.has_been_closed_server_side
@@ -193,7 +197,6 @@ def __init__(
         connection: Connection,
         execute_response: ExecuteResponse,
         thrift_client: ThriftDatabricksClient,
-        session_id_hex: Optional[str],
         buffer_size_bytes: int = 104857600,
         arraysize: int = 10000,
         use_cloud_fetch: bool = True,
@@ -217,7 +220,7 @@ def __init__(
             :param ssl_options: SSL options for cloud fetch
             :param has_more_rows: Whether there are more rows to fetch
         """
-        self.num_downloaded_chunks = 0
+        self.num_chunks = 0
 
         # Initialize ThriftResultSet-specific attributes
         self._use_cloud_fetch = use_cloud_fetch
@@ -237,12 +240,12 @@ def __init__(
                 lz4_compressed=execute_response.lz4_compressed,
                 description=execute_response.description,
                 ssl_options=ssl_options,
-                session_id_hex=session_id_hex,
+                session_id_hex=connection.get_session_id_hex(),
                 statement_id=execute_response.command_id.to_hex_guid(),
-                chunk_id=self.num_downloaded_chunks,
+                chunk_id=self.num_chunks,
             )
             if t_row_set.resultLinks:
-                self.num_downloaded_chunks += len(t_row_set.resultLinks)
+                self.num_chunks += len(t_row_set.resultLinks)
 
         # Call parent constructor with common attributes
         super().__init__(
@@ -275,11 +278,11 @@ def _fill_results_buffer(self):
             arrow_schema_bytes=self._arrow_schema_bytes,
             description=self.description,
             use_cloud_fetch=self._use_cloud_fetch,
-            chunk_id=self.num_downloaded_chunks,
+            chunk_id=self.num_chunks,
         )
         self.results = results
         self.has_more_rows = has_more_rows
-        self.num_downloaded_chunks += result_links_count
+        self.num_chunks += result_links_count
 
     def _convert_columnar_table(self, table):
         column_names = [c[0] for c in self.description]

src/databricks/sql/telemetry/telemetry_client.py

@@ -1,9 +1,9 @@
 import threading
 import time
-import requests
 import logging
 from concurrent.futures import ThreadPoolExecutor
 from typing import Dict, Optional
+from databricks.sql.common.http import TelemetryHttpClient
 from databricks.sql.telemetry.models.event import (
     TelemetryEvent,
     DriverSystemConfiguration,
@@ -159,6 +159,7 @@ def __init__(
         self._driver_connection_params = None
         self._host_url = host_url
         self._executor = executor
+        self._http_client = TelemetryHttpClient.get_instance()
 
     def _export_event(self, event):
         """Add an event to the batch queue and flush if batch is full"""
@@ -207,7 +208,7 @@ def _send_telemetry(self, events):
         try:
             logger.debug("Submitting telemetry request to thread pool")
             future = self._executor.submit(
-                requests.post,
+                self._http_client.post,
                 url,
                 data=request.to_json(),
                 headers=headers,
@@ -433,6 +434,7 @@ def close(session_id_hex):
                 )
                 try:
                     TelemetryClientFactory._executor.shutdown(wait=True)
+                    TelemetryHttpClient.close()
                 except Exception as e:
                     logger.debug("Failed to shutdown thread pool executor: %s", e)
                 TelemetryClientFactory._executor = None

tests/unit/test_client.py

@@ -266,7 +266,7 @@ def test_negative_fetch_throws_exception(self):
         mock_backend.fetch_results.return_value = (Mock(), False, 0)
 
         result_set = ThriftResultSet(
-            Mock(), Mock(), mock_backend, session_id_hex=Mock()
+            Mock(), Mock(), mock_backend
         )
 
         with self.assertRaises(ValueError) as e:

tests/unit/test_downloader.py

@@ -23,6 +23,17 @@ class DownloaderTests(unittest.TestCase):
     Unit tests for checking downloader logic.
     """
 
+    def _setup_time_mock_for_download(self, mock_time, end_time):
+        """Helper to setup time mock that handles logging system calls."""
+        call_count = [0]
+        def time_side_effect():
+            call_count[0] += 1
+            if call_count[0] <= 2:  # First two calls (validation, start_time)
+                return 1000
+            else:  # All subsequent calls (logging, duration calculation)
+                return end_time
+        mock_time.side_effect = time_side_effect
+
     @patch("time.time", return_value=1000)
     def test_run_link_expired(self, mock_time):
         settings = Mock()
@@ -90,13 +101,17 @@ def test_run_get_response_not_ok(self, mock_time):
                 d.run()
             self.assertTrue("404" in str(context.exception))
 
-    @patch("time.time", return_value=1000)
+    @patch("time.time")
     def test_run_uncompressed_successful(self, mock_time):
+        self._setup_time_mock_for_download(mock_time, 1000.5)
+        
         http_client = DatabricksHttpClient.get_instance()
         file_bytes = b"1234567890" * 10
         settings = Mock(link_expiry_buffer_secs=0, download_timeout=0, use_proxy=False)
         settings.is_lz4_compressed = False
+        settings.min_cloudfetch_download_speed = 1.0
         result_link = Mock(bytesNum=100, expiryTime=1001)
+        result_link.fileLink = "https://s3.amazonaws.com/bucket/file.arrow?token=abc123"
 
         with patch.object(
             http_client,
@@ -115,15 +130,19 @@ def test_run_uncompressed_successful(self, mock_time):
 
             assert file.file_bytes == b"1234567890" * 10
 
-    @patch("time.time", return_value=1000)
+    @patch("time.time")
     def test_run_compressed_successful(self, mock_time):
+        self._setup_time_mock_for_download(mock_time, 1000.2)
+        
         http_client = DatabricksHttpClient.get_instance()
         file_bytes = b"1234567890" * 10
         compressed_bytes = b'\x04"M\x18h@d\x00\x00\x00\x00\x00\x00\x00#\x14\x00\x00\x00\xaf1234567890\n\x00BP67890\x00\x00\x00\x00'
 
         settings = Mock(link_expiry_buffer_secs=0, download_timeout=0, use_proxy=False)
         settings.is_lz4_compressed = True
+        settings.min_cloudfetch_download_speed = 1.0
         result_link = Mock(bytesNum=100, expiryTime=1001)
+        result_link.fileLink = "https://s3.amazonaws.com/bucket/file.arrow?token=xyz789"
         with patch.object(
             http_client,
             "execute",

tests/unit/test_telemetry.py

@@ -1,6 +1,5 @@
 import uuid
 import pytest
-import requests
 from unittest.mock import patch, MagicMock
 
 from databricks.sql.telemetry.telemetry_client import (
@@ -91,7 +90,7 @@ def test_network_request_flow(self, mock_post, mock_telemetry_client):
         args, kwargs = client._executor.submit.call_args
 
         # Verify correct function and URL
-        assert args[0] == requests.post
+        assert args[0] == client._http_client.post
         assert args[1] == "https://test-host.com/telemetry-ext"
         assert kwargs["headers"]["Authorization"] == "Bearer test-token"