return None for out of range

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

Author: varun-edachali-dbx

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

@@ -231,17 +231,15 @@ def _trigger_next_batch_download(self) -> bool:
         )
         return True
 
-    def get_chunk_link(self, chunk_index: int) -> ExternalLink:
+    def get_chunk_link(self, chunk_index: int) -> Optional[ExternalLink]:
         """Return (blocking) the :class:`ExternalLink` associated with *chunk_index*."""
         logger.debug(
             "LinkFetcher[%s]: waiting for link of chunk %d",
             self._statement_id,
             chunk_index,
         )
         if chunk_index >= self.total_chunk_count:
-            raise ValueError(
-                f"Chunk index {chunk_index} is out of range for total chunk count {self.total_chunk_count}"
-            )
+            return None
 
         with self._link_data_update:
             while chunk_index not in self.chunk_index_to_link:
@@ -367,6 +365,8 @@ def _create_next_table(self) -> Union["pyarrow.Table", None]:
             return None
 
         chunk_link = self.link_fetcher.get_chunk_link(self._current_chunk_index)
+        if chunk_link is None:
+            return None
 
         row_offset = chunk_link.row_offset
         # NOTE: link has already been submitted to download manager at this point

tests/unit/test_sea_queue.py

@@ -351,15 +351,9 @@ def test_init_with_valid_initial_link(
                 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 attributes
         assert queue._current_chunk_index == 0
+        assert queue.link_fetcher is not None
 
     @patch("databricks.sql.backend.sea.queue.ResultFileDownloadManager")
     @patch("databricks.sql.backend.sea.queue.logger")
@@ -717,14 +711,10 @@ def _worker():
         # The thread should have finished and captured link1
         assert result_container.get("link") == link1
 
-    def test_get_chunk_link_out_of_range_raises_value_error(self, sample_links):
+    def test_get_chunk_link_out_of_range_returns_none(self, sample_links):
         """Requesting a chunk index >= total_chunk_count should immediately return None."""
         link0, _ = sample_links
 
         fetcher, _backend, _dm = self._create_fetcher([link0], total_chunk_count=1)
 
-        with pytest.raises(
-            ValueError,
-            match="Chunk index 10 is out of range for total chunk count 1",
-        ):
-            fetcher.get_chunk_link(10)
+        assert fetcher.get_chunk_link(10) is None