Decouple download error handling again

Author: TheByronHimes

src/ghga_connector/core/downloading/batch_processing.py

@@ -194,7 +194,7 @@ def _handle_failures(self) -> bool:
         return True
 
     async def manage_file_downloads(self, overwrite: bool) -> AsyncGenerator[FileInfo]:
-        """Manages file downloads by handling errors, checking for existing files,
+        """Manages file downloads by checking for existing files before download,
         printing messages to the display, and renaming files after they are downloaded.
 
         Yields file information.
@@ -203,19 +203,7 @@ async def manage_file_downloads(self, overwrite: bool) -> AsyncGenerator[FileInf
             staged_files = await self.get_staged_files()
             for file_info in staged_files:
                 utils.check_for_existing_file(file_info=file_info, overwrite=overwrite)
-                try:
-                    file_id = file_info.file_id
-                    yield file_info
-                except exceptions.GetEnvelopeError as error:
-                    CLIMessageDisplay.failure(
-                        f"The request to get an envelope for file '{file_id}' failed."
-                    )
-                    raise error
-                except exceptions.DownloadError as error:
-                    CLIMessageDisplay.failure(
-                        f"Failed downloading with id '{file_id}'."
-                    )
-                    raise error
+                yield file_info
                 file_info.path_during_download.rename(file_info.path_once_complete)
                 CLIMessageDisplay.success(
                     f"File with id '{file_info.file_id}' has been successfully downloaded."

src/ghga_connector/core/downloading/downloader.py

@@ -19,6 +19,7 @@
 import gc
 import logging
 from asyncio import PriorityQueue, Queue, Semaphore, Task
+from contextlib import contextmanager
 from io import BufferedWriter
 from pathlib import Path
 
@@ -28,11 +29,29 @@
 
 from ..progress_bar import DownloadProgressBar
 from .api_calls import DownloadClient, extract_download_url
-from .structs import RetryResponse
+from .structs import FileInfo, RetryResponse
 
 logger = logging.getLogger(__name__)
 # TODO: [later] More better logging
 
+__all__ = ["Downloader", "handle_download_errors"]
+
+
+@contextmanager
+def handle_download_errors(file_info: FileInfo):
+    """Used to handle download errors from `Downloader.download_file()`"""
+    file_id = file_info.file_id
+    try:
+        yield
+    except exceptions.GetEnvelopeError as error:
+        CLIMessageDisplay.failure(
+            f"The request to get an envelope for file '{file_id}' failed."
+        )
+        raise error
+    except exceptions.DownloadError as error:
+        CLIMessageDisplay.failure(f"Failed downloading with id '{file_id}'.")
+        raise error
+
 
 class Downloader:
     """Centralized high-level interface for downloading a single file.

src/ghga_connector/core/main.py

@@ -24,7 +24,10 @@
 from ghga_connector.core.client import async_client
 from ghga_connector.core.downloading.api_calls import DownloadClient
 from ghga_connector.core.downloading.batch_processing import FileStager
-from ghga_connector.core.downloading.downloader import Downloader
+from ghga_connector.core.downloading.downloader import (
+    Downloader,
+    handle_download_errors,
+)
 from ghga_connector.core.work_package import WorkPackageClient
 
 from .. import exceptions
@@ -153,9 +156,11 @@ async def async_download(
                 max_concurrent_downloads=CONFIG.max_concurrent_downloads,
             )
             CLIMessageDisplay.display(f"Downloading file with id '{file_id}'...")
-            await downloader.download_file(
-                output_path=file_info.path_during_download, part_size=CONFIG.part_size
-            )
+            with handle_download_errors(file_info):
+                await downloader.download_file(
+                    output_path=file_info.path_during_download,
+                    part_size=CONFIG.part_size,
+                )
 
 
 def decrypt_file(