account for empty response in LinkFetcher init

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

Author: varun-edachali-dbx

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

@@ -263,53 +263,50 @@ def __init__(
             description=description,
         )
 
-        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(
                 statement_id, total_chunk_count
             )
         )
 
         initial_links = result_data.external_links or []
-        self._chunk_index_to_link = {link.chunk_index: link for link in initial_links}
 
         # Track the current chunk we're processing
         self._current_chunk_index = 0
-        first_link = self._chunk_index_to_link.get(self._current_chunk_index, None)
-        if not first_link:
-            # possibly an empty response
-            return
-
-        self.current_chunk_index = 0
-
-        self.link_fetcher = LinkFetcher(
-            download_manager=self.download_manager,
-            backend=self._sea_client,
-            statement_id=self._statement_id,
-            initial_links=initial_links,
-            total_chunk_count=total_chunk_count,
-        )
-        self.link_fetcher.start()
+
+        if total_chunk_count < 1:
+            # an empty response
+            self.link_fetcher = None
+        else:
+            self.link_fetcher = LinkFetcher(
+                download_manager=self.download_manager,
+                backend=sea_client,
+                statement_id=statement_id,
+                initial_links=initial_links,
+                total_chunk_count=total_chunk_count,
+            )
+            self.link_fetcher.start()
 
         # Initialize table and position
         self.table = self._create_next_table()
 
     def _create_next_table(self) -> Union["pyarrow.Table", None]:
         """Create next table by retrieving the logical next downloaded file."""
-        chunk_link = self.link_fetcher.get_chunk_link(self.current_chunk_index)
+        if self.link_fetcher is None:
+            return None
+
+        chunk_link = self.link_fetcher.get_chunk_link(self._current_chunk_index)
         if not chunk_link:
             return None
 
         row_offset = chunk_link.row_offset
         arrow_table = self._create_table_at_offset(row_offset)
 
-        self.current_chunk_index += 1
+        self._current_chunk_index += 1
 
         return arrow_table
 
     def close(self):
         super().close()
-        self.link_fetcher.stop()
+        if self.link_fetcher:
+            self.link_fetcher.stop()

tests/unit/test_sea_queue.py

@@ -359,7 +359,6 @@ def test_init_with_valid_initial_link(
         )
 
         # Verify attributes
-        assert queue._statement_id == "test-statement-123"
         assert queue._current_chunk_index == 0
 
     @patch("databricks.sql.backend.sea.queue.ResultFileDownloadManager")
@@ -391,7 +390,7 @@ 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()
 
@@ -405,7 +404,7 @@ def test_create_next_table_success(self, mock_logger):
         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.link_fetcher.get_chunk_link.assert_called_once_with(0)