add_link -> add_links for less heavy link additions

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

Author: varun-edachali-dbx

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

@@ -188,9 +188,11 @@ def _add_links(self, links: List[ExternalLink]):
             len(links),
             ", ".join(str(l.chunk_index) for l in links) if links else "<none>",
         )
-        for link in links:
-            self.chunk_index_to_link[link.chunk_index] = link
-            self.download_manager.add_link(LinkFetcher._convert_to_thrift_link(link))
+
+        self.chunk_index_to_link.update({link.chunk_index: link for link in links})
+        self.download_manager.add_links(
+            [LinkFetcher._convert_to_thrift_link(link) for link in links]
+        )
 
     def _get_next_chunk_index(self) -> Optional[int]:
         """Return the next *chunk_index* that should be requested from the backend, or ``None`` if we have them all."""

src/databricks/sql/cloudfetch/download_manager.py

@@ -166,21 +166,22 @@ def _schedule_downloads(self):
         with self._download_condition:
             self._download_condition.notify_all()
 
-    def add_link(self, link: TSparkArrowResultLink):
+    def add_links(self, links: List[TSparkArrowResultLink]):
         """
         Add more links to the download manager.
 
         Args:
             link (TSparkArrowResultLink): The link to add to the download manager.
         """
-        if link.rowCount <= 0:
-            return
-        logger.debug(
-            "ResultFileDownloadManager: adding file link, start offset {}, row count: {}".format(
-                link.startRowOffset, link.rowCount
+        for link in links:
+            if link.rowCount <= 0:
+                continue
+            logger.debug(
+                "ResultFileDownloadManager: adding file link, start offset {}, row count: {}".format(
+                    link.startRowOffset, link.rowCount
+                )
             )
-        )
-        self._pending_links.append(link)
+            self._pending_links.append(link)
 
         self._schedule_downloads()
 

src/databricks/sql/utils.py

@@ -381,7 +381,7 @@ def __init__(
                         result_link.startRowOffset, result_link.rowCount
                     )
                 )
-                self.download_manager.add_link(result_link)
+            self.download_manager.add_links(self.result_links)
 
             # Initialize table and position
             self.table = self._create_next_table()

tests/unit/test_sea_queue.py

@@ -638,7 +638,7 @@ def test_add_links_and_get_next_chunk_index(self, sample_links):
         fetcher, _backend, download_manager = self._create_fetcher([link0])
 
         # add_link should have been called for the initial link
-        download_manager.add_link.assert_called_once()
+        download_manager.add_links.assert_called_once()
 
         # Internal mapping should contain the link
         assert fetcher.chunk_index_to_link[0] == link0
@@ -668,7 +668,7 @@ def test_trigger_next_batch_download_success(self, sample_links):
         backend.get_chunk_links.assert_called_once_with("statement-123", 1)
         assert fetcher.chunk_index_to_link[1] == link1
         # Two calls to add_link: one for initial link, one for new link
-        assert download_manager.add_link.call_count == 2
+        assert download_manager.add_links.call_count == 2
 
     def test_trigger_next_batch_download_error(self, sample_links):
         """Ensure that errors from backend are captured and surfaced."""