move link fetching immediately before table creation so link expiry is not an issue

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

Author: varun-edachali-dbx

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

@@ -156,10 +156,10 @@ def __init__(
         )
 
         # Track the current chunk we're processing
-        self._current_chunk_link: Optional["ExternalLink"] = initial_link
+        self._current_chunk_link = initial_link
 
         # Initialize table and position
-        self.table = self._create_next_table()
+        self.table = self._create_table_from_link(self._current_chunk_link)
 
     def _convert_to_thrift_link(self, link: "ExternalLink") -> TSparkArrowResultLink:
         """Convert SEA external links to Thrift format for compatibility with existing download manager."""
@@ -202,22 +202,30 @@ def _progress_chunk_link(self):
             f"SeaCloudFetchQueue: Progressed to link for chunk {next_chunk_index}: {self._current_chunk_link}"
         )
 
-    def _create_next_table(self) -> Union["pyarrow.Table", None]:
-        """Create next table by retrieving the logical next downloaded file."""
-        if not self._current_chunk_link:
-            logger.debug("SeaCloudFetchQueue: No current chunk link, returning")
-            return None
+    def _create_table_from_link(
+        self, link: "ExternalLink"
+    ) -> 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(self._current_chunk_link)
+        thrift_link = self._convert_to_thrift_link(link)
         self.download_manager.add_link(thrift_link)
 
-        row_offset = self._current_chunk_link.row_offset
+        row_offset = link.row_offset
         arrow_table = self._create_table_at_offset(row_offset)
 
+        return arrow_table
+
+    def _create_next_table(self) -> Union["pyarrow.Table", None]:
+        """Create next table by retrieving the logical next downloaded file."""
+
         self._progress_chunk_link()
 
-        return arrow_table
+        if not self._current_chunk_link:
+            logger.debug("SeaCloudFetchQueue: No current chunk link, returning")
+            return None
+
+        return self._create_table_from_link(self._current_chunk_link)

tests/unit/test_sea_queue.py

@@ -213,7 +213,9 @@ def test_build_queue_arrow_stream(
 
         with patch(
             "databricks.sql.backend.sea.queue.ResultFileDownloadManager"
-        ), patch.object(SeaCloudFetchQueue, "_create_next_table", return_value=None):
+        ), patch.object(
+            SeaCloudFetchQueue, "_create_table_from_link", return_value=None
+        ):
             queue = SeaResultSetQueueFactory.build_queue(
                 result_data=result_data,
                 manifest=arrow_manifest,
@@ -342,7 +344,9 @@ def test_init_with_valid_initial_link(
         mock_download_manager_class.return_value = mock_download_manager
 
         # Create a queue with valid initial link
-        with patch.object(SeaCloudFetchQueue, "_create_next_table", return_value=None):
+        with patch.object(
+            SeaCloudFetchQueue, "_create_table_from_link", return_value=None
+        ):
             queue = SeaCloudFetchQueue(
                 initial_links=[sample_external_link],
                 max_download_threads=5,
@@ -608,13 +612,14 @@ def test_create_next_table_success(self, mock_logger):
         # Mock the dependencies
         mock_table = Mock()
         queue._create_table_at_offset = Mock(return_value=mock_table)
+        queue._create_table_from_link = Mock(return_value=mock_table)
         queue._progress_chunk_link = Mock()
 
         # Call the method directly
         result = SeaCloudFetchQueue._create_next_table(queue)
 
         # Verify the table was created
-        queue._create_table_at_offset.assert_called_once_with(50)
+        queue._create_table_from_link.assert_called_once_with(queue._current_chunk_link)
 
         # Verify progress was called
         queue._progress_chunk_link.assert_called_once()