notify listeners on scheduling downloads

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

Author: varun-edachali-dbx

src/databricks/sql/cloudfetch/download_manager.py

@@ -52,8 +52,6 @@ def get_next_downloaded_file(self, next_row_offset: int) -> DownloadedFile:
         This function gets the next downloaded file in which its rows start at the specified next_row_offset
         in relation to the full result. File downloads are scheduled if not already, and once the correct
         download handler is located, the function waits for the download status and returns the resulting file.
-        If there are no more downloads, a download was not successful, or the correct file could not be located,
-        this function shuts down the thread pool and returns None.
 
         Args:
             next_row_offset (int): The offset of the starting row of the next file we want data from.
@@ -63,7 +61,7 @@ def get_next_downloaded_file(self, next_row_offset: int) -> DownloadedFile:
         self._schedule_downloads()
 
         # No more files to download from this batch of links
-        if len(self._download_tasks) == 0:
+        while len(self._download_tasks) == 0:
             if self._thread_pool._shutdown:
                 raise Error("download manager shut down before file was ready")
             with self._download_condition:
@@ -105,6 +103,9 @@ def _schedule_downloads(self):
             task = self._thread_pool.submit(handler.run)
             self._download_tasks.append(task)
 
+        with self._download_condition:
+            self._download_condition.notify_all()
+
     def add_link(self, link: TSparkArrowResultLink):
         """
         Add more links to the download manager.
@@ -123,6 +124,8 @@ def add_link(self, link: TSparkArrowResultLink):
         )
         self._pending_links.append(link)
 
+        self._schedule_downloads()
+
     def _shutdown_manager(self):
         # Clear download handlers and shutdown the thread pool
         self._pending_links = []