iterate over chunk indexes instead of link

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

Author: varun-edachali-dbx

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

@@ -140,6 +140,7 @@ def __init__(
 
         self._sea_client = sea_client
         self._statement_id = statement_id
+        self._total_chunk_count = total_chunk_count
 
         logger.debug(
             "SeaCloudFetchQueue: Initialize CloudFetch loader for statement {}, total chunks: {}".format(
@@ -154,10 +155,9 @@ def __init__(
             return None
 
         # Track the current chunk we're processing
-        self._current_chunk_link = first_link
-
+        self._current_chunk_index = 0
         # Initialize table and position
-        self.table = self._create_table_from_link(self._current_chunk_link)
+        self.table = self._create_table_from_link(first_link)
 
     def _convert_to_thrift_link(self, link: "ExternalLink") -> TSparkArrowResultLink:
         """Convert SEA external links to Thrift format for compatibility with existing download manager."""
@@ -172,34 +172,24 @@ def _convert_to_thrift_link(self, link: "ExternalLink") -> TSparkArrowResultLink
             httpHeaders=link.http_headers or {},
         )
 
-    def _progress_chunk_link(self):
+    def _get_chunk_link(self, chunk_index: int) -> Optional["ExternalLink"]:
         """Progress to the next chunk link."""
-        if not self._current_chunk_link:
-            return None
-
-        next_chunk_index = self._current_chunk_link.next_chunk_index
-
-        if next_chunk_index is None:
-            self._current_chunk_link = None
+        if chunk_index >= self._total_chunk_count:
             return None
 
         try:
-            self._current_chunk_link = self._sea_client.get_chunk_link(
-                self._statement_id, next_chunk_index
+            return self._sea_client.get_chunk_link(
+                self._statement_id, chunk_index
             )
         except Exception as e:
             raise ServerOperationError(
-                f"Error fetching link for chunk {next_chunk_index}: {e}",
+                f"Error fetching link for chunk {chunk_index}: {e}",
                 {
                     "operation-id": self._statement_id,
                     "diagnostic-info": None,
                 },
             )
 
-        logger.debug(
-            f"SeaCloudFetchQueue: Progressed to link for chunk {next_chunk_index}: {self._current_chunk_link}"
-        )
-
     def _create_table_from_link(
         self, link: "ExternalLink"
     ) -> Union["pyarrow.Table", None]:
@@ -216,10 +206,11 @@ def _create_table_from_link(
     def _create_next_table(self) -> Union["pyarrow.Table", None]:
         """Create next table by retrieving the logical next downloaded file."""
 
-        self._progress_chunk_link()
+        print(self._current_chunk_index)
+        self._current_chunk_index += 1
+        next_chunk_link = self._get_chunk_link(self._current_chunk_index)
 
-        if not self._current_chunk_link:
+        if not next_chunk_link:
             logger.debug("SeaCloudFetchQueue: No current chunk link, returning")
             return None
-
-        return self._create_table_from_link(self._current_chunk_link)
+        return self._create_table_from_link(next_chunk_link)

tests/unit/test_sea_queue.py

@@ -364,7 +364,7 @@ def test_init_with_valid_initial_link(
 
         # Verify attributes
         assert queue._statement_id == "test-statement-123"
-        assert queue._current_chunk_link == sample_external_link
+        assert queue._current_chunk_index == 0
 
     @patch("databricks.sql.backend.sea.queue.ResultFileDownloadManager")
     @patch("databricks.sql.backend.sea.queue.logger")
@@ -390,162 +390,31 @@ def test_init_no_initial_links(
         )
         assert queue.table is None
 
-    @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."""
-        # Create a queue instance without initializing
-        queue = Mock(spec=SeaCloudFetchQueue)
-        queue._current_chunk_link = None
-
-        # Call the method directly
-        result = SeaCloudFetchQueue._progress_chunk_link(queue)
-
-        # Verify the result is None
-        assert result is None
-
-    @patch("databricks.sql.backend.sea.queue.logger")
-    def test_progress_chunk_link_no_next_chunk(self, mock_logger):
-        """Test _progress_chunk_link with no next chunk index."""
-        # Create a queue instance without initializing
-        queue = Mock(spec=SeaCloudFetchQueue)
-        queue._current_chunk_link = ExternalLink(
-            external_link="https://example.com/data/chunk0",
-            expiration="2025-07-03T05:51:18.118009",
-            row_count=100,
-            byte_count=1024,
-            row_offset=0,
-            chunk_index=0,
-            next_chunk_index=None,
-            http_headers={"Authorization": "Bearer token123"},
-        )
-
-        # Call the method directly
-        result = SeaCloudFetchQueue._progress_chunk_link(queue)
-
-        # Verify the result is None
-        assert result is None
-        assert queue._current_chunk_link is None
-
-    @patch("databricks.sql.backend.sea.queue.logger")
-    def test_progress_chunk_link_success(self, mock_logger, mock_sea_client):
-        """Test _progress_chunk_link with successful progression."""
-        # Create a queue instance without initializing
-        queue = Mock(spec=SeaCloudFetchQueue)
-        queue._current_chunk_link = ExternalLink(
-            external_link="https://example.com/data/chunk0",
-            expiration="2025-07-03T05:51:18.118009",
-            row_count=100,
-            byte_count=1024,
-            row_offset=0,
-            chunk_index=0,
-            next_chunk_index=1,
-            http_headers={"Authorization": "Bearer token123"},
-        )
-        queue._sea_client = mock_sea_client
-        queue._statement_id = "test-statement-123"
-
-        # Setup the mock client to return a new link
-        next_link = ExternalLink(
-            external_link="https://example.com/data/chunk1",
-            expiration="2025-07-03T05:51:18.235843",
-            row_count=50,
-            byte_count=512,
-            row_offset=100,
-            chunk_index=1,
-            next_chunk_index=None,
-            http_headers={"Authorization": "Bearer token123"},
-        )
-        mock_sea_client.get_chunk_link.return_value = next_link
-
-        # Call the method directly
-        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 debug message was logged
-        mock_logger.debug.assert_called_with(
-            f"SeaCloudFetchQueue: Progressed to link for chunk 1: {next_link}"
-        )
-
-    @patch("databricks.sql.backend.sea.queue.logger")
-    def test_progress_chunk_link_error(self, mock_logger, mock_sea_client):
-        """Test _progress_chunk_link with error during chunk fetch."""
-        # Create a queue instance without initializing
-        queue = Mock(spec=SeaCloudFetchQueue)
-        queue._current_chunk_link = ExternalLink(
-            external_link="https://example.com/data/chunk0",
-            expiration="2025-07-03T05:51:18.118009",
-            row_count=100,
-            byte_count=1024,
-            row_offset=0,
-            chunk_index=0,
-            next_chunk_index=1,
-            http_headers={"Authorization": "Bearer token123"},
-        )
-        queue._sea_client = mock_sea_client
-        queue._statement_id = "test-statement-123"
-
-        # Setup the mock client to raise an error
-        error_message = "Network error"
-        mock_sea_client.get_chunk_link.side_effect = Exception(error_message)
-
-        # Call the method directly
-        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)
-
-    @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."""
-        # Create a queue instance without initializing
-        queue = Mock(spec=SeaCloudFetchQueue)
-        queue._current_chunk_link = None
-
-        # Call the method directly
-        result = SeaCloudFetchQueue._create_next_table(queue)
-
-        # Verify debug message was logged
-        mock_logger.debug.assert_called_with(
-            "SeaCloudFetchQueue: No current chunk link, returning"
-        )
-
-        # Verify the result is None
-        assert result is None
-
     @patch("databricks.sql.backend.sea.queue.logger")
     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_link = ExternalLink(
-            external_link="https://example.com/data/chunk0",
-            expiration="2025-07-03T05:51:18.118009",
-            row_count=100,
-            byte_count=1024,
-            row_offset=50,
-            chunk_index=0,
-            next_chunk_index=1,
-            http_headers={"Authorization": "Bearer token123"},
-        )
+        queue._current_chunk_index = 0
         queue.download_manager = Mock()
 
         # Mock the dependencies
         mock_table = Mock()
-        queue._create_table_at_offset = Mock(return_value=mock_table)
+        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._progress_chunk_link = Mock()
 
         # Call the method directly
         result = SeaCloudFetchQueue._create_next_table(queue)
 
-        # Verify the table was created
-        queue._create_table_from_link.assert_called_once_with(queue._current_chunk_link)
+        # Verify the chunk index was incremented
+        assert queue._current_chunk_index == 1
+
+        # Verify the chunk link was retrieved
+        queue._get_chunk_link.assert_called_once_with(1)
 
-        # Verify progress was called
-        queue._progress_chunk_link.assert_called_once()
+        # 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