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 <vu.phung@databricks.com>
Co-authored-by: Isaac <no-reply@databricks.com>

Author: vuanhphung

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

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 = [