pass result_data instead of "initial links" into SeaCloudFetchQueue

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

Author: varun-edachali-dbx

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

@@ -63,7 +63,7 @@ def build_queue(
         elif manifest.format == ResultFormat.ARROW_STREAM.value:
             # EXTERNAL_LINKS disposition
             return SeaCloudFetchQueue(
-                initial_links=result_data.external_links or [],
+                result_data=result_data,
                 max_download_threads=max_download_threads,
                 ssl_options=ssl_options,
                 sea_client=sea_client,
@@ -103,7 +103,7 @@ class SeaCloudFetchQueue(CloudFetchQueue):
 
     def __init__(
         self,
-        initial_links: List["ExternalLink"],
+        result_data: ResultData,
         max_download_threads: int,
         ssl_options: SSLOptions,
         sea_client: "SeaDatabricksClient",
@@ -144,8 +144,9 @@ def __init__(
             )
         )
 
-        initial_link = next((l for l in initial_links if l.chunk_index == 0), None)
-        if not initial_link:
+        initial_links = result_data.external_links or []
+        first_link = next((l for l in initial_links if l.chunk_index == 0), None)
+        if not first_link:
             return
 
         self.download_manager = ResultFileDownloadManager(
@@ -156,7 +157,7 @@ def __init__(
         )
 
         # Track the current chunk we're processing
-        self._current_chunk_link = initial_link
+        self._current_chunk_link = first_link
 
         # Initialize table and position
         self.table = self._create_table_from_link(self._current_chunk_link)

tests/unit/test_sea_queue.py

@@ -348,7 +348,7 @@ def test_init_with_valid_initial_link(
             SeaCloudFetchQueue, "_create_table_from_link", return_value=None
         ):
             queue = SeaCloudFetchQueue(
-                initial_links=[sample_external_link],
+                result_data=ResultData(external_links=[sample_external_link]),
                 max_download_threads=5,
                 ssl_options=ssl_options,
                 sea_client=mock_sea_client,
@@ -386,7 +386,7 @@ def test_init_no_initial_links(
         """Test initialization with no initial links."""
         # Create a queue with empty initial links
         queue = SeaCloudFetchQueue(
-            initial_links=[],
+            result_data=ResultData(external_links=[]),
             max_download_threads=5,
             ssl_options=ssl_options,
             sea_client=mock_sea_client,
@@ -413,51 +413,6 @@ def test_init_no_initial_links(
             or queue._current_chunk_link is None
         )
 
-    @patch("databricks.sql.backend.sea.queue.ResultFileDownloadManager")
-    @patch("databricks.sql.backend.sea.queue.logger")
-    def test_init_non_zero_chunk_index(
-        self,
-        mock_logger,
-        mock_download_manager_class,
-        mock_sea_client,
-        ssl_options,
-        description,
-    ):
-        """Test initialization with non-zero chunk index initial link."""
-        # Create a link with chunk_index != 0
-        non_zero_link = ExternalLink(
-            external_link="https://example.com/data/chunk1",
-            expiration="2025-07-03T05:51:18.118009",
-            row_count=100,
-            byte_count=1024,
-            row_offset=100,
-            chunk_index=1,
-            next_chunk_index=2,
-            http_headers={"Authorization": "Bearer token123"},
-        )
-
-        # Create a queue with non-zero chunk index
-        queue = SeaCloudFetchQueue(
-            initial_links=[non_zero_link],
-            max_download_threads=5,
-            ssl_options=ssl_options,
-            sea_client=mock_sea_client,
-            statement_id="test-statement-123",
-            total_chunk_count=1,
-            lz4_compressed=False,
-            description=description,
-        )
-
-        # Verify debug message was logged
-        mock_logger.debug.assert_called_with(
-            "SeaCloudFetchQueue: Initialize CloudFetch loader for statement {}, total chunks: {}".format(
-                "test-statement-123", 1
-            )
-        )
-
-        # Verify download manager wasn't created (no chunk 0)
-        mock_download_manager_class.assert_not_called()
-
     @patch("databricks.sql.backend.sea.queue.logger")
     def test_progress_chunk_link_no_current_link(self, mock_logger):
         """Test _progress_chunk_link with no current link."""