move download_manager init into parent CloudFetchQueue

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

Author: varun-edachali-dbx

src/databricks/sql/backend/sea/queue.py

@@ -19,7 +19,7 @@
     ResultManifest,
 )
 from databricks.sql.backend.sea.utils.constants import ResultFormat
-from databricks.sql.exc import ProgrammingError
+from databricks.sql.exc import ProgrammingError, ServerOperationError
 from databricks.sql.thrift_api.TCLIService.ttypes import TSparkArrowResultLink
 from databricks.sql.types import SSLOptions
 from databricks.sql.utils import CloudFetchQueue, ResultSetQueue
@@ -149,13 +149,6 @@ def __init__(
         if not first_link:
             return
 
-        self.download_manager = ResultFileDownloadManager(
-            links=[],
-            max_download_threads=max_download_threads,
-            lz4_compressed=lz4_compressed,
-            ssl_options=ssl_options,
-        )
-
         # Track the current chunk we're processing
         self._current_chunk_link = first_link
 
@@ -191,13 +184,13 @@ def _progress_chunk_link(self):
                 self._statement_id, next_chunk_index
             )
         except Exception as e:
-            logger.error(
-                "SeaCloudFetchQueue: Error fetching link for chunk {}: {}".format(
-                    next_chunk_index, e
-                )
+            raise ServerOperationError(
+                f"Error fetching link for chunk {next_chunk_index}: {e}",
+                {
+                    "operation-id": self._statement_id,
+                    "diagnostic-info": None,
+                },
             )
-            self._current_chunk_link = None
-            return None
 
         logger.debug(
             f"SeaCloudFetchQueue: Progressed to link for chunk {next_chunk_index}: {self._current_chunk_link}"
@@ -208,10 +201,6 @@ def _create_table_from_link(
     ) -> Union["pyarrow.Table", None]:
         """Create a table from a link."""
 
-        if not self.download_manager:
-            logger.debug("SeaCloudFetchQueue: No download manager, returning")
-            return None
-
         thrift_link = self._convert_to_thrift_link(link)
         self.download_manager.add_link(thrift_link)
 

src/databricks/sql/utils.py

@@ -230,7 +230,12 @@ def __init__(
         self.table_row_index = 0
 
         # Initialize download manager
-        self.download_manager: Optional["ResultFileDownloadManager"] = None
+        self.download_manager = ResultFileDownloadManager(
+            links=[],
+            max_download_threads=max_download_threads,
+            lz4_compressed=lz4_compressed,
+            ssl_options=ssl_options,
+        )
 
     def next_n_rows(self, num_rows: int) -> "pyarrow.Table":
         """
@@ -287,11 +292,8 @@ def remaining_rows(self) -> "pyarrow.Table":
 
     def _create_table_at_offset(self, offset: int) -> Union["pyarrow.Table", None]:
         """Create next table at the given row offset"""
-        # Create next table by retrieving the logical next downloaded file, or return None to signal end of queue
-        if not self.download_manager:
-            logger.debug("CloudFetchQueue: No download manager available")
-            return None
 
+        # Create next table by retrieving the logical next downloaded file, or return None to signal end of queue
         downloaded_file = self.download_manager.get_next_downloaded_file(offset)
         if not downloaded_file:
             logger.debug(
@@ -373,14 +375,7 @@ def __init__(
                         result_link.startRowOffset, result_link.rowCount
                     )
                 )
-
-        # Initialize download manager
-        self.download_manager = ResultFileDownloadManager(
-            links=self.result_links,
-            max_download_threads=self.max_download_threads,
-            lz4_compressed=self.lz4_compressed,
-            ssl_options=self._ssl_options,
-        )
+                self.download_manager.add_link(result_link)
 
         # Initialize table and position
         self.table = self._create_next_table()

tests/unit/test_sea_queue.py

@@ -18,7 +18,7 @@
     ExternalLink,
 )
 from databricks.sql.backend.sea.utils.constants import ResultFormat
-from databricks.sql.exc import ProgrammingError
+from databricks.sql.exc import ProgrammingError, ServerOperationError
 from databricks.sql.types import SSLOptions
 
 
@@ -340,9 +340,6 @@ def test_init_with_valid_initial_link(
         sample_external_link,
     ):
         """Test initialization with valid initial link."""
-        mock_download_manager = Mock()
-        mock_download_manager_class.return_value = mock_download_manager
-
         # Create a queue with valid initial link
         with patch.object(
             SeaCloudFetchQueue, "_create_table_from_link", return_value=None
@@ -365,13 +362,9 @@ def test_init_with_valid_initial_link(
             )
         )
 
-        # Verify download manager was created
-        mock_download_manager_class.assert_called_once()
-
         # Verify attributes
         assert queue._statement_id == "test-statement-123"
         assert queue._current_chunk_link == sample_external_link
-        assert queue.download_manager == mock_download_manager
 
     @patch("databricks.sql.backend.sea.queue.ResultFileDownloadManager")
     @patch("databricks.sql.backend.sea.queue.logger")
@@ -514,21 +507,12 @@ def test_progress_chunk_link_error(self, mock_logger, mock_sea_client):
         mock_sea_client.get_chunk_link.side_effect = Exception(error_message)
 
         # Call the method directly
-        result = SeaCloudFetchQueue._progress_chunk_link(queue)
+        with pytest.raises(ServerOperationError, match=error_message):
+            SeaCloudFetchQueue._progress_chunk_link(queue)
 
         # Verify the client was called
         mock_sea_client.get_chunk_link.assert_called_once_with("test-statement-123", 1)
 
-        # Verify error message was logged
-        mock_logger.error.assert_called_with(
-            "SeaCloudFetchQueue: Error fetching link for chunk {}: {}".format(
-                1, error_message
-            )
-        )
-
-        # Verify the result is None
-        assert result is None
-
     @patch("databricks.sql.backend.sea.queue.logger")
     def test_create_next_table_no_current_link(self, mock_logger):
         """Test _create_next_table with no current link."""