make LinkFetcher convert link static

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

Author: varun-edachali-dbx

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

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
 from abc import ABC
-from typing import List, Optional, Tuple, Union, TYPE_CHECKING
+import threading
+from typing import Dict, List, Optional, Tuple, Union, TYPE_CHECKING
 
 from databricks.sql.cloudfetch.download_manager import ResultFileDownloadManager
 
@@ -142,7 +143,7 @@ def __init__(
 
         for link in initial_links:
             self.chunk_index_to_link[link.chunk_index] = link
-            self.download_manager.add_link(self._convert_to_thrift_link(link))
+            self.download_manager.add_link(LinkFetcher._convert_to_thrift_link(link))
         self.total_chunk_count = total_chunk_count
 
     def _get_next_chunk_index(self) -> Optional[int]:
@@ -188,7 +189,8 @@ def get_chunk_link(self, chunk_index: int) -> Optional["ExternalLink"]:
 
             return self.chunk_index_to_link.get(chunk_index, None)
 
-    def _convert_to_thrift_link(self, link: "ExternalLink") -> TSparkArrowResultLink:
+    @staticmethod
+    def _convert_to_thrift_link(link: "ExternalLink") -> TSparkArrowResultLink:
         """Convert SEA external links to Thrift format for compatibility with existing download manager."""
         # Parse the ISO format expiration time
         expiry_time = int(dateutil.parser.parse(link.expiration).timestamp())

tests/unit/test_sea_queue.py

@@ -11,6 +11,7 @@
 
 from databricks.sql.backend.sea.queue import (
     JsonQueue,
+    LinkFetcher,
     SeaResultSetQueueFactory,
     SeaCloudFetchQueue,
 )
@@ -216,9 +217,7 @@ def test_build_queue_arrow_stream(
 
         with patch(
             "databricks.sql.backend.sea.queue.ResultFileDownloadManager"
-        ), patch.object(
-            SeaCloudFetchQueue, "_create_table_from_link", return_value=None
-        ):
+        ), patch.object(SeaCloudFetchQueue, "_create_next_table", return_value=None):
             queue = SeaResultSetQueueFactory.build_queue(
                 result_data=result_data,
                 manifest=arrow_manifest,
@@ -303,10 +302,8 @@ def sample_external_link_no_headers(self):
 
     def test_convert_to_thrift_link(self, sample_external_link):
         """Test conversion of ExternalLink to TSparkArrowResultLink."""
-        queue = Mock(spec=SeaCloudFetchQueue)
-
         # Call the method directly
-        result = SeaCloudFetchQueue._convert_to_thrift_link(queue, sample_external_link)
+        result = LinkFetcher._convert_to_thrift_link(sample_external_link)
 
         # Verify the conversion
         assert result.fileLink == sample_external_link.external_link
@@ -317,12 +314,8 @@ def test_convert_to_thrift_link(self, sample_external_link):
 
     def test_convert_to_thrift_link_no_headers(self, sample_external_link_no_headers):
         """Test conversion of ExternalLink with no headers to TSparkArrowResultLink."""
-        queue = Mock(spec=SeaCloudFetchQueue)
-
         # Call the method directly
-        result = SeaCloudFetchQueue._convert_to_thrift_link(
-            queue, sample_external_link_no_headers
-        )
+        result = LinkFetcher._convert_to_thrift_link(sample_external_link_no_headers)
 
         # Verify the conversion
         assert result.fileLink == sample_external_link_no_headers.external_link
@@ -344,9 +337,7 @@ def test_init_with_valid_initial_link(
     ):
         """Test initialization with valid initial link."""
         # Create a queue with valid initial link
-        with patch.object(
-            SeaCloudFetchQueue, "_create_table_from_link", return_value=None
-        ):
+        with patch.object(SeaCloudFetchQueue, "_create_next_table", return_value=None):
             queue = SeaCloudFetchQueue(
                 result_data=ResultData(external_links=[sample_external_link]),
                 max_download_threads=5,
@@ -398,29 +389,29 @@ def test_create_next_table_success(self, mock_logger):
         """Test _create_next_table with successful table creation."""
         # Create a queue instance without initializing
         queue = Mock(spec=SeaCloudFetchQueue)
-        queue._current_chunk_index = 0
+        queue.current_chunk_index = 0
         queue.download_manager = Mock()
+        queue.link_fetcher = Mock()
 
         # Mock the dependencies
         mock_table = Mock()
         mock_chunk_link = Mock()
-        queue._get_chunk_link = Mock(return_value=mock_chunk_link)
-        queue._create_table_from_link = Mock(return_value=mock_table)
+        queue.link_fetcher.get_chunk_link = Mock(return_value=mock_chunk_link)
+        queue._create_table_at_offset = Mock(return_value=mock_table)
 
         # Call the method directly
-        result = SeaCloudFetchQueue._create_next_table(queue)
+        SeaCloudFetchQueue._create_next_table(queue)
 
         # Verify the chunk index was incremented
-        assert queue._current_chunk_index == 1
+        assert queue.current_chunk_index == 1
 
         # Verify the chunk link was retrieved
-        queue._get_chunk_link.assert_called_once_with(1)
+        queue.link_fetcher.get_chunk_link.assert_called_once_with(0)
 
         # Verify the table was created from the link
-        queue._create_table_from_link.assert_called_once_with(mock_chunk_link)
-
-        # Verify the result is the table
-        assert result == mock_table
+        queue._create_table_at_offset.assert_called_once_with(
+            mock_chunk_link.row_offset
+        )
 
 
 class TestHybridDisposition:
@@ -494,7 +485,7 @@ def test_hybrid_disposition_with_attachment(
         mock_create_table.assert_called_once_with(attachment_data, description)
 
     @patch("databricks.sql.backend.sea.queue.ResultFileDownloadManager")
-    @patch.object(SeaCloudFetchQueue, "_create_table_from_link", return_value=None)
+    @patch.object(SeaCloudFetchQueue, "_create_next_table", return_value=None)
     def test_hybrid_disposition_with_external_links(
         self,
         mock_create_table,