diff --git a/CHANGELOG.md b/CHANGELOG.md index 4038e882d..f4675550b 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 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/backend/thrift_backend.py b/src/databricks/sql/backend/thrift_backend.py index ee904c368..7b7643fc2 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,24 @@ def __init__( # Cloud fetch self._max_download_threads = kwargs.get("max_download_threads", 10) + 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 + value = prefetch.strip().lower() + if value in ("false", "0", "no", "off"): + prefetch = False + 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; turning CloudFetch " + "prefetch off (use 'true' or 'false')", + prefetch, + ) + prefetch = False + self._cloudfetch_prefetch = bool(prefetch) self._ssl_options = ssl_options self._auth_provider = auth_provider @@ -1001,6 +1022,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 +1177,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 +1216,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 +1263,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 +1314,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 +1365,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 +1401,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 +1427,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 +1435,52 @@ 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 + ) + 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 results.rows + 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, + 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/client.py b/src/databricks/sql/client.py index 613b62386..e49e35ba6 100755 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -1958,7 +1958,17 @@ 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) + 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, 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 ece7406b7..2a6cfc355 100644 --- a/src/databricks/sql/cloudfetch/download_manager.py +++ b/src/databricks/sql/cloudfetch/download_manager.py @@ -1,10 +1,12 @@ from __future__ import annotations import logging +import threading -from concurrent.futures import ThreadPoolExecutor, Future -from typing import List, Union, Tuple, Optional, TYPE_CHECKING +from concurrent.futures import CancelledError, ThreadPoolExecutor, Future +from typing import Callable, List, NoReturn, Union, Tuple, Optional, TYPE_CHECKING +from databricks.sql.exc import OperationalError from databricks.sql.cloudfetch.downloader import ( ResultSetDownloadHandler, DownloadableResultSettings, @@ -21,6 +23,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 +61,37 @@ 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() + # Set by cancel(): reads fail once they need data the cancel stopped. + self._cancel_requested = False + # 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. + + 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 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 @@ -73,11 +109,19 @@ def get_next_downloaded_file( next_row_offset (int): The offset of the starting row of the next file we want data from. """ - # Make sure the download queue is always full - self._schedule_downloads() + 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 + ): + # The cancel stopped part of the result: fail rather than end early + raise OperationalError("CloudFetch download cancelled") self._shutdown_manager() return None @@ -85,7 +129,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 ): @@ -97,7 +144,49 @@ 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 ( + self._link_source is not None + and not self._pending_links + and len(self._download_tasks) < self._max_download_threads + ): + scheduled = len(self._download_tasks) + try: + links, has_more = self._link_source( + self._max_download_threads - scheduled + ) + except Exception as e: + # Stop scheduling; files already downloading are still returned and + # the read that needs rows past them raises (_raise_link_error) + self._link_error = e + return + 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: + self.add_link(link) + self._submit_pending_links() + 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): """ While download queue has a capacity, peek pending links and submit them to thread pool. """ @@ -119,8 +208,17 @@ 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) + try: + task = self._thread_pool.submit(handler.run) + except RuntimeError: + # 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) def add_link(self, link: TSparkArrowResultLink): @@ -144,6 +242,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..ed134e58e 100644 --- a/src/databricks/sql/cloudfetch/downloader.py +++ b/src/databricks/sql/cloudfetch/downloader.py @@ -19,10 +19,23 @@ # 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 +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, after a backoff that doubles each time. +_BODY_READ_ATTEMPTS = 3 +_BODY_RETRY_DELAY_SECS = 0.5 + @dataclass class DownloadedFile: @@ -70,6 +83,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 +92,22 @@ 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 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: @@ -100,17 +130,44 @@ 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( - 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 + for attempt in range(1, _BODY_READ_ATTEMPTS + 1): + if attempt > 1: + # 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 + ) + 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, + 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 2a2599259..e4a0c3b92 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 @@ -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. @@ -214,6 +217,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. @@ -229,7 +233,11 @@ 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 """ self.num_chunks = 0 @@ -238,6 +246,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 +293,56 @@ 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 + ) + 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 + # 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 _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: @@ -322,6 +382,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..bb2a5c9e6 --- /dev/null +++ b/tests/unit/test_cloudfetch_prefetch.py @@ -0,0 +1,816 @@ +"""CloudFetch link prefetching across batches and cancel-on-close.""" + +import threading +import unittest +from concurrent.futures import Future +from contextlib import contextmanager +from unittest.mock import MagicMock, Mock, patch + +import pytest +from urllib3.exceptions import ProtocolError + +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 +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, OperationalError, RequestError +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 _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( + 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) + + @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_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(10) + self.assertIs(raised.exception, error) + 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(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) + + def source(free): + manager._shutdown_manager() + 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_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_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(10) + + @patch.object(download_manager, "ResultSetDownloadHandler", _InstantHandler) + 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() # 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) + + 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 + 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: + manager._schedule_downloads() + self.assertIs(handler.call_args.kwargs["cancel_event"], manager._cancelled) + manager._shutdown_manager() + + +class DownloaderCancelTests(unittest.TestCase): + 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 + + @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 **_: self._request_context( + pending.pop(0) + ) + return http_client + + def _handler(self, http_client, event, lz4_compressed=True): + return downloader.ResultSetDownloadHandler( + downloader.DownloadableResultSettings(is_lz4_compressed=lz4_compressed), + _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(OperationalError): + self._handler(http_client, event).run() + http_client.request_context.assert_not_called() + + 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) + + file = self._handler(http_client, threading.Event(), False).run() + + self.assertEqual(file.file_bytes, b"abcd") + self.assertTrue(response.released) + self.assertIs( + http_client.request_context.call_args.kwargs["preload_content"], False + ) + + 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(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() + + @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, 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() + 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)] + ) + 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, 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): + 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_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)) + + 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) + + +@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() + + 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_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 = [ + DataError("FetchResults failed"), + ([_link(20, self.ROWS_PER_LINK)], False), + ] + + 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(OperationalError) as raised: + rs.fetchmany_arrow(25) + self.assertIsInstance(raised.exception.__cause__, DataError) + self.assertEqual(backend.fetch_result_links.call_count, 1) + + +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) + for off in (False, "false", "False", "0", "no", "off"): + self.assertFalse( + self._client(_cloudfetch_prefetch=off)._cloudfetch_prefetch + ) + for on in (True, "true", "1", "yes", "on", None): + self.assertTrue(self._client(_cloudfetch_prefetch=on)._cloudfetch_prefetch) + + 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") + @patch("databricks.sql.backend.thrift_backend.ThriftResultSet") + @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True) + 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._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): + 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, + ) + + @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, + ), + 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): + client.fetch_result_links( + command_id=Mock(), + max_rows=1, + max_bytes=1, + expected_row_start_offset=0, + ) + + +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_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_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() + 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 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