From 9c6cb9eab11aae17d6f34d85b6e98390324a104e Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 12 Nov 2026 10:39:12 +0000 Subject: [PATCH 01/17] perf(thrift): prefetch CloudFetch links while earlier files download Each FetchResults batch waited for its slowest file before the next batch was requested, so on average only about half of the download threads were busy. The download manager now asks the result set for more links whenever a slot frees up and no links are pending (one per free slot plus half a pool), using one thread pool for the whole result. _cloudfetch_prefetch=False restores per-batch fetching. Closing a result set now sets a cancel event that downloads check before the GET and before decompression, and cancels queued downloads. Signed-off-by: Vu Anh Phung Co-authored-by: Isaac --- CHANGELOG.md | 1 + src/databricks/sql/backend/thrift_backend.py | 52 ++- .../sql/cloudfetch/download_manager.py | 41 ++- src/databricks/sql/cloudfetch/downloader.py | 11 + src/databricks/sql/result_set.py | 53 ++- tests/unit/test_cloudfetch_prefetch.py | 334 ++++++++++++++++++ tests/unit/test_fetches.py | 3 +- 7 files changed, 483 insertions(+), 12 deletions(-) create mode 100644 tests/unit/test_cloudfetch_prefetch.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 4038e882d..d6240c029 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,7 @@ - Fix inline parameters (`use_inline_params=True`) so a sequence directly after `IN` renders as a list again: `x IN %(ids)s` renders as `x IN (1,2,3)` instead of `x IN ARRAY(1,2,3)`, which Spark rejects. Inside the list, tuples render as rows for a multi-column `IN`, and a `set` or `frozenset` is accepted again. Sequences anywhere else still render as `ARRAY(...)` ([#629](https://github.com/databricks/databricks-sql-python/issues/629)) - Speed up Thrift response handling: read replies in 1 MiB chunks and decode them with Thrift's C-accelerated protocol, and stop stringifying every request/response for the default no-op debug logger. Wide result schemas, metadata calls with many rows, and inline Arrow results benefit most. - Speed up large CloudFetch results on Thrift: each FetchResults call now asks for about one result file per download thread (at most 4x `buffer_size_bytes`) instead of 100 MiB of files, so all download threads stay busy and the result takes fewer round trips. +- Speed up large CloudFetch results on Thrift further by requesting the next result links while earlier files are still downloading, and stop downloads promptly when a result set is closed early. # 4.9.0 (2026-11-04) - Upgrade Databricks SQL Kernel to 1.4.0, which keeps Lakehouse Real-Time session affinity and flushes shared telemetry only after the final session closes. diff --git a/src/databricks/sql/backend/thrift_backend.py b/src/databricks/sql/backend/thrift_backend.py index ee904c368..05f929f77 100644 --- a/src/databricks/sql/backend/thrift_backend.py +++ b/src/databricks/sql/backend/thrift_backend.py @@ -6,7 +6,7 @@ import re import time import threading -from typing import Dict, List, Optional, Union, Any, TYPE_CHECKING +from typing import Dict, List, Optional, Tuple, Union, Any, TYPE_CHECKING from uuid import UUID from databricks.sql.common.unified_http_client import UnifiedHttpClient @@ -160,6 +160,9 @@ def __init__( # (defaults to None) # max_download_threads # Number of threads for handling cloud fetch downloads. Defaults to 10 + # _cloudfetch_prefetch + # Request the next CloudFetch links while earlier files still download. + # False fetches one batch of links at a time. Defaults to True logger.debug( "ThriftBackend.__init__(server_hostname=%s, port=%s, http_path=%s)" @@ -192,6 +195,7 @@ def __init__( # Cloud fetch self._max_download_threads = kwargs.get("max_download_threads", 10) + self._cloudfetch_prefetch = bool(kwargs.get("_cloudfetch_prefetch", True)) self._ssl_options = ssl_options self._auth_provider = auth_provider @@ -1001,6 +1005,7 @@ def get_execution_result( t_row_set=resp.results, max_download_threads=self.max_download_threads, ssl_options=self._ssl_options, + cloudfetch_prefetch=self._cloudfetch_prefetch, has_more_rows=has_more_rows, ) @@ -1155,6 +1160,7 @@ def execute_command( t_row_set=t_row_set, max_download_threads=self.max_download_threads, ssl_options=self._ssl_options, + cloudfetch_prefetch=self._cloudfetch_prefetch, has_more_rows=has_more_rows, ) @@ -1193,6 +1199,7 @@ def get_catalogs( t_row_set=t_row_set, max_download_threads=self.max_download_threads, ssl_options=self._ssl_options, + cloudfetch_prefetch=self._cloudfetch_prefetch, has_more_rows=has_more_rows, ) @@ -1239,6 +1246,7 @@ def get_schemas( t_row_set=t_row_set, max_download_threads=self.max_download_threads, ssl_options=self._ssl_options, + cloudfetch_prefetch=self._cloudfetch_prefetch, has_more_rows=has_more_rows, ) @@ -1289,6 +1297,7 @@ def get_tables( t_row_set=t_row_set, max_download_threads=self.max_download_threads, ssl_options=self._ssl_options, + cloudfetch_prefetch=self._cloudfetch_prefetch, has_more_rows=has_more_rows, ) @@ -1339,6 +1348,7 @@ def get_columns( t_row_set=t_row_set, max_download_threads=self.max_download_threads, ssl_options=self._ssl_options, + cloudfetch_prefetch=self._cloudfetch_prefetch, has_more_rows=has_more_rows, ) @@ -1374,17 +1384,13 @@ def _handle_execute_response_async(self, resp, cursor): cursor.active_command_id = command_id self._check_direct_results_for_error(resp.directResults, self._host) - def fetch_results( + def _fetch_next( self, command_id: CommandId, max_rows: int, max_bytes: int, expected_row_start_offset: int, - lz4_compressed: bool, - arrow_schema_bytes, - description, - chunk_id: int, - use_cloud_fetch=True, + retryable: bool, ): thrift_handle = command_id.to_thrift_handle() if not thrift_handle: @@ -1404,7 +1410,7 @@ def fetch_results( ) # Fetch results in Inline mode with FETCH_NEXT orientation are not idempotent and hence not retried - resp = self.make_request(self._client.FetchResults, req, use_cloud_fetch) + resp = self.make_request(self._client.FetchResults, req, retryable) if resp.results.startRowOffset > expected_row_start_offset: raise DataError( "fetch_results failed due to inconsistency in the state between the client and the server. Expected results to start from {} but they instead start at {}, some result batches must have been skipped".format( @@ -1412,6 +1418,36 @@ def fetch_results( ), host_url=self._host, ) + return resp + + def fetch_result_links( + self, + command_id: CommandId, + max_rows: int, + max_bytes: int, + expected_row_start_offset: int, + ) -> Tuple[List[ttypes.TSparkArrowResultLink], bool]: + """Fetch the next CloudFetch links of a result: (links, has_more_rows).""" + resp = self._fetch_next( + command_id, max_rows, max_bytes, expected_row_start_offset, True + ) + return resp.results.resultLinks or [], resp.hasMoreRows + + def fetch_results( + self, + command_id: CommandId, + max_rows: int, + max_bytes: int, + expected_row_start_offset: int, + lz4_compressed: bool, + arrow_schema_bytes, + description, + chunk_id: int, + use_cloud_fetch=True, + ): + resp = self._fetch_next( + command_id, max_rows, max_bytes, expected_row_start_offset, use_cloud_fetch + ) queue = ThriftResultSetQueueFactory.build_queue( row_set_type=resp.resultSetMetadata.resultFormat, diff --git a/src/databricks/sql/cloudfetch/download_manager.py b/src/databricks/sql/cloudfetch/download_manager.py index ece7406b7..8b00731e1 100644 --- a/src/databricks/sql/cloudfetch/download_manager.py +++ b/src/databricks/sql/cloudfetch/download_manager.py @@ -1,9 +1,10 @@ from __future__ import annotations import logging +import threading from concurrent.futures import ThreadPoolExecutor, Future -from typing import List, Union, Tuple, Optional, TYPE_CHECKING +from typing import Callable, List, Union, Tuple, Optional, TYPE_CHECKING from databricks.sql.cloudfetch.downloader import ( ResultSetDownloadHandler, @@ -21,6 +22,9 @@ logger = logging.getLogger(__name__) +# Called with the number of free download slots; returns (new links, has more rows). +LinkSource = Callable[[int], Tuple[List["TSparkArrowResultLink"], bool]] + class ResultFileDownloadManager: def __init__( @@ -56,6 +60,17 @@ def __init__( self.session_id_hex = session_id_hex self.statement_id = statement_id self._http_client = http_client + self._link_source: Optional[LinkSource] = None + # Set on shutdown so downloads that have not started (or finished) stop early. + self._cancelled = threading.Event() + + def set_link_source(self, link_source: LinkSource) -> None: + """Fetch further links on demand. + + Whenever a download slot is free and no links are pending, ``link_source`` is + asked for more, so the next links arrive while earlier files still download. + """ + self._link_source = link_source def get_next_downloaded_file( self, next_row_offset: int @@ -98,6 +113,25 @@ def get_next_downloaded_file( return file def _schedule_downloads(self): + self._submit_pending_links() + while ( + self._link_source is not None + and not self._pending_links + and len(self._download_tasks) < self._max_download_threads + ): + links, has_more = self._link_source( + self._max_download_threads - len(self._download_tasks) + ) + if not has_more: + self._link_source = None + for link in links: + self.add_link(link) + self._submit_pending_links() + if not links and self._download_tasks: + # Ask again on the next call instead of polling the server now + break + + def _submit_pending_links(self): """ While download queue has a capacity, peek pending links and submit them to thread pool. """ @@ -119,6 +153,7 @@ def _schedule_downloads(self): session_id_hex=self.session_id_hex, statement_id=self.statement_id, http_client=self._http_client, + cancel_event=self._cancelled, ) task = self._thread_pool.submit(handler.run) self._download_tasks.append(task) @@ -144,6 +179,8 @@ def add_link(self, link: TSparkArrowResultLink): def _shutdown_manager(self): # Clear download handlers and shutdown the thread pool + self._link_source = None self._pending_links = [] self._download_tasks = [] - self._thread_pool.shutdown(wait=False) + self._cancelled.set() + self._thread_pool.shutdown(wait=False, cancel_futures=True) diff --git a/src/databricks/sql/cloudfetch/downloader.py b/src/databricks/sql/cloudfetch/downloader.py index 295f147cd..9a140ea63 100644 --- a/src/databricks/sql/cloudfetch/downloader.py +++ b/src/databricks/sql/cloudfetch/downloader.py @@ -19,6 +19,8 @@ # runtime -- which keeps the (Apache Thrift) ``thrift`` package out of the # cloud-fetch code path used by the SEA and kernel backends. See the # ``test_lazy_thrift_import`` regression test. + import threading + from databricks.sql.thrift_api.TCLIService.ttypes import TSparkArrowResultLink logger = logging.getLogger(__name__) @@ -70,6 +72,7 @@ def __init__( session_id_hex: Optional[str], statement_id: str, http_client, + cancel_event: Optional[threading.Event] = None, ): self.settings = settings self.link = link @@ -78,6 +81,11 @@ def __init__( self.chunk_id = chunk_id self.session_id_hex = session_id_hex self.statement_id = statement_id + self._cancel_event = cancel_event + + def _raise_if_cancelled(self): + if self._cancel_event is not None and self._cancel_event.is_set(): + raise Error("CloudFetch download cancelled: the result set was closed") @log_latency(StatementType.QUERY) def run(self) -> DownloadedFile: @@ -100,6 +108,7 @@ def run(self) -> DownloadedFile: self.link, self.settings.link_expiry_buffer_secs ) + self._raise_if_cancelled() start_time = time.time() with self._http_client.request_context( @@ -111,6 +120,8 @@ def run(self) -> DownloadedFile: if response.status >= 400: raise Exception(f"HTTP {response.status}: {response.data.decode()}") compressed_data = response.data + # Skip decompressing a file nobody will read + self._raise_if_cancelled() # Log download metrics download_duration = time.time() - start_time diff --git a/src/databricks/sql/result_set.py b/src/databricks/sql/result_set.py index 2a2599259..2212aa6b2 100644 --- a/src/databricks/sql/result_set.py +++ b/src/databricks/sql/result_set.py @@ -1,7 +1,7 @@ from __future__ import annotations from abc import ABC, abstractmethod -from typing import List, Optional, TYPE_CHECKING, Tuple +from typing import List, Optional, TYPE_CHECKING, Tuple, cast import logging import pandas @@ -214,6 +214,7 @@ def __init__( max_download_threads: int = 10, ssl_options=None, has_more_rows: bool = True, + cloudfetch_prefetch: bool = True, ): """ Initialize a ThriftResultSet with direct access to the ThriftDatabricksClient. @@ -230,6 +231,8 @@ def __init__( :param max_download_threads: Maximum number of download threads for cloud fetch :param ssl_options: SSL options for cloud fetch :param has_more_rows: Whether there are more rows to fetch + :param cloudfetch_prefetch: Request further CloudFetch links while earlier + files still download, instead of one batch at a time """ self.num_chunks = 0 @@ -238,6 +241,8 @@ def __init__( self.has_more_rows = has_more_rows self._max_download_threads = max_download_threads self._max_link_bytes = 0 + self._cloudfetch_prefetch = cloudfetch_prefetch + self._links_end_offset = 0 # Build the results queue if t_row_set is provided results_queue = None @@ -283,6 +288,51 @@ def __init__( # Initialize results queue if not provided if not self.results: self._fill_results_buffer() + else: + self._start_prefetch() + + def _start_prefetch(self) -> None: + # Hand the remaining CloudFetch links to the queue: its download manager then + # asks for more links whenever a download slot frees up, rather than this result + # set fetching the next batch only after the current one is fully consumed. + queue = self.results + if ( + not self._cloudfetch_prefetch + or not isinstance(queue, ThriftCloudFetchQueue) + or not queue.result_links + or not self.has_more_rows + or self.has_been_closed_server_side + ): + return + self._links_end_offset = max( + link.startRowOffset + link.rowCount for link in queue.result_links + ) + queue.download_manager.set_link_source(self._fetch_more_links) + self.has_more_rows = False + + def _fetch_more_links(self, free_slots: int): + # One link per free slot plus half a pool queued ahead, so slots refill without a + # round trip each time. A queued link waits for at most that many files to be + # consumed before it starts downloading. + wanted = (free_slots + self._max_download_threads // 2) * self._max_link_bytes + max_bytes = min( + max(self.buffer_size_bytes, wanted), + _CLOUDFETCH_MAX_BUFFERS_PER_FETCH * self.buffer_size_bytes, + ) + backend = cast("ThriftDatabricksClient", self.backend) + links, has_more_rows = backend.fetch_result_links( + command_id=self.command_id, + max_rows=self.arraysize, + max_bytes=max_bytes, + expected_row_start_offset=self._links_end_offset, + ) + if links: + self._note_link_sizes(links) + self.num_chunks += len(links) + self._links_end_offset = max( + link.startRowOffset + link.rowCount for link in links + ) + return links, has_more_rows def _note_link_sizes(self, links) -> None: for link in links: @@ -322,6 +372,7 @@ def _fill_results_buffer(self): self.num_chunks += result_links_count if isinstance(results, ThriftCloudFetchQueue): self._note_link_sizes(results.result_links) + self._start_prefetch() def _convert_columnar_table(self, table): column_names = [c[0] for c in self.description] diff --git a/tests/unit/test_cloudfetch_prefetch.py b/tests/unit/test_cloudfetch_prefetch.py new file mode 100644 index 000000000..ff113ecbd --- /dev/null +++ b/tests/unit/test_cloudfetch_prefetch.py @@ -0,0 +1,334 @@ +"""CloudFetch link prefetching across batches and cancel-on-close.""" + +import threading +import unittest +from unittest.mock import MagicMock, Mock, patch + +import databricks.sql.cloudfetch.download_manager as download_manager +import databricks.sql.cloudfetch.downloader as downloader +from databricks.sql.auth.authenticators import AuthProvider +from databricks.sql.backend.thrift_backend import ThriftDatabricksClient +from databricks.sql.backend.types import ExecuteResponse +from databricks.sql.cloudfetch.downloader import DownloadedFile +from databricks.sql.exc import DataError, Error +from databricks.sql.result_set import ThriftResultSet +from databricks.sql.thrift_api.TCLIService import ttypes +from databricks.sql.types import SSLOptions +from databricks.sql.utils import ArrowQueue, ThriftCloudFetchQueue + +MIB = 2**20 + + +def _link(start, rows=10, size=20 * MIB): + return ttypes.TSparkArrowResultLink( + fileLink=f"https://example.com/{start}", + expiryTime=2**62, + startRowOffset=start, + rowCount=rows, + bytesNum=size, + ) + + +class _InstantHandler: + """Stands in for ResultSetDownloadHandler: "downloads" a link immediately.""" + + def __init__(self, link, **_): + self.link = link + + def run(self): + return DownloadedFile(b"", self.link.startRowOffset, self.link.rowCount) + + +class DownloadManagerPrefetchTests(unittest.TestCase): + def _manager(self, links, max_download_threads): + return download_manager.ResultFileDownloadManager( + links, + max_download_threads, + lz4_compressed=False, + ssl_options=SSLOptions(), + session_id_hex=None, + statement_id="s", + chunk_id=0, + http_client=MagicMock(), + ) + + @patch.object(download_manager, "ResultSetDownloadHandler", _InstantHandler) + def test_refills_from_link_source_whenever_a_slot_frees(self): + manager = self._manager([_link(0), _link(10)], max_download_threads=4) + replies = [ + ([_link(20), _link(30), _link(40)], True), + ([_link(50)], True), + ([], False), + ] + free_slots = [] + + def source(free): + free_slots.append(free) + self.assertLessEqual(len(manager._download_tasks), 4) + return replies.pop(0) + + manager.set_link_source(source) + offsets = [] + while True: + file = manager.get_next_downloaded_file(len(offsets) * 10) + if file is None: + break + self.assertLessEqual(len(manager._download_tasks), 4) + offsets.append(file.start_row_offset) + + self.assertEqual(offsets, [0, 10, 20, 30, 40, 50]) + # Asked only when a slot was free and no link was pending + self.assertEqual(free_slots, [2, 1, 1]) + + @patch.object(download_manager, "ResultSetDownloadHandler", _InstantHandler) + def test_empty_reply_is_retried_on_the_next_call_not_in_a_loop(self): + manager = self._manager([_link(0)], max_download_threads=2) + source = Mock(side_effect=[([], True), ([_link(10)], False)]) + manager.set_link_source(source) + + self.assertEqual(manager.get_next_downloaded_file(0).start_row_offset, 0) + self.assertEqual(source.call_count, 1) + self.assertEqual(manager.get_next_downloaded_file(10).start_row_offset, 10) + self.assertIsNone(manager.get_next_downloaded_file(20)) + self.assertEqual(source.call_count, 2) + + def test_shutdown_cancels_downloads_and_stops_fetching_links(self): + manager = self._manager([_link(0), _link(10)], max_download_threads=1) + manager.set_link_source(Mock()) + with patch.object( + manager._thread_pool, "shutdown", wraps=manager._thread_pool.shutdown + ) as shutdown: + manager._shutdown_manager() + + self.assertTrue(manager._cancelled.is_set()) + self.assertIsNone(manager._link_source) + self.assertEqual(manager._pending_links, []) + shutdown.assert_called_once_with(wait=False, cancel_futures=True) + + def test_handlers_share_the_managers_cancel_event(self): + manager = self._manager([_link(0)], max_download_threads=1) + with patch.object(download_manager, "ResultSetDownloadHandler") as handler: + manager._schedule_downloads() + self.assertIs(handler.call_args.kwargs["cancel_event"], manager._cancelled) + manager._shutdown_manager() + + +class DownloaderCancelTests(unittest.TestCase): + def _handler(self, http_client, event): + return downloader.ResultSetDownloadHandler( + downloader.DownloadableResultSettings(is_lz4_compressed=True), + _link(0), + ssl_options=SSLOptions(), + chunk_id=0, + session_id_hex=None, + statement_id="s", + http_client=http_client, + cancel_event=event, + ) + + def test_cancelled_before_the_request_skips_the_get(self): + http_client = MagicMock() + event = threading.Event() + event.set() + + with self.assertRaises(Error): + self._handler(http_client, event).run() + http_client.request_context.assert_not_called() + + def test_cancelled_during_the_download_skips_decompression(self): + event = threading.Event() + + class Response: + status = 200 + + @property + def data(self): + event.set() + return b"compressed" + + http_client = MagicMock() + http_client.request_context.return_value.__enter__.return_value = Response() + with patch.object( + downloader.ResultSetDownloadHandler, "_decompress_data" + ) as decompress: + with self.assertRaises(Error): + self._handler(http_client, event).run() + decompress.assert_not_called() + + +class ResultSetPrefetchTests(unittest.TestCase): + BUFFER = 100 * MIB + + def _result_set(self, initial_links, prefetch=True, has_more_rows=True, queue=None): + backend = Mock(spec=ThriftDatabricksClient) + if queue is None: + queue = Mock(spec=ThriftCloudFetchQueue) + queue.result_links = initial_links + queue.download_manager = Mock() + t_row_set = ttypes.TRowSet(startRowOffset=0, rows=[], resultLinks=initial_links) + with patch( + "databricks.sql.utils.ThriftResultSetQueueFactory.build_queue", + return_value=queue, + ): + rs = ThriftResultSet( + connection=Mock(), + execute_response=ExecuteResponse( + command_id=Mock(), + status=None, + has_been_closed_server_side=False, + description=[], + result_format=ttypes.TSparkRowSetType.URL_BASED_SET, + ), + thrift_client=backend, + buffer_size_bytes=self.BUFFER, + arraysize=1000, + t_row_set=t_row_set, + max_download_threads=10, + has_more_rows=has_more_rows, + cloudfetch_prefetch=prefetch, + ) + return rs, backend, queue + + def test_hands_link_fetching_to_the_queue(self): + rs, backend, queue = self._result_set([_link(0, 100), _link(100, 100)]) + + queue.download_manager.set_link_source.assert_called_once_with( + rs._fetch_more_links + ) + self.assertFalse(rs.has_more_rows) + backend.fetch_results.assert_not_called() + + def test_link_source_asks_for_free_slots_plus_half_a_pool(self): + rs, backend, queue = self._result_set([_link(0, 100), _link(100, 100)]) + source = queue.download_manager.set_link_source.call_args.args[0] + backend.fetch_result_links.side_effect = [ + ([_link(200, 50)], True), + ([_link(250, 50)], False), + ] + + self.assertEqual(source(3), ([_link(200, 50)], True)) + self.assertEqual(source(1), ([_link(250, 50)], False)) + + first, second = backend.fetch_result_links.call_args_list + self.assertEqual(first.kwargs["max_bytes"], 8 * 20 * MIB) + self.assertEqual(first.kwargs["expected_row_start_offset"], 200) + self.assertEqual(first.kwargs["max_rows"], 1000) + self.assertEqual(second.kwargs["max_bytes"], 6 * 20 * MIB) + self.assertEqual(second.kwargs["expected_row_start_offset"], 250) + + def test_link_source_budget_stays_between_one_and_four_buffers(self): + rs, backend, queue = self._result_set([_link(0, size=60 * MIB)]) + source = queue.download_manager.set_link_source.call_args.args[0] + backend.fetch_result_links.return_value = ([], True) + + source(10) + rs._max_link_bytes = MIB + source(1) + + budgets = [ + c.kwargs["max_bytes"] for c in backend.fetch_result_links.call_args_list + ] + self.assertEqual(budgets, [4 * self.BUFFER, self.BUFFER]) + + def test_kill_switch_keeps_per_batch_fetching(self): + rs, backend, queue = self._result_set([_link(0)], prefetch=False) + + queue.download_manager.set_link_source.assert_not_called() + self.assertTrue(rs.has_more_rows) + + def test_inline_results_are_not_prefetched(self): + rs, backend, _ = self._result_set([], queue=Mock(spec=ArrowQueue)) + + self.assertTrue(rs.has_more_rows) + backend.fetch_result_links.assert_not_called() + + def test_complete_results_are_not_prefetched(self): + rs, _, queue = self._result_set([_link(0)], has_more_rows=False) + + queue.download_manager.set_link_source.assert_not_called() + + def test_prefetch_starts_after_the_first_fetch_without_direct_results(self): + backend = Mock(spec=ThriftDatabricksClient) + queue = Mock(spec=ThriftCloudFetchQueue) + queue.result_links = [_link(0)] + queue.download_manager = Mock() + backend.fetch_results.return_value = (queue, True, 1) + + rs = ThriftResultSet( + connection=Mock(), + execute_response=ExecuteResponse( + command_id=Mock(), + status=None, + has_been_closed_server_side=False, + description=[], + ), + thrift_client=backend, + ) + + queue.download_manager.set_link_source.assert_called_once() + self.assertFalse(rs.has_more_rows) + + +class ThriftClientPrefetchTests(unittest.TestCase): + def _client(self, **kwargs): + return ThriftDatabricksClient( + "foobar", + 443, + "path", + [], + auth_provider=AuthProvider(), + ssl_options=SSLOptions(), + http_client=MagicMock(), + **kwargs, + ) + + def test_prefetch_is_on_by_default_and_can_be_disabled(self): + self.assertTrue(self._client()._cloudfetch_prefetch) + self.assertFalse(self._client(_cloudfetch_prefetch=False)._cloudfetch_prefetch) + + @patch("databricks.sql.backend.thrift_backend.ThriftResultSet") + @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True) + def test_execute_passes_the_setting_to_the_result_set(self, _, result_set): + client = self._client(_cloudfetch_prefetch=False) + client._handle_execute_response = Mock(return_value=(Mock(), Mock())) + + client.execute_command("foo", Mock(), 100, 200, Mock(), Mock(), Mock()) + + self.assertIs(result_set.call_args.kwargs["cloudfetch_prefetch"], False) + + @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True) + def test_fetch_result_links(self, tcli_service_class): + links = [_link(200, 50)] + tcli_service_class.return_value.FetchResults.return_value = ( + ttypes.TFetchResultsResp( + status=ttypes.TStatus(statusCode=ttypes.TStatusCode.SUCCESS_STATUS), + hasMoreRows=True, + results=ttypes.TRowSet(startRowOffset=200, rows=[], resultLinks=links), + ) + ) + client = self._client() + + got = client.fetch_result_links( + command_id=Mock(), + max_rows=1000, + max_bytes=123, + expected_row_start_offset=200, + ) + + self.assertEqual(got, (links, True)) + req = tcli_service_class.return_value.FetchResults.call_args.args[0] + self.assertEqual( + (req.maxRows, req.maxBytes, req.orientation), + (1000, 123, ttypes.TFetchOrientation.FETCH_NEXT), + ) + with self.assertRaises(DataError): + client.fetch_result_links( + command_id=Mock(), + max_rows=1, + max_bytes=1, + expected_row_start_offset=100, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/test_fetches.py b/tests/unit/test_fetches.py index d0d74df1d..8bffc8183 100644 --- a/tests/unit/test_fetches.py +++ b/tests/unit/test_fetches.py @@ -393,7 +393,7 @@ def test_fetchmany_arrow_all_empty_returns_zero_row_table(self): class CloudFetchBatchSizeTests(unittest.TestCase): - """maxBytes sent by ThriftResultSet._fill_results_buffer.""" + """maxBytes sent by ThriftResultSet._fill_results_buffer (prefetch disabled).""" MIB = 2**20 BUFFER = 100 * MIB @@ -457,6 +457,7 @@ def fetch_results(**kwargs): arraysize=1000, t_row_set=t_row_set, max_download_threads=max_download_threads, + cloudfetch_prefetch=False, ) return rs, calls From 117a095c63c568caf0b24899bd51bc510fa845f4 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 12 Nov 2026 11:18:00 +0000 Subject: [PATCH 02/17] Guard prefetched link fetches and a shut-down queue fetch_result_links raises DataError if a reply carries rows that are not CloudFetch links, instead of dropping them. set_link_source refuses a manager that is already shut down, so such a result keeps per-batch fetching. Pin that a failed link fetch reaches the caller and that fetching can resume. Signed-off-by: Vu Anh Phung Co-authored-by: Isaac --- src/databricks/sql/backend/thrift_backend.py | 17 ++++- .../sql/cloudfetch/download_manager.py | 6 +- src/databricks/sql/result_set.py | 4 +- tests/unit/test_cloudfetch_prefetch.py | 69 +++++++++++++++++++ 4 files changed, 92 insertions(+), 4 deletions(-) diff --git a/src/databricks/sql/backend/thrift_backend.py b/src/databricks/sql/backend/thrift_backend.py index 05f929f77..5a42058cd 100644 --- a/src/databricks/sql/backend/thrift_backend.py +++ b/src/databricks/sql/backend/thrift_backend.py @@ -1431,7 +1431,22 @@ def fetch_result_links( resp = self._fetch_next( command_id, max_rows, max_bytes, expected_row_start_offset, True ) - return resp.results.resultLinks or [], resp.hasMoreRows + results, metadata = resp.results, resp.resultSetMetadata + # Rows in any other form would be dropped silently, so fail loudly instead + if ( + results.arrowBatches + or results.columns + or ( + metadata is not None + and metadata.resultFormat is not None + and metadata.resultFormat != ttypes.TSparkRowSetType.URL_BASED_SET + ) + ): + raise DataError( + "Expected only CloudFetch links when fetching more results of a CloudFetch result", + host_url=self._host, + ) + return results.resultLinks or [], resp.hasMoreRows def fetch_results( self, diff --git a/src/databricks/sql/cloudfetch/download_manager.py b/src/databricks/sql/cloudfetch/download_manager.py index 8b00731e1..f40dee9a1 100644 --- a/src/databricks/sql/cloudfetch/download_manager.py +++ b/src/databricks/sql/cloudfetch/download_manager.py @@ -64,13 +64,17 @@ def __init__( # Set on shutdown so downloads that have not started (or finished) stop early. self._cancelled = threading.Event() - def set_link_source(self, link_source: LinkSource) -> None: + def set_link_source(self, link_source: LinkSource) -> bool: """Fetch further links on demand. Whenever a download slot is free and no links are pending, ``link_source`` is asked for more, so the next links arrive while earlier files still download. + Returns False, without attaching it, if the manager is already shut down. """ + if self._cancelled.is_set(): + return False self._link_source = link_source + return True def get_next_downloaded_file( self, next_row_offset: int diff --git a/src/databricks/sql/result_set.py b/src/databricks/sql/result_set.py index 2212aa6b2..73c3c0811 100644 --- a/src/databricks/sql/result_set.py +++ b/src/databricks/sql/result_set.py @@ -307,8 +307,8 @@ def _start_prefetch(self) -> None: self._links_end_offset = max( link.startRowOffset + link.rowCount for link in queue.result_links ) - queue.download_manager.set_link_source(self._fetch_more_links) - self.has_more_rows = False + if queue.download_manager.set_link_source(self._fetch_more_links): + self.has_more_rows = False def _fetch_more_links(self, free_slots: int): # One link per free slot plus half a pool queued ahead, so slots refill without a diff --git a/tests/unit/test_cloudfetch_prefetch.py b/tests/unit/test_cloudfetch_prefetch.py index ff113ecbd..c176946d2 100644 --- a/tests/unit/test_cloudfetch_prefetch.py +++ b/tests/unit/test_cloudfetch_prefetch.py @@ -105,6 +105,32 @@ def test_shutdown_cancels_downloads_and_stops_fetching_links(self): self.assertEqual(manager._pending_links, []) shutdown.assert_called_once_with(wait=False, cancel_futures=True) + @patch.object(download_manager, "ResultSetDownloadHandler", _InstantHandler) + def test_link_source_errors_reach_the_caller_and_fetching_can_resume(self): + manager = self._manager([_link(0)], max_download_threads=2) + source = Mock( + side_effect=[RuntimeError("FetchResults failed"), ([_link(10)], False)] + ) + manager.set_link_source(source) + + with self.assertRaises(RuntimeError): + manager.get_next_downloaded_file(0) + # The scheduled download is kept and the next call asks for links again + offsets = [manager.get_next_downloaded_file(0).start_row_offset] + offsets.append(manager.get_next_downloaded_file(10).start_row_offset) + self.assertIsNone(manager.get_next_downloaded_file(20)) + self.assertEqual(offsets, [0, 10]) + self.assertEqual(source.call_count, 2) + + @patch.object(download_manager, "ResultSetDownloadHandler", _InstantHandler) + def test_link_source_is_refused_after_shutdown(self): + # No schedulable links: the first request already shuts the manager down + manager = self._manager([_link(0, rows=0)], max_download_threads=2) + self.assertIsNone(manager.get_next_downloaded_file(0)) + + self.assertFalse(manager.set_link_source(Mock())) + self.assertIsNone(manager._link_source) + def test_handlers_share_the_managers_cancel_event(self): manager = self._manager([_link(0)], max_download_threads=1) with patch.object(download_manager, "ResultSetDownloadHandler") as handler: @@ -236,6 +262,16 @@ def test_kill_switch_keeps_per_batch_fetching(self): queue.download_manager.set_link_source.assert_not_called() self.assertTrue(rs.has_more_rows) + def test_shut_down_queue_keeps_per_batch_fetching(self): + queue = Mock(spec=ThriftCloudFetchQueue) + queue.result_links = [_link(0)] + queue.download_manager = Mock() + queue.download_manager.set_link_source.return_value = False + + rs, _, _ = self._result_set([_link(0)], queue=queue) + + self.assertTrue(rs.has_more_rows) + def test_inline_results_are_not_prefetched(self): rs, backend, _ = self._result_set([], queue=Mock(spec=ArrowQueue)) @@ -329,6 +365,39 @@ def test_fetch_result_links(self, tcli_service_class): expected_row_start_offset=100, ) + @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True) + def test_fetch_result_links_rejects_rows_that_are_not_links( + self, tcli_service_class + ): + ok = ttypes.TStatus(statusCode=ttypes.TStatusCode.SUCCESS_STATUS) + inline_rows = ttypes.TRowSet( + startRowOffset=0, + rows=[], + arrowBatches=[ttypes.TSparkArrowBatch(batch=b"x", rowCount=1)], + ) + arrow_format = ttypes.TGetResultSetMetadataResp( + status=ok, resultFormat=ttypes.TSparkRowSetType.ARROW_BASED_SET + ) + client = self._client() + + for resp in ( + ttypes.TFetchResultsResp(status=ok, hasMoreRows=False, results=inline_rows), + ttypes.TFetchResultsResp( + status=ok, + hasMoreRows=False, + results=ttypes.TRowSet(startRowOffset=0, rows=[]), + resultSetMetadata=arrow_format, + ), + ): + tcli_service_class.return_value.FetchResults.return_value = resp + with self.assertRaises(DataError): + client.fetch_result_links( + command_id=Mock(), + max_rows=1, + max_bytes=1, + expected_row_start_offset=0, + ) + if __name__ == "__main__": unittest.main() From 0dede5d210960646b2e3e253b032a9c1292a8580 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 12 Nov 2026 11:24:43 +0000 Subject: [PATCH 03/17] Don't re-poll for links when a reply only had empty links Base the early exit on whether a download was scheduled, so a reply whose links all have zero rows waits for the next call like an empty reply. Add end-to-end tests that run ThriftResultSet over a real CloudFetch queue and download manager across several link fetches. Signed-off-by: Vu Anh Phung Co-authored-by: Isaac --- .../sql/cloudfetch/download_manager.py | 10 +- tests/unit/test_cloudfetch_prefetch.py | 109 ++++++++++++++++++ 2 files changed, 114 insertions(+), 5 deletions(-) diff --git a/src/databricks/sql/cloudfetch/download_manager.py b/src/databricks/sql/cloudfetch/download_manager.py index f40dee9a1..1696ef302 100644 --- a/src/databricks/sql/cloudfetch/download_manager.py +++ b/src/databricks/sql/cloudfetch/download_manager.py @@ -123,16 +123,16 @@ def _schedule_downloads(self): and not self._pending_links and len(self._download_tasks) < self._max_download_threads ): - links, has_more = self._link_source( - self._max_download_threads - len(self._download_tasks) - ) + scheduled = len(self._download_tasks) + links, has_more = self._link_source(self._max_download_threads - scheduled) if not has_more: self._link_source = None for link in links: self.add_link(link) self._submit_pending_links() - if not links and self._download_tasks: - # Ask again on the next call instead of polling the server now + if self._download_tasks and len(self._download_tasks) == scheduled: + # Nothing new to download (no links, or only empty ones): ask again + # on the next call instead of polling the server now break def _submit_pending_links(self): diff --git a/tests/unit/test_cloudfetch_prefetch.py b/tests/unit/test_cloudfetch_prefetch.py index c176946d2..48e481db2 100644 --- a/tests/unit/test_cloudfetch_prefetch.py +++ b/tests/unit/test_cloudfetch_prefetch.py @@ -4,6 +4,13 @@ import unittest from unittest.mock import MagicMock, Mock, patch +import pytest + +try: + import pyarrow as pa +except ImportError: + pa = None + import databricks.sql.cloudfetch.download_manager as download_manager import databricks.sql.cloudfetch.downloader as downloader from databricks.sql.auth.authenticators import AuthProvider @@ -39,6 +46,18 @@ def run(self): return DownloadedFile(b"", self.link.startRowOffset, self.link.rowCount) +class _ArrowFileHandler(_InstantHandler): + """Returns an Arrow IPC stream with ids startRowOffset..startRowOffset+rowCount-1.""" + + def run(self): + start, rows = self.link.startRowOffset, self.link.rowCount + table = pa.table({"id": pa.array(range(start, start + rows), pa.int64())}) + sink = pa.BufferOutputStream() + with pa.ipc.new_stream(sink, table.schema) as writer: + writer.write_table(table) + return DownloadedFile(sink.getvalue().to_pybytes(), start, rows) + + class DownloadManagerPrefetchTests(unittest.TestCase): def _manager(self, links, max_download_threads): return download_manager.ResultFileDownloadManager( @@ -105,6 +124,17 @@ def test_shutdown_cancels_downloads_and_stops_fetching_links(self): self.assertEqual(manager._pending_links, []) shutdown.assert_called_once_with(wait=False, cancel_futures=True) + @patch.object(download_manager, "ResultSetDownloadHandler", _InstantHandler) + def test_reply_with_only_empty_links_is_retried_on_the_next_call(self): + manager = self._manager([_link(0)], max_download_threads=2) + source = Mock(side_effect=[([_link(10, rows=0)], True), ([_link(10)], False)]) + manager.set_link_source(source) + + self.assertEqual(manager.get_next_downloaded_file(0).start_row_offset, 0) + self.assertEqual(source.call_count, 1) + self.assertEqual(manager.get_next_downloaded_file(10).start_row_offset, 10) + self.assertIsNone(manager.get_next_downloaded_file(20)) + @patch.object(download_manager, "ResultSetDownloadHandler", _InstantHandler) def test_link_source_errors_reach_the_caller_and_fetching_can_resume(self): manager = self._manager([_link(0)], max_download_threads=2) @@ -305,6 +335,85 @@ def test_prefetch_starts_after_the_first_fetch_without_direct_results(self): self.assertFalse(rs.has_more_rows) +@pytest.mark.skipif(pa is None, reason="PyArrow is not installed") +@patch.object(download_manager, "ResultSetDownloadHandler", _ArrowFileHandler) +class ResultSetPrefetchEndToEndTests(unittest.TestCase): + """Real ThriftCloudFetchQueue and download manager; only the download is stubbed.""" + + ROWS_PER_LINK = 10 + + def _result_set(self, initial_links=2, total_links=11, links_per_reply=3): + links = [ + _link(i * self.ROWS_PER_LINK, self.ROWS_PER_LINK) + for i in range(total_links) + ] + rest = links[initial_links:] + expected_offsets = [] + + def fetch_result_links(**kwargs): + expected_offsets.append(kwargs["expected_row_start_offset"]) + reply = rest[:links_per_reply] + del rest[:links_per_reply] + return reply, bool(rest) + + backend = Mock(spec=ThriftDatabricksClient) + backend.fetch_result_links.side_effect = fetch_result_links + rs = ThriftResultSet( + connection=Mock(), + execute_response=ExecuteResponse( + command_id=Mock(), + status=None, + has_been_closed_server_side=False, + description=[("id", "bigint", None, None, None, None, None)], + lz4_compressed=False, + arrow_schema_bytes=pa.schema([("id", pa.int64())]) + .serialize() + .to_pybytes(), + result_format=ttypes.TSparkRowSetType.URL_BASED_SET, + ), + thrift_client=backend, + arraysize=1000, + t_row_set=ttypes.TRowSet( + startRowOffset=0, rows=[], resultLinks=links[:initial_links] + ), + max_download_threads=4, + ) + return rs, backend, expected_offsets + + def test_fetchall_arrow_reads_every_link_in_order(self): + rs, backend, expected_offsets = self._result_set() + + self.assertFalse(rs.has_more_rows) + ids = rs.fetchall_arrow().column(0).to_pylist() + + self.assertEqual(ids, list(range(110))) + self.assertEqual(expected_offsets, [20, 50, 80]) + backend.fetch_results.assert_not_called() + + def test_fetchmany_and_fetchone_cross_link_batches(self): + rs, _, _ = self._result_set() + + ids = [rs.fetchone()[0] for _ in range(15)] + while True: + table = rs.fetchmany_arrow(25) + if table.num_rows == 0: + break + ids += table.column(0).to_pylist() + + self.assertEqual(ids, list(range(110))) + + def test_close_mid_stream_cancels_the_download_manager(self): + rs, backend, _ = self._result_set() + manager = rs.results.download_manager + rs.fetchmany_arrow(35) + + rs.close() + + self.assertTrue(manager._cancelled.is_set()) + self.assertIsNone(manager._link_source) + backend.close_command.assert_called_once() + + class ThriftClientPrefetchTests(unittest.TestCase): def _client(self, **kwargs): return ThriftDatabricksClient( From e114c4830322a6ba518ea50d03c6010cdc4d2c84 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 12 Nov 2026 11:29:21 +0000 Subject: [PATCH 04/17] Stop scheduling if the manager is shut down during a link fetch A close() on another thread while FetchResults is in flight now ends the stream instead of submitting to a shut-down pool. Signed-off-by: Vu Anh Phung Co-authored-by: Isaac --- src/databricks/sql/cloudfetch/download_manager.py | 3 +++ tests/unit/test_cloudfetch_prefetch.py | 13 +++++++++++++ 2 files changed, 16 insertions(+) diff --git a/src/databricks/sql/cloudfetch/download_manager.py b/src/databricks/sql/cloudfetch/download_manager.py index 1696ef302..0ca491a64 100644 --- a/src/databricks/sql/cloudfetch/download_manager.py +++ b/src/databricks/sql/cloudfetch/download_manager.py @@ -125,6 +125,9 @@ def _schedule_downloads(self): ): scheduled = len(self._download_tasks) links, has_more = self._link_source(self._max_download_threads - scheduled) + if self._cancelled.is_set(): + # Shut down while the request was in flight + return if not has_more: self._link_source = None for link in links: diff --git a/tests/unit/test_cloudfetch_prefetch.py b/tests/unit/test_cloudfetch_prefetch.py index 48e481db2..7ace54cf0 100644 --- a/tests/unit/test_cloudfetch_prefetch.py +++ b/tests/unit/test_cloudfetch_prefetch.py @@ -152,6 +152,19 @@ def test_link_source_errors_reach_the_caller_and_fetching_can_resume(self): self.assertEqual(offsets, [0, 10]) self.assertEqual(source.call_count, 2) + @patch.object(download_manager, "ResultSetDownloadHandler", _InstantHandler) + def test_shutdown_during_a_link_fetch_ends_the_stream_cleanly(self): + manager = self._manager([_link(0)], max_download_threads=2) + + def source(free): + manager._shutdown_manager() # e.g. close() on another thread + return [_link(10)], True + + manager.set_link_source(source) + + self.assertIsNone(manager.get_next_downloaded_file(0)) + self.assertEqual(manager._pending_links, []) + @patch.object(download_manager, "ResultSetDownloadHandler", _InstantHandler) def test_link_source_is_refused_after_shutdown(self): # No schedulable links: the first request already shuts the manager down From 5ef9b34cf889f92954b997f2056176f23c5fb7ec Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 12 Nov 2026 11:35:13 +0000 Subject: [PATCH 05/17] Accept string values for _cloudfetch_prefetch "false", "0", "no" and "off" now turn prefetch off, so the kill switch also works when set from a URL query. Add a queue-level test that a failed link fetch propagates out of the fetch call. Signed-off-by: Vu Anh Phung Co-authored-by: Isaac --- src/databricks/sql/backend/thrift_backend.py | 5 ++++- tests/unit/test_cloudfetch_prefetch.py | 18 +++++++++++++++--- 2 files changed, 19 insertions(+), 4 deletions(-) diff --git a/src/databricks/sql/backend/thrift_backend.py b/src/databricks/sql/backend/thrift_backend.py index 5a42058cd..e03bce69a 100644 --- a/src/databricks/sql/backend/thrift_backend.py +++ b/src/databricks/sql/backend/thrift_backend.py @@ -195,7 +195,10 @@ def __init__( # Cloud fetch self._max_download_threads = kwargs.get("max_download_threads", 10) - self._cloudfetch_prefetch = bool(kwargs.get("_cloudfetch_prefetch", True)) + prefetch = kwargs.get("_cloudfetch_prefetch", True) + if isinstance(prefetch, str): # e.g. from a SQLAlchemy URL query + prefetch = prefetch.strip().lower() not in ("false", "0", "no", "off") + self._cloudfetch_prefetch = bool(prefetch) self._ssl_options = ssl_options self._auth_provider = auth_provider diff --git a/tests/unit/test_cloudfetch_prefetch.py b/tests/unit/test_cloudfetch_prefetch.py index 7ace54cf0..260f4d1a5 100644 --- a/tests/unit/test_cloudfetch_prefetch.py +++ b/tests/unit/test_cloudfetch_prefetch.py @@ -136,7 +136,7 @@ def test_reply_with_only_empty_links_is_retried_on_the_next_call(self): self.assertIsNone(manager.get_next_downloaded_file(20)) @patch.object(download_manager, "ResultSetDownloadHandler", _InstantHandler) - def test_link_source_errors_reach_the_caller_and_fetching_can_resume(self): + def test_link_source_errors_reach_the_caller_and_the_next_call_asks_again(self): manager = self._manager([_link(0)], max_download_threads=2) source = Mock( side_effect=[RuntimeError("FetchResults failed"), ([_link(10)], False)] @@ -157,7 +157,7 @@ def test_shutdown_during_a_link_fetch_ends_the_stream_cleanly(self): manager = self._manager([_link(0)], max_download_threads=2) def source(free): - manager._shutdown_manager() # e.g. close() on another thread + manager._shutdown_manager() return [_link(10)], True manager.set_link_source(source) @@ -426,6 +426,13 @@ def test_close_mid_stream_cancels_the_download_manager(self): self.assertIsNone(manager._link_source) backend.close_command.assert_called_once() + def test_link_fetch_errors_propagate_out_of_the_fetch_call(self): + rs, backend, _ = self._result_set() + backend.fetch_result_links.side_effect = DataError("FetchResults failed") + + with self.assertRaises(DataError): + rs.fetchall_arrow() + class ThriftClientPrefetchTests(unittest.TestCase): def _client(self, **kwargs): @@ -442,7 +449,12 @@ def _client(self, **kwargs): def test_prefetch_is_on_by_default_and_can_be_disabled(self): self.assertTrue(self._client()._cloudfetch_prefetch) - self.assertFalse(self._client(_cloudfetch_prefetch=False)._cloudfetch_prefetch) + for off in (False, "false", "False", "0", "no", "off"): + self.assertFalse( + self._client(_cloudfetch_prefetch=off)._cloudfetch_prefetch + ) + for on in (True, "true", "1"): + self.assertTrue(self._client(_cloudfetch_prefetch=on)._cloudfetch_prefetch) @patch("databricks.sql.backend.thrift_backend.ThriftResultSet") @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True) From 65128208dd350bc1058397e7e56bc66292b81b36 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 12 Nov 2026 12:22:46 +0000 Subject: [PATCH 06/17] Make CloudFetch cancellation stop downloads in flight Downloads now read the body in 1 MiB chunks and check the cancel event between chunks, so closing a result set aborts a GET instead of waiting for it; a body read that fails mid-way is retried (up to 3 attempts), as urllib3 did when it read the body itself. Cursor.cancel() also cancels the active result set's downloads: later reads raise OperationalError rather than returning a truncated result. Signed-off-by: Vu Anh Phung Co-authored-by: Isaac --- CHANGELOG.md | 2 +- src/databricks/sql/client.py | 4 + .../sql/cloudfetch/download_manager.py | 34 +++- src/databricks/sql/cloudfetch/downloader.py | 54 ++++-- src/databricks/sql/result_set.py | 8 + tests/unit/test_cloudfetch_prefetch.py | 166 ++++++++++++++++-- tests/unit/test_downloader.py | 1 + 7 files changed, 238 insertions(+), 31 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d6240c029..f4675550b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,7 +4,7 @@ - Fix inline parameters (`use_inline_params=True`) so a sequence directly after `IN` renders as a list again: `x IN %(ids)s` renders as `x IN (1,2,3)` instead of `x IN ARRAY(1,2,3)`, which Spark rejects. Inside the list, tuples render as rows for a multi-column `IN`, and a `set` or `frozenset` is accepted again. Sequences anywhere else still render as `ARRAY(...)` ([#629](https://github.com/databricks/databricks-sql-python/issues/629)) - Speed up Thrift response handling: read replies in 1 MiB chunks and decode them with Thrift's C-accelerated protocol, and stop stringifying every request/response for the default no-op debug logger. Wide result schemas, metadata calls with many rows, and inline Arrow results benefit most. - Speed up large CloudFetch results on Thrift: each FetchResults call now asks for about one result file per download thread (at most 4x `buffer_size_bytes`) instead of 100 MiB of files, so all download threads stay busy and the result takes fewer round trips. -- Speed up large CloudFetch results on Thrift further by requesting the next result links while earlier files are still downloading, and stop downloads promptly when a result set is closed early. +- Speed up large CloudFetch results on Thrift further by requesting the next result links while earlier files are still downloading, and stop downloads in flight when a result set is closed early or its statement is cancelled. # 4.9.0 (2026-11-04) - Upgrade Databricks SQL Kernel to 1.4.0, which keeps Lakehouse Real-Time session affinity and flushes shared telemetry only after the final session closes. diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index 613b62386..1841c8034 100755 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -1959,6 +1959,10 @@ def cancel(self) -> None: """ if self.active_command_id is not None: self.backend.cancel_command(self.active_command_id) + result_set = self.active_result_set + if result_set is not None: + # Stop downloads already under way; reads then fail like the server would + result_set._cancel_downloads() return # No command id yet. A backend whose synchronous ``execute()`` # blocks without first publishing a command id (the kernel diff --git a/src/databricks/sql/cloudfetch/download_manager.py b/src/databricks/sql/cloudfetch/download_manager.py index 0ca491a64..c7a356bec 100644 --- a/src/databricks/sql/cloudfetch/download_manager.py +++ b/src/databricks/sql/cloudfetch/download_manager.py @@ -3,9 +3,10 @@ import logging import threading -from concurrent.futures import ThreadPoolExecutor, Future +from concurrent.futures import CancelledError, ThreadPoolExecutor, Future from typing import Callable, List, Union, Tuple, Optional, TYPE_CHECKING +from databricks.sql.exc import OperationalError from databricks.sql.cloudfetch.downloader import ( ResultSetDownloadHandler, DownloadableResultSettings, @@ -63,6 +64,22 @@ def __init__( self._link_source: Optional[LinkSource] = None # Set on shutdown so downloads that have not started (or finished) stop early. self._cancelled = threading.Event() + # Set by cancel(): the result is incomplete, so reads must fail, not just end. + self._cancel_requested = False + + def cancel(self) -> None: + """Stop downloading because the statement was cancelled. + + May be called from another thread while the consumer is reading. Later reads + raise instead of reporting the end of the result. + """ + self._cancel_requested = True + self._cancelled.set() + self._thread_pool.shutdown(wait=False, cancel_futures=True) + + def _raise_if_cancel_requested(self) -> None: + if self._cancel_requested: + raise OperationalError("CloudFetch download cancelled") def set_link_source(self, link_source: LinkSource) -> bool: """Fetch further links on demand. @@ -92,8 +109,10 @@ def get_next_downloaded_file( next_row_offset (int): The offset of the starting row of the next file we want data from. """ + self._raise_if_cancel_requested() # Make sure the download queue is always full self._schedule_downloads() + self._raise_if_cancel_requested() # No more files to download from this batch of links if len(self._download_tasks) == 0: @@ -104,7 +123,10 @@ def get_next_downloaded_file( # Future's `result()` method will wait for the call to complete, and return # the value returned by the call. If the call throws an exception - `result()` # will throw the same exception - file = task.result() + try: + file = task.result() + except CancelledError: + raise OperationalError("CloudFetch download cancelled") from None if (next_row_offset < file.start_row_offset) or ( next_row_offset > file.start_row_offset + file.row_count ): @@ -162,7 +184,13 @@ def _submit_pending_links(self): http_client=self._http_client, cancel_event=self._cancelled, ) - task = self._thread_pool.submit(handler.run) + try: + task = self._thread_pool.submit(handler.run) + except RuntimeError: + # cancel() shut the pool down from another thread + if self._cancel_requested: + return + raise self._download_tasks.append(task) def add_link(self, link: TSparkArrowResultLink): diff --git a/src/databricks/sql/cloudfetch/downloader.py b/src/databricks/sql/cloudfetch/downloader.py index 9a140ea63..170033e38 100644 --- a/src/databricks/sql/cloudfetch/downloader.py +++ b/src/databricks/sql/cloudfetch/downloader.py @@ -23,8 +23,18 @@ from databricks.sql.thrift_api.TCLIService.ttypes import TSparkArrowResultLink +from urllib3.exceptions import HTTPError as Urllib3HTTPError + +from databricks.sql.exc import OperationalError + logger = logging.getLogger(__name__) +# Bodies are read in chunks so a download stops soon after its result set is closed. +_READ_CHUNK_BYTES = 1024 * 1024 +# urllib3 retries a failed body read only when it reads the body itself, so a GET +# interrupted mid-body is retried here. +_BODY_READ_ATTEMPTS = 3 + @dataclass class DownloadedFile: @@ -85,7 +95,18 @@ def __init__( def _raise_if_cancelled(self): if self._cancel_event is not None and self._cancel_event.is_set(): - raise Error("CloudFetch download cancelled: the result set was closed") + raise OperationalError("CloudFetch download cancelled") + + def _read_body(self, response) -> Optional[bytes]: + """Returns None if cancelled part-way; the unread connection is then closed.""" + chunks = [] + for chunk in response.stream(_READ_CHUNK_BYTES): + if self._cancel_event is not None and self._cancel_event.is_set(): + return None + chunks.append(chunk) + # Fully read, so the connection can be reused for the next download + response.release_conn() + return b"".join(chunks) @log_latency(StatementType.QUERY) def run(self) -> DownloadedFile: @@ -111,17 +132,28 @@ def run(self) -> DownloadedFile: self._raise_if_cancelled() start_time = time.time() - with self._http_client.request_context( - method=HttpMethod.GET, - url=self.link.fileLink, - timeout=self.settings.download_timeout, - headers=self.link.httpHeaders, - ) as response: - if response.status >= 400: - raise Exception(f"HTTP {response.status}: {response.data.decode()}") - compressed_data = response.data - # Skip decompressing a file nobody will read + for attempt in range(1, _BODY_READ_ATTEMPTS + 1): + with self._http_client.request_context( + method=HttpMethod.GET, + url=self.link.fileLink, + timeout=self.settings.download_timeout, + headers=self.link.httpHeaders, + preload_content=False, + ) as response: + if response.status >= 400: + raise Exception(f"HTTP {response.status}: {response.data.decode()}") + try: + compressed_data = self._read_body(response) + break + except Urllib3HTTPError as e: + if attempt == _BODY_READ_ATTEMPTS: + raise + logger.debug("CloudFetch download interrupted, retrying: %s", e) + # Raised outside request_context, which would log it as an HTTP error. + # Also skips decompressing a file nobody will read. self._raise_if_cancelled() + if compressed_data is None: + raise OperationalError("CloudFetch download cancelled") # Log download metrics download_duration = time.time() - start_time diff --git a/src/databricks/sql/result_set.py b/src/databricks/sql/result_set.py index 73c3c0811..f40f62b73 100644 --- a/src/databricks/sql/result_set.py +++ b/src/databricks/sql/result_set.py @@ -172,6 +172,9 @@ def fetchall_arrow(self) -> "pyarrow.Table": """Fetch all remaining rows as an Arrow table.""" pass + def _cancel_downloads(self) -> None: + """Stop client-side downloads after Cursor.cancel(); may run on another thread.""" + def close(self) -> None: """ Close the result set. @@ -334,6 +337,11 @@ def _fetch_more_links(self, free_slots: int): ) return links, has_more_rows + def _cancel_downloads(self) -> None: + queue = self.results + if isinstance(queue, ThriftCloudFetchQueue): + queue.download_manager.cancel() + def _note_link_sizes(self, links) -> None: for link in links: # The reader doesn't enforce required fields; a missing size only diff --git a/tests/unit/test_cloudfetch_prefetch.py b/tests/unit/test_cloudfetch_prefetch.py index 260f4d1a5..4bceae595 100644 --- a/tests/unit/test_cloudfetch_prefetch.py +++ b/tests/unit/test_cloudfetch_prefetch.py @@ -2,9 +2,12 @@ import threading import unittest +from concurrent.futures import Future +from contextlib import nullcontext from unittest.mock import MagicMock, Mock, patch import pytest +from urllib3.exceptions import ProtocolError try: import pyarrow as pa @@ -16,8 +19,9 @@ from databricks.sql.auth.authenticators import AuthProvider from databricks.sql.backend.thrift_backend import ThriftDatabricksClient from databricks.sql.backend.types import ExecuteResponse +from databricks.sql.client import Cursor from databricks.sql.cloudfetch.downloader import DownloadedFile -from databricks.sql.exc import DataError, Error +from databricks.sql.exc import DataError, OperationalError from databricks.sql.result_set import ThriftResultSet from databricks.sql.thrift_api.TCLIService import ttypes from databricks.sql.types import SSLOptions @@ -165,6 +169,50 @@ def source(free): self.assertIsNone(manager.get_next_downloaded_file(0)) self.assertEqual(manager._pending_links, []) + @patch.object(download_manager, "ResultSetDownloadHandler", _InstantHandler) + def test_cancel_makes_later_reads_fail_instead_of_ending(self): + manager = self._manager([_link(0), _link(10)], max_download_threads=2) + self.assertEqual(manager.get_next_downloaded_file(0).start_row_offset, 0) + + manager.cancel() + + self.assertTrue(manager._cancelled.is_set()) + with self.assertRaises(OperationalError): + manager.get_next_downloaded_file(10) + + @patch.object(download_manager, "ResultSetDownloadHandler", _InstantHandler) + def test_cancel_during_a_link_fetch_fails_the_read(self): + manager = self._manager([_link(0)], max_download_threads=2) + + def source(free): + manager.cancel() # Cursor.cancel() running on another thread + return [_link(10)], True + + manager.set_link_source(source) + + with self.assertRaises(OperationalError): + manager.get_next_downloaded_file(0) + + @patch.object(download_manager, "ResultSetDownloadHandler", _InstantHandler) + def test_scheduling_after_cancel_does_not_raise_runtime_error(self): + manager = self._manager([_link(0), _link(10)], max_download_threads=2) + manager.cancel() + + manager._submit_pending_links() + + self.assertEqual(manager._download_tasks, []) + with self.assertRaises(OperationalError): + manager.get_next_downloaded_file(0) + + def test_download_cancelled_before_it_started_fails_the_read(self): + manager = self._manager([], max_download_threads=1) + future = Future() + future.cancel() + manager._download_tasks.append(future) + + with self.assertRaises(OperationalError): + manager.get_next_downloaded_file(0) + @patch.object(download_manager, "ResultSetDownloadHandler", _InstantHandler) def test_link_source_is_refused_after_shutdown(self): # No schedulable links: the first request already shuts the manager down @@ -183,9 +231,39 @@ def test_handlers_share_the_managers_cancel_event(self): class DownloaderCancelTests(unittest.TestCase): - def _handler(self, http_client, event): + class _Response: + """Streams the given chunks; a chunk that is an exception is raised instead.""" + + status = 200 + + def __init__(self, chunks, before_chunk=None): + self.chunks = chunks + self.before_chunk = before_chunk or (lambda i: None) + self.read = 0 + self.released = False + + def stream(self, amt): + for i, chunk in enumerate(self.chunks): + self.before_chunk(i) + if isinstance(chunk, Exception): + raise chunk + self.read += 1 + yield chunk + + def release_conn(self): + self.released = True + + def _http_client(self, *responses): + http_client = MagicMock() + pending = list(responses) + http_client.request_context.side_effect = lambda **_: nullcontext( + pending.pop(0) + ) + return http_client + + def _handler(self, http_client, event, lz4_compressed=True): return downloader.ResultSetDownloadHandler( - downloader.DownloadableResultSettings(is_lz4_compressed=True), + downloader.DownloadableResultSettings(is_lz4_compressed=lz4_compressed), _link(0), ssl_options=SSLOptions(), chunk_id=0, @@ -200,30 +278,55 @@ def test_cancelled_before_the_request_skips_the_get(self): event = threading.Event() event.set() - with self.assertRaises(Error): + with self.assertRaises(OperationalError): self._handler(http_client, event).run() http_client.request_context.assert_not_called() - def test_cancelled_during_the_download_skips_decompression(self): - event = threading.Event() + def test_body_is_read_in_chunks_and_the_connection_released(self): + response = self._Response([b"ab", b"cd"]) + http_client = self._http_client(response) - class Response: - status = 200 + file = self._handler(http_client, threading.Event(), False).run() - @property - def data(self): - event.set() - return b"compressed" + self.assertEqual(file.file_bytes, b"abcd") + self.assertTrue(response.released) + self.assertIs( + http_client.request_context.call_args.kwargs["preload_content"], False + ) - http_client = MagicMock() - http_client.request_context.return_value.__enter__.return_value = Response() + def test_cancel_mid_body_stops_reading_and_skips_decompression(self): + event = threading.Event() + response = self._Response( + [b"a", b"b", b"c", b"d"], before_chunk=lambda i: i == 1 and event.set() + ) with patch.object( downloader.ResultSetDownloadHandler, "_decompress_data" ) as decompress: - with self.assertRaises(Error): - self._handler(http_client, event).run() + with self.assertRaises(OperationalError): + self._handler(self._http_client(response), event).run() + + self.assertEqual(response.read, 2) + self.assertFalse(response.released) decompress.assert_not_called() + def test_interrupted_body_is_retried(self): + broken = self._Response([b"ab", ProtocolError("connection reset")]) + http_client = self._http_client(broken, self._Response([b"ab", b"cd"])) + + file = self._handler(http_client, threading.Event(), False).run() + + self.assertEqual(file.file_bytes, b"abcd") + self.assertEqual(http_client.request_context.call_count, 2) + + def test_body_errors_give_up_after_three_attempts(self): + http_client = self._http_client( + *[self._Response([ProtocolError("connection reset")]) for _ in range(3)] + ) + + with self.assertRaises(ProtocolError): + self._handler(http_client, threading.Event(), False).run() + self.assertEqual(http_client.request_context.call_count, 3) + class ResultSetPrefetchTests(unittest.TestCase): BUFFER = 100 * MIB @@ -426,6 +529,15 @@ def test_close_mid_stream_cancels_the_download_manager(self): self.assertIsNone(manager._link_source) backend.close_command.assert_called_once() + def test_cancelled_downloads_fail_the_next_fetch_instead_of_truncating(self): + rs, _, _ = self._result_set() + rs.fetchmany_arrow(15) + + rs._cancel_downloads() + + with self.assertRaises(OperationalError): + rs.fetchall_arrow() + def test_link_fetch_errors_propagate_out_of_the_fetch_call(self): rs, backend, _ = self._result_set() backend.fetch_result_links.side_effect = DataError("FetchResults failed") @@ -533,5 +645,27 @@ def test_fetch_result_links_rejects_rows_that_are_not_links( ) +class CursorCancelTests(unittest.TestCase): + def test_cancel_stops_the_active_result_sets_downloads(self): + backend = Mock() + cursor = Cursor(Mock(), backend) + cursor.active_command_id = Mock() + cursor.active_result_set = Mock() + + cursor.cancel() + + backend.cancel_command.assert_called_once_with(cursor.active_command_id) + cursor.active_result_set._cancel_downloads.assert_called_once_with() + + def test_thrift_result_set_cancels_its_cloudfetch_downloads(self): + queue = Mock(spec=ThriftCloudFetchQueue) + queue.download_manager = Mock() + rs = Mock(spec=ThriftResultSet, results=queue) + + ThriftResultSet._cancel_downloads(rs) + + queue.download_manager.cancel.assert_called_once_with() + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/test_downloader.py b/tests/unit/test_downloader.py index 00b1b849a..2d9d174e3 100644 --- a/tests/unit/test_downloader.py +++ b/tests/unit/test_downloader.py @@ -26,6 +26,7 @@ def _setup_mock_http_response(self, mock_http_client, status=200, data=b""): mock_response = MagicMock() mock_response.status = status mock_response.data = data + mock_response.stream.side_effect = lambda *args, **kwargs: iter([data]) mock_context_manager = MagicMock() mock_context_manager.__enter__.return_value = mock_response mock_context_manager.__exit__.return_value = None From a86a60930c006f16291e40461526884837e1759f Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 12 Nov 2026 13:06:22 +0000 Subject: [PATCH 07/17] Stop downloads even if the cancel request fails Cursor.cancel() now cancels the active result set's downloads in a finally block. The downloader tests use a stand-in that wraps errors like the real request_context, so they check the exception types callers actually see. Signed-off-by: Vu Anh Phung Co-authored-by: Isaac --- src/databricks/sql/client.py | 13 ++++++---- tests/unit/test_cloudfetch_prefetch.py | 34 ++++++++++++++++++++++---- 2 files changed, 37 insertions(+), 10 deletions(-) diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index 1841c8034..d3b8b6530 100755 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -1958,11 +1958,14 @@ def cancel(self) -> None: This method can be called from another thread. """ if self.active_command_id is not None: - self.backend.cancel_command(self.active_command_id) - result_set = self.active_result_set - if result_set is not None: - # Stop downloads already under way; reads then fail like the server would - result_set._cancel_downloads() + try: + self.backend.cancel_command(self.active_command_id) + finally: + result_set = self.active_result_set + if result_set is not None: + # Stop downloads already under way; reads then fail like the server + # would. Done even if the cancel RPC fails. + result_set._cancel_downloads() return # No command id yet. A backend whose synchronous ``execute()`` # blocks without first publishing a command id (the kernel diff --git a/tests/unit/test_cloudfetch_prefetch.py b/tests/unit/test_cloudfetch_prefetch.py index 4bceae595..5a7f79898 100644 --- a/tests/unit/test_cloudfetch_prefetch.py +++ b/tests/unit/test_cloudfetch_prefetch.py @@ -3,7 +3,7 @@ import threading import unittest from concurrent.futures import Future -from contextlib import nullcontext +from contextlib import contextmanager from unittest.mock import MagicMock, Mock, patch import pytest @@ -21,7 +21,7 @@ from databricks.sql.backend.types import ExecuteResponse from databricks.sql.client import Cursor from databricks.sql.cloudfetch.downloader import DownloadedFile -from databricks.sql.exc import DataError, OperationalError +from databricks.sql.exc import DataError, OperationalError, RequestError from databricks.sql.result_set import ThriftResultSet from databricks.sql.thrift_api.TCLIService import ttypes from databricks.sql.types import SSLOptions @@ -253,10 +253,19 @@ def stream(self, amt): def release_conn(self): self.released = True + @staticmethod + @contextmanager + def _request_context(response): + # Like UnifiedHttpClient.request_context: errors in the block become RequestError + try: + yield response + except Exception as e: + raise RequestError(f"HTTP request error: {e}") + def _http_client(self, *responses): http_client = MagicMock() pending = list(responses) - http_client.request_context.side_effect = lambda **_: nullcontext( + http_client.request_context.side_effect = lambda **_: self._request_context( pending.pop(0) ) return http_client @@ -302,9 +311,11 @@ def test_cancel_mid_body_stops_reading_and_skips_decompression(self): with patch.object( downloader.ResultSetDownloadHandler, "_decompress_data" ) as decompress: - with self.assertRaises(OperationalError): + with self.assertRaises(OperationalError) as raised: self._handler(self._http_client(response), event).run() + # Raised outside the HTTP client, so it is not wrapped (or logged) as an HTTP error + self.assertNotIsInstance(raised.exception, RequestError) self.assertEqual(response.read, 2) self.assertFalse(response.released) decompress.assert_not_called() @@ -323,7 +334,9 @@ def test_body_errors_give_up_after_three_attempts(self): *[self._Response([ProtocolError("connection reset")]) for _ in range(3)] ) - with self.assertRaises(ProtocolError): + # The last failure surfaces as the client's RequestError, as on main when + # urllib3 runs out of retries + with self.assertRaises(RequestError): self._handler(http_client, threading.Event(), False).run() self.assertEqual(http_client.request_context.call_count, 3) @@ -657,6 +670,17 @@ def test_cancel_stops_the_active_result_sets_downloads(self): backend.cancel_command.assert_called_once_with(cursor.active_command_id) cursor.active_result_set._cancel_downloads.assert_called_once_with() + def test_downloads_stop_even_if_the_cancel_request_fails(self): + backend = Mock() + backend.cancel_command.side_effect = RequestError("session gone") + cursor = Cursor(Mock(), backend) + cursor.active_command_id = Mock() + cursor.active_result_set = Mock() + + with self.assertRaises(RequestError): + cursor.cancel() + cursor.active_result_set._cancel_downloads.assert_called_once_with() + def test_thrift_result_set_cancels_its_cloudfetch_downloads(self): queue = Mock(spec=ThriftCloudFetchQueue) queue.download_manager = Mock() From e1ce85786b7dca67073433f812a5d515fc6990f3 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 12 Nov 2026 13:18:03 +0000 Subject: [PATCH 08/17] After cancel, still return files that finished downloading A read after cancel() now fails only when the cancel left the result incomplete (links pending, more on the server, or a download it stopped). Cancelling a statement whose result had already arrived no longer turns the end of the result into an error. A link whose submit races with cancel() is kept, so it still counts as missing. Signed-off-by: Vu Anh Phung Co-authored-by: Isaac --- .../sql/cloudfetch/download_manager.py | 27 +++++++++-------- tests/unit/test_cloudfetch_prefetch.py | 30 +++++++++++++++---- 2 files changed, 39 insertions(+), 18 deletions(-) diff --git a/src/databricks/sql/cloudfetch/download_manager.py b/src/databricks/sql/cloudfetch/download_manager.py index c7a356bec..5482f241f 100644 --- a/src/databricks/sql/cloudfetch/download_manager.py +++ b/src/databricks/sql/cloudfetch/download_manager.py @@ -64,23 +64,20 @@ def __init__( self._link_source: Optional[LinkSource] = None # Set on shutdown so downloads that have not started (or finished) stop early. self._cancelled = threading.Event() - # Set by cancel(): the result is incomplete, so reads must fail, not just end. + # Set by cancel(): reads fail once they need data the cancel stopped. self._cancel_requested = False def cancel(self) -> None: """Stop downloading because the statement was cancelled. - May be called from another thread while the consumer is reading. Later reads - raise instead of reporting the end of the result. + May be called from another thread while the consumer is reading. Files that + finished downloading are still returned; a read that needs anything else + raises instead of reporting the end of the result. """ self._cancel_requested = True self._cancelled.set() self._thread_pool.shutdown(wait=False, cancel_futures=True) - def _raise_if_cancel_requested(self) -> None: - if self._cancel_requested: - raise OperationalError("CloudFetch download cancelled") - def set_link_source(self, link_source: LinkSource) -> bool: """Fetch further links on demand. @@ -109,13 +106,17 @@ def get_next_downloaded_file( next_row_offset (int): The offset of the starting row of the next file we want data from. """ - self._raise_if_cancel_requested() - # Make sure the download queue is always full - self._schedule_downloads() - self._raise_if_cancel_requested() + if not self._cancel_requested: + # Make sure the download queue is always full + self._schedule_downloads() # No more files to download from this batch of links if len(self._download_tasks) == 0: + if self._cancel_requested and ( + self._pending_links or self._link_source is not None + ): + # The cancel stopped part of the result: fail rather than end early + raise OperationalError("CloudFetch download cancelled") self._shutdown_manager() return None @@ -187,8 +188,10 @@ def _submit_pending_links(self): try: task = self._thread_pool.submit(handler.run) except RuntimeError: - # cancel() shut the pool down from another thread + # cancel() shut the pool down from another thread; keep the link so + # the result still counts as incomplete if self._cancel_requested: + self._pending_links.insert(0, (chunk_id, link)) return raise self._download_tasks.append(task) diff --git a/tests/unit/test_cloudfetch_prefetch.py b/tests/unit/test_cloudfetch_prefetch.py index 5a7f79898..fba07e952 100644 --- a/tests/unit/test_cloudfetch_prefetch.py +++ b/tests/unit/test_cloudfetch_prefetch.py @@ -170,37 +170,55 @@ def source(free): self.assertEqual(manager._pending_links, []) @patch.object(download_manager, "ResultSetDownloadHandler", _InstantHandler) - def test_cancel_makes_later_reads_fail_instead_of_ending(self): + def test_cancel_after_the_whole_result_arrived_still_returns_it(self): manager = self._manager([_link(0), _link(10)], max_download_threads=2) self.assertEqual(manager.get_next_downloaded_file(0).start_row_offset, 0) + manager._download_tasks[0].result() # the last file finished downloading manager.cancel() self.assertTrue(manager._cancelled.is_set()) + self.assertEqual(manager.get_next_downloaded_file(10).start_row_offset, 10) + self.assertIsNone(manager.get_next_downloaded_file(20)) + + @patch.object(download_manager, "ResultSetDownloadHandler", _InstantHandler) + def test_cancel_with_links_not_yet_downloaded_fails_the_read(self): + manager = self._manager( + [_link(0), _link(10), _link(20)], max_download_threads=1 + ) + self.assertEqual(manager.get_next_downloaded_file(0).start_row_offset, 0) + + manager.cancel() + with self.assertRaises(OperationalError): manager.get_next_downloaded_file(10) @patch.object(download_manager, "ResultSetDownloadHandler", _InstantHandler) - def test_cancel_during_a_link_fetch_fails_the_read(self): + def test_cancel_during_a_link_fetch_fails_once_the_downloaded_files_run_out( + self, + ): manager = self._manager([_link(0)], max_download_threads=2) def source(free): + manager._download_tasks[0].result() # the first file is downloaded manager.cancel() # Cursor.cancel() running on another thread return [_link(10)], True manager.set_link_source(source) + self.assertEqual(manager.get_next_downloaded_file(0).start_row_offset, 0) with self.assertRaises(OperationalError): - manager.get_next_downloaded_file(0) + manager.get_next_downloaded_file(10) @patch.object(download_manager, "ResultSetDownloadHandler", _InstantHandler) - def test_scheduling_after_cancel_does_not_raise_runtime_error(self): - manager = self._manager([_link(0), _link(10)], max_download_threads=2) + def test_link_whose_submit_raced_with_cancel_still_counts_as_missing(self): + manager = self._manager([_link(0)], max_download_threads=2) manager.cancel() - manager._submit_pending_links() + manager._submit_pending_links() # submit fails: the pool is shut down self.assertEqual(manager._download_tasks, []) + self.assertEqual(len(manager._pending_links), 1) with self.assertRaises(OperationalError): manager.get_next_downloaded_file(0) From 30aaac39b8da3f3518d872bc489198412b85f89f Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 12 Nov 2026 13:27:59 +0000 Subject: [PATCH 09/17] Re-check cancellation and link expiry before a retried GET Signed-off-by: Vu Anh Phung Co-authored-by: Isaac --- src/databricks/sql/cloudfetch/downloader.py | 6 ++++++ tests/unit/test_cloudfetch_prefetch.py | 11 +++++++++++ 2 files changed, 17 insertions(+) diff --git a/src/databricks/sql/cloudfetch/downloader.py b/src/databricks/sql/cloudfetch/downloader.py index 170033e38..b444d8349 100644 --- a/src/databricks/sql/cloudfetch/downloader.py +++ b/src/databricks/sql/cloudfetch/downloader.py @@ -133,6 +133,12 @@ def run(self) -> DownloadedFile: start_time = time.time() for attempt in range(1, _BODY_READ_ATTEMPTS + 1): + if attempt > 1: + # Re-check before the retry's GET + ResultSetDownloadHandler._validate_link( + self.link, self.settings.link_expiry_buffer_secs + ) + self._raise_if_cancelled() with self._http_client.request_context( method=HttpMethod.GET, url=self.link.fileLink, diff --git a/tests/unit/test_cloudfetch_prefetch.py b/tests/unit/test_cloudfetch_prefetch.py index fba07e952..667037775 100644 --- a/tests/unit/test_cloudfetch_prefetch.py +++ b/tests/unit/test_cloudfetch_prefetch.py @@ -347,6 +347,17 @@ def test_interrupted_body_is_retried(self): self.assertEqual(file.file_bytes, b"abcd") self.assertEqual(http_client.request_context.call_count, 2) + def test_no_retry_after_a_cancel(self): + event = threading.Event() + broken = self._Response( + [ProtocolError("connection reset")], before_chunk=lambda i: event.set() + ) + http_client = self._http_client(broken, self._Response([b"ab"])) + + with self.assertRaises(OperationalError): + self._handler(http_client, event, False).run() + self.assertEqual(http_client.request_context.call_count, 1) + def test_body_errors_give_up_after_three_attempts(self): http_client = self._http_client( *[self._Response([ProtocolError("connection reset")]) for _ in range(3)] From a2643e85804467b33e72a599b0853429137f8d78 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 12 Nov 2026 13:40:24 +0000 Subject: [PATCH 10/17] Make a failed mid-stream link fetch fail every later read A link fetch that raises part-way through a read loses the rows that read had collected, so a retry would continue past them. The download manager now raises the same error on every later read. Test the kill switch at all six ThriftResultSet entry points. Signed-off-by: Vu Anh Phung Co-authored-by: Isaac --- .../sql/cloudfetch/download_manager.py | 14 +++- tests/unit/test_cloudfetch_prefetch.py | 67 +++++++++++++------ 2 files changed, 60 insertions(+), 21 deletions(-) diff --git a/src/databricks/sql/cloudfetch/download_manager.py b/src/databricks/sql/cloudfetch/download_manager.py index 5482f241f..bca369a8a 100644 --- a/src/databricks/sql/cloudfetch/download_manager.py +++ b/src/databricks/sql/cloudfetch/download_manager.py @@ -66,6 +66,8 @@ def __init__( self._cancelled = threading.Event() # Set by cancel(): reads fail once they need data the cancel stopped. self._cancel_requested = False + # A failed link fetch, raised again by every later read (see _schedule_downloads) + self._link_error: Optional[BaseException] = None def cancel(self) -> None: """Stop downloading because the statement was cancelled. @@ -106,6 +108,8 @@ def get_next_downloaded_file( next_row_offset (int): The offset of the starting row of the next file we want data from. """ + if self._link_error is not None: + raise self._link_error if not self._cancel_requested: # Make sure the download queue is always full self._schedule_downloads() @@ -147,7 +151,15 @@ def _schedule_downloads(self): and len(self._download_tasks) < self._max_download_threads ): scheduled = len(self._download_tasks) - links, has_more = self._link_source(self._max_download_threads - scheduled) + try: + links, has_more = self._link_source( + self._max_download_threads - scheduled + ) + except Exception as e: + # This read is failing part-way, so the caller loses the rows it had + # collected. Fail every later read too, so a retry cannot skip them. + self._link_error = e + raise if self._cancelled.is_set(): # Shut down while the request was in flight return diff --git a/tests/unit/test_cloudfetch_prefetch.py b/tests/unit/test_cloudfetch_prefetch.py index 667037775..5160f540f 100644 --- a/tests/unit/test_cloudfetch_prefetch.py +++ b/tests/unit/test_cloudfetch_prefetch.py @@ -140,21 +140,17 @@ def test_reply_with_only_empty_links_is_retried_on_the_next_call(self): self.assertIsNone(manager.get_next_downloaded_file(20)) @patch.object(download_manager, "ResultSetDownloadHandler", _InstantHandler) - def test_link_source_errors_reach_the_caller_and_the_next_call_asks_again(self): + def test_a_failed_link_fetch_fails_every_later_read(self): manager = self._manager([_link(0)], max_download_threads=2) - source = Mock( - side_effect=[RuntimeError("FetchResults failed"), ([_link(10)], False)] - ) + error = RuntimeError("FetchResults failed") + source = Mock(side_effect=[error, ([_link(10)], False)]) manager.set_link_source(source) - with self.assertRaises(RuntimeError): - manager.get_next_downloaded_file(0) - # The scheduled download is kept and the next call asks for links again - offsets = [manager.get_next_downloaded_file(0).start_row_offset] - offsets.append(manager.get_next_downloaded_file(10).start_row_offset) - self.assertIsNone(manager.get_next_downloaded_file(20)) - self.assertEqual(offsets, [0, 10]) - self.assertEqual(source.call_count, 2) + for _ in range(2): + with self.assertRaises(RuntimeError) as raised: + manager.get_next_downloaded_file(0) + self.assertIs(raised.exception, error) + self.assertEqual(source.call_count, 1) @patch.object(download_manager, "ResultSetDownloadHandler", _InstantHandler) def test_shutdown_during_a_link_fetch_ends_the_stream_cleanly(self): @@ -580,12 +576,19 @@ def test_cancelled_downloads_fail_the_next_fetch_instead_of_truncating(self): with self.assertRaises(OperationalError): rs.fetchall_arrow() - def test_link_fetch_errors_propagate_out_of_the_fetch_call(self): + def test_a_retry_after_a_failed_link_fetch_fails_instead_of_skipping_rows(self): rs, backend, _ = self._result_set() - backend.fetch_result_links.side_effect = DataError("FetchResults failed") + backend.fetch_result_links.side_effect = [ + DataError("FetchResults failed"), + ([_link(20, self.ROWS_PER_LINK)], False), + ] with self.assertRaises(DataError): - rs.fetchall_arrow() + rs.fetchmany_arrow(25) + # The rows collected by the failed call are gone, so retrying must fail too + with self.assertRaises(DataError): + rs.fetchmany_arrow(25) + self.assertEqual(backend.fetch_result_links.call_count, 1) class ThriftClientPrefetchTests(unittest.TestCase): @@ -610,15 +613,39 @@ def test_prefetch_is_on_by_default_and_can_be_disabled(self): for on in (True, "true", "1"): self.assertTrue(self._client(_cloudfetch_prefetch=on)._cloudfetch_prefetch) + # get_schemas/get_tables/get_columns import ThriftResultSet locally + @patch("databricks.sql.result_set.ThriftResultSet") @patch("databricks.sql.backend.thrift_backend.ThriftResultSet") @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True) - def test_execute_passes_the_setting_to_the_result_set(self, _, result_set): + def test_every_entry_point_passes_the_setting_to_the_result_set( + self, _, backend_result_set, module_result_set + ): client = self._client(_cloudfetch_prefetch=False) client._handle_execute_response = Mock(return_value=(Mock(), Mock())) - - client.execute_command("foo", Mock(), 100, 200, Mock(), Mock(), Mock()) - - self.assertIs(result_set.call_args.kwargs["cloudfetch_prefetch"], False) + client._hive_schema_to_description = Mock(return_value=[]) + entry_points = { + "execute_command": lambda: client.execute_command( + "foo", Mock(), 100, 200, Mock(), Mock(), Mock() + ), + "get_execution_result": lambda: client.get_execution_result(Mock(), Mock()), + "get_catalogs": lambda: client.get_catalogs(Mock(), 100, 200, Mock()), + "get_schemas": lambda: client.get_schemas(Mock(), 100, 200, Mock()), + "get_tables": lambda: client.get_tables(Mock(), 100, 200, Mock()), + "get_columns": lambda: client.get_columns(Mock(), 100, 200, Mock()), + } + for name, call in entry_points.items(): + with self.subTest(name): + backend_result_set.reset_mock() + module_result_set.reset_mock() + + call() + + result_set = ( + backend_result_set + if backend_result_set.called + else module_result_set + ) + self.assertIs(result_set.call_args.kwargs["cloudfetch_prefetch"], False) @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True) def test_fetch_result_links(self, tcli_service_class): From 181cb825e58a810384c8121d1416b138f3dbdfd8 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 12 Nov 2026 13:51:35 +0000 Subject: [PATCH 11/17] Back off between body-read retries; chain the stored link error Interrupted body reads wait 0.5 s, then 1 s, before retrying; a cancel ends the wait. After a failed mid-stream link fetch, later reads raise a fresh OperationalError chained to it that says to re-execute. A failure while cancelling downloads no longer hides the cancel RPC's own error. Signed-off-by: Vu Anh Phung Co-authored-by: Isaac --- src/databricks/sql/client.py | 7 ++-- .../sql/cloudfetch/download_manager.py | 5 ++- src/databricks/sql/cloudfetch/downloader.py | 10 ++++-- tests/unit/test_cloudfetch_prefetch.py | 35 ++++++++++++++++--- 4 files changed, 47 insertions(+), 10 deletions(-) diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index d3b8b6530..e49e35ba6 100755 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -1964,8 +1964,11 @@ def cancel(self) -> None: result_set = self.active_result_set if result_set is not None: # Stop downloads already under way; reads then fail like the server - # would. Done even if the cancel RPC fails. - result_set._cancel_downloads() + # would. Done even if the cancel RPC fails, without hiding its error. + try: + result_set._cancel_downloads() + except Exception as e: + logger.debug("Failed to cancel CloudFetch downloads: %s", e) return # No command id yet. A backend whose synchronous ``execute()`` # blocks without first publishing a command id (the kernel diff --git a/src/databricks/sql/cloudfetch/download_manager.py b/src/databricks/sql/cloudfetch/download_manager.py index bca369a8a..7de639a36 100644 --- a/src/databricks/sql/cloudfetch/download_manager.py +++ b/src/databricks/sql/cloudfetch/download_manager.py @@ -109,7 +109,10 @@ def get_next_downloaded_file( """ if self._link_error is not None: - raise self._link_error + raise OperationalError( + "Fetching more CloudFetch results failed earlier; the result can't " + "be read further, execute the statement again" + ) from self._link_error if not self._cancel_requested: # Make sure the download queue is always full self._schedule_downloads() diff --git a/src/databricks/sql/cloudfetch/downloader.py b/src/databricks/sql/cloudfetch/downloader.py index b444d8349..dbaca451f 100644 --- a/src/databricks/sql/cloudfetch/downloader.py +++ b/src/databricks/sql/cloudfetch/downloader.py @@ -32,8 +32,9 @@ # Bodies are read in chunks so a download stops soon after its result set is closed. _READ_CHUNK_BYTES = 1024 * 1024 # urllib3 retries a failed body read only when it reads the body itself, so a GET -# interrupted mid-body is retried here. +# interrupted mid-body is retried here, after a backoff that doubles each time. _BODY_READ_ATTEMPTS = 3 +_BODY_RETRY_DELAY_SECS = 0.5 @dataclass @@ -134,7 +135,12 @@ def run(self) -> DownloadedFile: for attempt in range(1, _BODY_READ_ATTEMPTS + 1): if attempt > 1: - # Re-check before the retry's GET + # Back off (a cancel ends the wait early), then re-check before the GET + delay = _BODY_RETRY_DELAY_SECS * 2 ** (attempt - 2) + if self._cancel_event is not None: + self._cancel_event.wait(delay) + else: + time.sleep(delay) ResultSetDownloadHandler._validate_link( self.link, self.settings.link_expiry_buffer_secs ) diff --git a/tests/unit/test_cloudfetch_prefetch.py b/tests/unit/test_cloudfetch_prefetch.py index 5160f540f..b5aa660f9 100644 --- a/tests/unit/test_cloudfetch_prefetch.py +++ b/tests/unit/test_cloudfetch_prefetch.py @@ -146,10 +146,14 @@ def test_a_failed_link_fetch_fails_every_later_read(self): source = Mock(side_effect=[error, ([_link(10)], False)]) manager.set_link_source(source) + with self.assertRaises(RuntimeError) as raised: + manager.get_next_downloaded_file(0) + self.assertIs(raised.exception, error) + # Later reads raise a fresh error that points at the original one for _ in range(2): - with self.assertRaises(RuntimeError) as raised: + with self.assertRaises(OperationalError) as raised: manager.get_next_downloaded_file(0) - self.assertIs(raised.exception, error) + self.assertIs(raised.exception.__cause__, error) self.assertEqual(source.call_count, 1) @patch.object(download_manager, "ResultSetDownloadHandler", _InstantHandler) @@ -334,14 +338,21 @@ def test_cancel_mid_body_stops_reading_and_skips_decompression(self): self.assertFalse(response.released) decompress.assert_not_called() + @staticmethod + def _unset_event(): + """A cancel event that is never set and records the backoff waits.""" + return Mock(is_set=Mock(return_value=False)) + def test_interrupted_body_is_retried(self): broken = self._Response([b"ab", ProtocolError("connection reset")]) http_client = self._http_client(broken, self._Response([b"ab", b"cd"])) + event = self._unset_event() - file = self._handler(http_client, threading.Event(), False).run() + file = self._handler(http_client, event, False).run() self.assertEqual(file.file_bytes, b"abcd") self.assertEqual(http_client.request_context.call_count, 2) + event.wait.assert_called_once_with(0.5) def test_no_retry_after_a_cancel(self): event = threading.Event() @@ -358,12 +369,14 @@ def test_body_errors_give_up_after_three_attempts(self): http_client = self._http_client( *[self._Response([ProtocolError("connection reset")]) for _ in range(3)] ) + event = self._unset_event() # The last failure surfaces as the client's RequestError, as on main when # urllib3 runs out of retries with self.assertRaises(RequestError): - self._handler(http_client, threading.Event(), False).run() + self._handler(http_client, event, False).run() self.assertEqual(http_client.request_context.call_count, 3) + self.assertEqual([c.args[0] for c in event.wait.call_args_list], [0.5, 1.0]) class ResultSetPrefetchTests(unittest.TestCase): @@ -586,8 +599,9 @@ def test_a_retry_after_a_failed_link_fetch_fails_instead_of_skipping_rows(self): with self.assertRaises(DataError): rs.fetchmany_arrow(25) # The rows collected by the failed call are gone, so retrying must fail too - with self.assertRaises(DataError): + with self.assertRaises(OperationalError) as raised: rs.fetchmany_arrow(25) + self.assertIsInstance(raised.exception.__cause__, DataError) self.assertEqual(backend.fetch_result_links.call_count, 1) @@ -737,6 +751,17 @@ def test_downloads_stop_even_if_the_cancel_request_fails(self): cursor.cancel() cursor.active_result_set._cancel_downloads.assert_called_once_with() + def test_a_failing_download_cancel_does_not_hide_the_cancel_error(self): + backend = Mock() + backend.cancel_command.side_effect = RequestError("session gone") + cursor = Cursor(Mock(), backend) + cursor.active_command_id = Mock() + cursor.active_result_set = Mock() + cursor.active_result_set._cancel_downloads.side_effect = RuntimeError("boom") + + with self.assertRaises(RequestError): + cursor.cancel() + def test_thrift_result_set_cancels_its_cloudfetch_downloads(self): queue = Mock(spec=ThriftCloudFetchQueue) queue.download_manager = Mock() From 8d864348c769d74423efde7a081c693ad33c1685 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 12 Nov 2026 14:00:22 +0000 Subject: [PATCH 12/17] Treat _cloudfetch_prefetch=None as unset Document what has_more_rows means after the handoff to the download manager. Signed-off-by: Vu Anh Phung Co-authored-by: Isaac --- src/databricks/sql/backend/thrift_backend.py | 6 ++++-- src/databricks/sql/result_set.py | 4 +++- tests/unit/test_cloudfetch_prefetch.py | 2 +- 3 files changed, 8 insertions(+), 4 deletions(-) diff --git a/src/databricks/sql/backend/thrift_backend.py b/src/databricks/sql/backend/thrift_backend.py index e03bce69a..401fb1ffa 100644 --- a/src/databricks/sql/backend/thrift_backend.py +++ b/src/databricks/sql/backend/thrift_backend.py @@ -195,8 +195,10 @@ def __init__( # Cloud fetch self._max_download_threads = kwargs.get("max_download_threads", 10) - prefetch = kwargs.get("_cloudfetch_prefetch", True) - if isinstance(prefetch, str): # e.g. from a SQLAlchemy URL query + prefetch = kwargs.get("_cloudfetch_prefetch") + if prefetch is None: # unset, or forwarded as None by a wrapper + prefetch = True + elif isinstance(prefetch, str): # e.g. from a SQLAlchemy URL query prefetch = prefetch.strip().lower() not in ("false", "0", "no", "off") self._cloudfetch_prefetch = bool(prefetch) diff --git a/src/databricks/sql/result_set.py b/src/databricks/sql/result_set.py index f40f62b73..e4a0c3b92 100644 --- a/src/databricks/sql/result_set.py +++ b/src/databricks/sql/result_set.py @@ -233,7 +233,9 @@ def __init__( :param t_row_set: The TRowSet containing result data (if available) :param max_download_threads: Maximum number of download threads for cloud fetch :param ssl_options: SSL options for cloud fetch - :param has_more_rows: Whether there are more rows to fetch + :param has_more_rows: Whether this result set still has to fetch more rows + itself. Set to False once CloudFetch link fetching moves to the + download manager, which then fetches the rest. :param cloudfetch_prefetch: Request further CloudFetch links while earlier files still download, instead of one batch at a time """ diff --git a/tests/unit/test_cloudfetch_prefetch.py b/tests/unit/test_cloudfetch_prefetch.py index b5aa660f9..534cf7e72 100644 --- a/tests/unit/test_cloudfetch_prefetch.py +++ b/tests/unit/test_cloudfetch_prefetch.py @@ -624,7 +624,7 @@ def test_prefetch_is_on_by_default_and_can_be_disabled(self): self.assertFalse( self._client(_cloudfetch_prefetch=off)._cloudfetch_prefetch ) - for on in (True, "true", "1"): + for on in (True, "true", "1", None): self.assertTrue(self._client(_cloudfetch_prefetch=on)._cloudfetch_prefetch) # get_schemas/get_tables/get_columns import ThriftResultSet locally From 779182a435d2248d2a09624796c2dfa906091653 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 12 Nov 2026 14:11:17 +0000 Subject: [PATCH 13/17] Warn on unrecognized _cloudfetch_prefetch strings; time only the last attempt Unrecognized string values keep prefetch on but log a warning that names the value. Download speed metrics now cover only the attempt that succeeded, not earlier failed attempts or the backoff. Signed-off-by: Vu Anh Phung Co-authored-by: Isaac --- src/databricks/sql/backend/thrift_backend.py | 13 ++++++++++++- src/databricks/sql/cloudfetch/downloader.py | 2 ++ tests/unit/test_cloudfetch_prefetch.py | 11 ++++++++++- 3 files changed, 24 insertions(+), 2 deletions(-) diff --git a/src/databricks/sql/backend/thrift_backend.py b/src/databricks/sql/backend/thrift_backend.py index 401fb1ffa..a6f913274 100644 --- a/src/databricks/sql/backend/thrift_backend.py +++ b/src/databricks/sql/backend/thrift_backend.py @@ -199,7 +199,18 @@ def __init__( if prefetch is None: # unset, or forwarded as None by a wrapper prefetch = True elif isinstance(prefetch, str): # e.g. from a SQLAlchemy URL query - prefetch = prefetch.strip().lower() not in ("false", "0", "no", "off") + value = prefetch.strip().lower() + if value in ("false", "0", "no", "off"): + prefetch = False + elif value in ("true", "1", "yes", "on"): + prefetch = True + else: + logger.warning( + "Unrecognized _cloudfetch_prefetch=%r; keeping CloudFetch " + "prefetch on (use 'false' to turn it off)", + prefetch, + ) + prefetch = True self._cloudfetch_prefetch = bool(prefetch) self._ssl_options = ssl_options diff --git a/src/databricks/sql/cloudfetch/downloader.py b/src/databricks/sql/cloudfetch/downloader.py index dbaca451f..ed134e58e 100644 --- a/src/databricks/sql/cloudfetch/downloader.py +++ b/src/databricks/sql/cloudfetch/downloader.py @@ -145,6 +145,8 @@ def run(self) -> DownloadedFile: self.link, self.settings.link_expiry_buffer_secs ) self._raise_if_cancelled() + # Speed metrics cover the successful attempt only + start_time = time.time() with self._http_client.request_context( method=HttpMethod.GET, url=self.link.fileLink, diff --git a/tests/unit/test_cloudfetch_prefetch.py b/tests/unit/test_cloudfetch_prefetch.py index 534cf7e72..a654f347d 100644 --- a/tests/unit/test_cloudfetch_prefetch.py +++ b/tests/unit/test_cloudfetch_prefetch.py @@ -624,9 +624,18 @@ def test_prefetch_is_on_by_default_and_can_be_disabled(self): self.assertFalse( self._client(_cloudfetch_prefetch=off)._cloudfetch_prefetch ) - for on in (True, "true", "1", None): + for on in (True, "true", "1", "yes", "on", None): self.assertTrue(self._client(_cloudfetch_prefetch=on)._cloudfetch_prefetch) + def test_unrecognized_strings_keep_prefetch_on_with_a_warning(self): + with self.assertLogs( + "databricks.sql.backend.thrift_backend", level="WARNING" + ) as logs: + client = self._client(_cloudfetch_prefetch="disabled") + + self.assertTrue(client._cloudfetch_prefetch) + self.assertIn("_cloudfetch_prefetch='disabled'", logs.output[0]) + # get_schemas/get_tables/get_columns import ThriftResultSet locally @patch("databricks.sql.result_set.ThriftResultSet") @patch("databricks.sql.backend.thrift_backend.ThriftResultSet") From 66eeda5d531117ab6403f99386363d7695f46df7 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 12 Nov 2026 14:21:55 +0000 Subject: [PATCH 14/17] Stop downloads when a mid-stream link fetch fails Signed-off-by: Vu Anh Phung Co-authored-by: Isaac --- src/databricks/sql/cloudfetch/download_manager.py | 4 +++- tests/unit/test_cloudfetch_prefetch.py | 2 ++ 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/src/databricks/sql/cloudfetch/download_manager.py b/src/databricks/sql/cloudfetch/download_manager.py index 7de639a36..be3b1315d 100644 --- a/src/databricks/sql/cloudfetch/download_manager.py +++ b/src/databricks/sql/cloudfetch/download_manager.py @@ -160,8 +160,10 @@ def _schedule_downloads(self): ) except Exception as e: # This read is failing part-way, so the caller loses the rows it had - # collected. Fail every later read too, so a retry cannot skip them. + # collected. Fail every later read too, so a retry cannot skip them, + # and stop the downloads nobody will read. self._link_error = e + self._shutdown_manager() raise if self._cancelled.is_set(): # Shut down while the request was in flight diff --git a/tests/unit/test_cloudfetch_prefetch.py b/tests/unit/test_cloudfetch_prefetch.py index a654f347d..f55f8e4ad 100644 --- a/tests/unit/test_cloudfetch_prefetch.py +++ b/tests/unit/test_cloudfetch_prefetch.py @@ -149,6 +149,8 @@ def test_a_failed_link_fetch_fails_every_later_read(self): with self.assertRaises(RuntimeError) as raised: manager.get_next_downloaded_file(0) self.assertIs(raised.exception, error) + # Downloads nobody can read any more are stopped + self.assertTrue(manager._cancelled.is_set()) # Later reads raise a fresh error that points at the original one for _ in range(2): with self.assertRaises(OperationalError) as raised: From 677f4bcf8a49b7627aaa8d66c4cac525bf37e350 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 12 Nov 2026 14:33:39 +0000 Subject: [PATCH 15/17] Turn prefetch off for unrecognized _cloudfetch_prefetch strings A kill switch should fail safe: an unparsable value falls back to per-batch fetching and logs a warning. None still means unset (prefetch on). Signed-off-by: Vu Anh Phung Co-authored-by: Isaac --- src/databricks/sql/backend/thrift_backend.py | 7 ++++--- tests/unit/test_cloudfetch_prefetch.py | 18 ++++++++++-------- 2 files changed, 14 insertions(+), 11 deletions(-) diff --git a/src/databricks/sql/backend/thrift_backend.py b/src/databricks/sql/backend/thrift_backend.py index a6f913274..fc4f5c11c 100644 --- a/src/databricks/sql/backend/thrift_backend.py +++ b/src/databricks/sql/backend/thrift_backend.py @@ -205,12 +205,13 @@ def __init__( elif value in ("true", "1", "yes", "on"): prefetch = True else: + # Someone setting the kill switch likely wants it off: fail safe logger.warning( - "Unrecognized _cloudfetch_prefetch=%r; keeping CloudFetch " - "prefetch on (use 'false' to turn it off)", + "Unrecognized _cloudfetch_prefetch=%r; turning CloudFetch " + "prefetch off (use 'true' or 'false')", prefetch, ) - prefetch = True + prefetch = False self._cloudfetch_prefetch = bool(prefetch) self._ssl_options = ssl_options diff --git a/tests/unit/test_cloudfetch_prefetch.py b/tests/unit/test_cloudfetch_prefetch.py index f55f8e4ad..ba2a76324 100644 --- a/tests/unit/test_cloudfetch_prefetch.py +++ b/tests/unit/test_cloudfetch_prefetch.py @@ -629,14 +629,16 @@ def test_prefetch_is_on_by_default_and_can_be_disabled(self): for on in (True, "true", "1", "yes", "on", None): self.assertTrue(self._client(_cloudfetch_prefetch=on)._cloudfetch_prefetch) - def test_unrecognized_strings_keep_prefetch_on_with_a_warning(self): - with self.assertLogs( - "databricks.sql.backend.thrift_backend", level="WARNING" - ) as logs: - client = self._client(_cloudfetch_prefetch="disabled") - - self.assertTrue(client._cloudfetch_prefetch) - self.assertIn("_cloudfetch_prefetch='disabled'", logs.output[0]) + def test_unrecognized_strings_turn_prefetch_off_with_a_warning(self): + for value in ("disabled", "flase", ""): + with self.subTest(value): + with self.assertLogs( + "databricks.sql.backend.thrift_backend", level="WARNING" + ) as logs: + client = self._client(_cloudfetch_prefetch=value) + + self.assertFalse(client._cloudfetch_prefetch) + self.assertIn(f"_cloudfetch_prefetch={value!r}", logs.output[0]) # get_schemas/get_tables/get_columns import ThriftResultSet locally @patch("databricks.sql.result_set.ThriftResultSet") From 6e14f55b8715f4abb5f55d27af56c0bd23a5842a Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 12 Nov 2026 14:45:29 +0000 Subject: [PATCH 16/17] Reject row-based replies in fetch_result_links too Signed-off-by: Vu Anh Phung Co-authored-by: Isaac --- src/databricks/sql/backend/thrift_backend.py | 1 + tests/unit/test_cloudfetch_prefetch.py | 7 +++++++ 2 files changed, 8 insertions(+) diff --git a/src/databricks/sql/backend/thrift_backend.py b/src/databricks/sql/backend/thrift_backend.py index fc4f5c11c..7b7643fc2 100644 --- a/src/databricks/sql/backend/thrift_backend.py +++ b/src/databricks/sql/backend/thrift_backend.py @@ -1453,6 +1453,7 @@ def fetch_result_links( if ( results.arrowBatches or results.columns + or results.rows or ( metadata is not None and metadata.resultFormat is not None diff --git a/tests/unit/test_cloudfetch_prefetch.py b/tests/unit/test_cloudfetch_prefetch.py index ba2a76324..e7cb9aaab 100644 --- a/tests/unit/test_cloudfetch_prefetch.py +++ b/tests/unit/test_cloudfetch_prefetch.py @@ -730,6 +730,13 @@ def test_fetch_result_links_rejects_rows_that_are_not_links( results=ttypes.TRowSet(startRowOffset=0, rows=[]), resultSetMetadata=arrow_format, ), + ttypes.TFetchResultsResp( + status=ok, + hasMoreRows=False, + results=ttypes.TRowSet( + startRowOffset=0, rows=[ttypes.TRow(colVals=[])] + ), + ), ): tcli_service_class.return_value.FetchResults.return_value = resp with self.assertRaises(DataError): From a83f50f7f306c81c88c5840550db2b64e2caaa45 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 12 Nov 2026 14:59:03 +0000 Subject: [PATCH 17/17] Keep returning downloaded files after a failed link fetch A failed mid-stream link fetch now stops scheduling and is raised only once the files already downloading run out, so a reader that never needs the missing rows still succeeds, as on main. The first read that needs them gets the original error and later reads fail too. Signed-off-by: Vu Anh Phung Co-authored-by: Isaac --- .../sql/cloudfetch/download_manager.py | 34 ++++++++++++------- tests/unit/test_cloudfetch_prefetch.py | 28 ++++++++++++--- 2 files changed, 45 insertions(+), 17 deletions(-) diff --git a/src/databricks/sql/cloudfetch/download_manager.py b/src/databricks/sql/cloudfetch/download_manager.py index be3b1315d..2a6cfc355 100644 --- a/src/databricks/sql/cloudfetch/download_manager.py +++ b/src/databricks/sql/cloudfetch/download_manager.py @@ -4,7 +4,7 @@ import threading from concurrent.futures import CancelledError, ThreadPoolExecutor, Future -from typing import Callable, List, Union, Tuple, Optional, TYPE_CHECKING +from typing import Callable, List, NoReturn, Union, Tuple, Optional, TYPE_CHECKING from databricks.sql.exc import OperationalError from databricks.sql.cloudfetch.downloader import ( @@ -66,8 +66,9 @@ def __init__( self._cancelled = threading.Event() # Set by cancel(): reads fail once they need data the cancel stopped. self._cancel_requested = False - # A failed link fetch, raised again by every later read (see _schedule_downloads) + # A failed link fetch, raised once the downloaded files run out self._link_error: Optional[BaseException] = None + self._link_error_raised = False def cancel(self) -> None: """Stop downloading because the statement was cancelled. @@ -108,17 +109,14 @@ def get_next_downloaded_file( next_row_offset (int): The offset of the starting row of the next file we want data from. """ - if self._link_error is not None: - raise OperationalError( - "Fetching more CloudFetch results failed earlier; the result can't " - "be read further, execute the statement again" - ) from self._link_error - if not self._cancel_requested: + if not self._cancel_requested and self._link_error is None: # Make sure the download queue is always full self._schedule_downloads() # No more files to download from this batch of links if len(self._download_tasks) == 0: + if self._link_error is not None: + self._raise_link_error(self._link_error) if self._cancel_requested and ( self._pending_links or self._link_source is not None ): @@ -146,6 +144,18 @@ def get_next_downloaded_file( return file + def _raise_link_error(self, error: BaseException) -> NoReturn: + # That read loses the rows it had collected, so every later read fails too + # and a retry can't continue past them. + self._shutdown_manager() + if not self._link_error_raised: + self._link_error_raised = True + raise error + raise OperationalError( + "Fetching more CloudFetch results failed earlier; the result can't " + "be read further, execute the statement again" + ) from error + def _schedule_downloads(self): self._submit_pending_links() while ( @@ -159,12 +169,10 @@ def _schedule_downloads(self): self._max_download_threads - scheduled ) except Exception as e: - # This read is failing part-way, so the caller loses the rows it had - # collected. Fail every later read too, so a retry cannot skip them, - # and stop the downloads nobody will read. + # Stop scheduling; files already downloading are still returned and + # the read that needs rows past them raises (_raise_link_error) self._link_error = e - self._shutdown_manager() - raise + return if self._cancelled.is_set(): # Shut down while the request was in flight return diff --git a/tests/unit/test_cloudfetch_prefetch.py b/tests/unit/test_cloudfetch_prefetch.py index e7cb9aaab..bb2a5c9e6 100644 --- a/tests/unit/test_cloudfetch_prefetch.py +++ b/tests/unit/test_cloudfetch_prefetch.py @@ -140,24 +140,34 @@ def test_reply_with_only_empty_links_is_retried_on_the_next_call(self): self.assertIsNone(manager.get_next_downloaded_file(20)) @patch.object(download_manager, "ResultSetDownloadHandler", _InstantHandler) - def test_a_failed_link_fetch_fails_every_later_read(self): + def test_a_failed_link_fetch_fails_reads_once_downloaded_files_run_out(self): manager = self._manager([_link(0)], max_download_threads=2) error = RuntimeError("FetchResults failed") source = Mock(side_effect=[error, ([_link(10)], False)]) manager.set_link_source(source) + # The file already downloading is still returned + self.assertEqual(manager.get_next_downloaded_file(0).start_row_offset, 0) with self.assertRaises(RuntimeError) as raised: - manager.get_next_downloaded_file(0) + manager.get_next_downloaded_file(10) self.assertIs(raised.exception, error) - # Downloads nobody can read any more are stopped self.assertTrue(manager._cancelled.is_set()) # Later reads raise a fresh error that points at the original one for _ in range(2): with self.assertRaises(OperationalError) as raised: - manager.get_next_downloaded_file(0) + manager.get_next_downloaded_file(10) self.assertIs(raised.exception.__cause__, error) self.assertEqual(source.call_count, 1) + @patch.object(download_manager, "ResultSetDownloadHandler", _InstantHandler) + def test_with_nothing_in_flight_empty_replies_are_retried_at_once(self): + manager = self._manager([], max_download_threads=2) + source = Mock(side_effect=[([], True), ([], True), ([_link(0)], False)]) + manager.set_link_source(source) + + self.assertEqual(manager.get_next_downloaded_file(0).start_row_offset, 0) + self.assertEqual(source.call_count, 3) + @patch.object(download_manager, "ResultSetDownloadHandler", _InstantHandler) def test_shutdown_during_a_link_fetch_ends_the_stream_cleanly(self): manager = self._manager([_link(0)], max_download_threads=2) @@ -591,6 +601,16 @@ def test_cancelled_downloads_fail_the_next_fetch_instead_of_truncating(self): with self.assertRaises(OperationalError): rs.fetchall_arrow() + def test_a_reader_that_never_needs_the_failed_links_is_not_failed(self): + rs, backend, _ = self._result_set() + backend.fetch_result_links.side_effect = DataError("FetchResults failed") + + # Rows 0-14 come from the two initial links, which were already downloading + table = rs.fetchmany_arrow(15) + + self.assertEqual(table.column(0).to_pylist(), list(range(15))) + rs.close() + def test_a_retry_after_a_failed_link_fetch_fails_instead_of_skipping_rows(self): rs, backend, _ = self._result_set() backend.fetch_result_links.side_effect = [