Flatten CLI and further refactor top of download path

Author: TheByronHimes

src/ghga_connector/cli.py

@@ -19,39 +19,17 @@
 import logging
 import os
 import sys
-from dataclasses import dataclass
 from functools import partial
 from pathlib import Path
 from types import TracebackType
 
-import crypt4gh.keys
 import typer
-from ghga_service_commons.utils import crypt
 
 from ghga_connector import exceptions
 from ghga_connector.config import CONFIG, set_runtime_config
 from ghga_connector.constants import C4GH
-from ghga_connector.core import (
-    CLIMessageDisplay,
-    WorkPackageClient,
-    async_client,
-)
-from ghga_connector.core.downloading.api_calls import DownloadClient
-from ghga_connector.core.downloading.batch_processing import FileInfo, FileStager
-from ghga_connector.core.main import (
-    decrypt_file,
-    download_file,
-    get_work_package_token,
-    upload_file,
-)
-
-
-@dataclass
-class WorkPackageInformation:
-    """Wraps decrypted work package token and id to pass to other functions"""
-
-    decrypted_token: str
-    package_id: str
+from ghga_connector.core import CLIMessageDisplay, async_client
+from ghga_connector.core.main import async_download, decrypt_file, upload_file
 
 
 def strtobool(value: str) -> bool:
@@ -86,16 +64,6 @@ def modify_for_debug(debug: bool):
         sys.excepthook = partial(exception_hook)
 
 
-def get_work_package_information(my_private_key: bytes):
-    """Fetch a work package id and work package token and decrypt the token"""
-    # get work package access token and id from user input
-    work_package_id, work_package_token = get_work_package_token(max_tries=3)
-    decrypted_token = crypt.decrypt(data=work_package_token, key=my_private_key)
-    return WorkPackageInformation(
-        decrypted_token=decrypted_token, package_id=work_package_id
-    )
-
-
 cli = typer.Typer(no_args_is_help=True)
 
 
@@ -203,113 +171,6 @@ def download(  # noqa: PLR0913
     )
 
 
-def get_public_key(my_public_key_path: Path) -> bytes:
-    """Get the user's private key from the path supplied"""
-    if not my_public_key_path.is_file():
-        raise exceptions.PubKeyFileDoesNotExistError(public_key_path=my_public_key_path)
-
-    return crypt4gh.keys.get_public_key(filepath=my_public_key_path)
-
-
-def get_private_key(my_private_key_path: Path, passphrase: str | None = None) -> bytes:
-    """Get the user's private key, using the passphrase if supplied/needed."""
-    if passphrase:
-        my_private_key = crypt4gh.keys.get_private_key(
-            filepath=my_private_key_path, callback=lambda: passphrase
-        )
-    else:
-        my_private_key = crypt4gh.keys.get_private_key(
-            filepath=my_private_key_path, callback=None
-        )
-    return my_private_key
-
-
-async def async_download(
-    *,
-    output_dir: Path,
-    my_public_key_path: Path,
-    my_private_key_path: Path,
-    passphrase: str | None = None,
-    overwrite: bool = False,
-):
-    """Download files asynchronously"""
-    if not output_dir.is_dir():
-        raise exceptions.DirectoryDoesNotExistError(directory=output_dir)
-
-    my_public_key = get_public_key(my_public_key_path)
-    my_private_key = get_private_key(my_private_key_path, passphrase)
-
-    CLIMessageDisplay.display("\nFetching work package token...")
-    work_package_information = get_work_package_information(
-        my_private_key=my_private_key
-    )
-
-    async with async_client() as client, set_runtime_config(client=client):
-        CLIMessageDisplay.display("Retrieving API configuration information...")
-        work_package_client = WorkPackageClient(
-            access_token=work_package_information.decrypted_token,
-            client=client,
-            package_id=work_package_information.package_id,
-            my_private_key=my_private_key,
-            my_public_key=my_public_key,
-        )
-
-        file_ids_with_extension = await work_package_client.get_package_files()
-
-        download_client = DownloadClient(
-            client=client, work_package_client=work_package_client
-        )
-
-        CLIMessageDisplay.display("Preparing files for download...")
-        stager = FileStager(
-            wanted_files=file_ids_with_extension,
-            output_dir=output_dir,
-            work_package_client=work_package_client,
-            download_client=download_client,
-            config=CONFIG,
-        )
-        while not stager.finished:
-            staged_files = await stager.get_staged_files()
-            for file_info in staged_files:
-                check_for_existing_file(file_info=file_info, overwrite=overwrite)
-                await download_file(
-                    download_client=download_client,
-                    file_info=file_info,
-                    max_concurrent_downloads=CONFIG.max_concurrent_downloads,
-                    part_size=CONFIG.part_size,
-                )
-                finalize_download(file_info)
-            staged_files.clear()
-
-
-def check_for_existing_file(*, file_info: FileInfo, overwrite: bool):
-    """Check if a file with the given name already exists and conditionally overwrite it."""
-    # check output file
-    output_file = file_info.path_once_complete
-    if output_file.exists():
-        if overwrite:
-            CLIMessageDisplay.display(
-                f"A file with name '{output_file}' already exists and will be overwritten."
-            )
-        else:
-            CLIMessageDisplay.failure(
-                f"A file with name '{output_file}' already exists. Skipping."
-            )
-            return
-
-    output_file_ongoing = file_info.path_during_download
-    if output_file_ongoing.exists():
-        output_file_ongoing.unlink()
-
-
-def finalize_download(file_info: FileInfo):
-    """Rename a file after downloading and announce completion"""
-    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."
-    )
-
-
 @cli.command(no_args_is_help=True)
 def decrypt(  # noqa: PLR0912, C901
     *,

src/ghga_connector/core/downloading/api_calls.py

@@ -25,7 +25,10 @@
 from ghga_connector.config import get_download_api_url
 from ghga_connector.constants import TIMEOUT_LONG
 from ghga_connector.core import RetryHandler
-from ghga_connector.core.api_calls.utils import modify_headers_for_cache_refresh
+from ghga_connector.core.api_calls.utils import (
+    is_service_healthy,
+    modify_headers_for_cache_refresh,
+)
 from ghga_connector.core.work_package import WorkPackageClient
 
 from .structs import RetryResponse
@@ -100,6 +103,14 @@ async def _retrieve_drs_object_using_params(
 
         return _handle_drs_object_response(url=url, response=response)
 
+    def check_download_api_is_reachable(self):
+        """Verify that the download API is reachable.
+
+        Raises an `ApiNotReachableError` if it is not reachable.
+        """
+        if not is_service_healthy(self._download_api_url):
+            raise exceptions.ApiNotReachableError(api_url=self._download_api_url)
+
     async def get_envelope_authorization_headers(
         self, *, file_id: str
     ) -> httpx.Headers:

src/ghga_connector/core/downloading/batch_processing.py

@@ -14,18 +14,19 @@
 # limitations under the License.
 """Module for batch processing related code"""
 
+from collections.abc import AsyncGenerator
 from dataclasses import dataclass, field
 from pathlib import Path
 from time import perf_counter, sleep
 
 from ghga_connector import exceptions
 from ghga_connector.config import Config, get_download_api_url
 from ghga_connector.constants import C4GH
-from ghga_connector.core import CLIMessageDisplay, WorkPackageClient
+from ghga_connector.core import CLIMessageDisplay, WorkPackageClient, utils
 from ghga_connector.core.api_calls import is_service_healthy
 
 from .api_calls import DownloadClient
-from .structs import RetryResponse
+from .structs import FileInfo, RetryResponse
 
 
 @dataclass
@@ -60,36 +61,6 @@ def handle_response(self, *, response: str):
             raise exceptions.AbortBatchProcessError()
 
 
-@dataclass
-class FileInfo:
-    """Information about a file to be downloaded"""
-
-    file_id: str
-    file_extension: str
-    file_size: int
-    output_dir: Path
-
-    @property
-    def file_name(self) -> str:
-        """Construct file name with suffix, if given"""
-        file_name = f"{self.file_id}"
-        if self.file_extension:
-            file_name = f"{self.file_id}{self.file_extension}"
-        return file_name
-
-    @property
-    def path_during_download(self) -> Path:
-        """The file path while the file download is still in progress"""
-        # with_suffix() might overwrite existing suffixes, do this instead:
-        output_file = self.path_once_complete
-        return output_file.parent / (output_file.name + ".part")
-
-    @property
-    def path_once_complete(self) -> Path:
-        """The file path once the download is complete"""
-        return self.output_dir / f"{self.file_name}{C4GH}"
-
-
 class FileStager:
     """Utility class to deal with file staging in batch processing."""
 
@@ -130,7 +101,7 @@ def __init__(
     async def get_staged_files(self) -> list[FileInfo]:
         """Get files that are already staged.
 
-        Returns a dict with file IDs as keys and FileInfo as values.
+        Returns a list of `FileInfo` instances.
         These values contain the download URLs and file sizes.
         The dict should be cleared after these files have been downloaded.
         """
@@ -221,3 +192,31 @@ def _handle_failures(self) -> bool:
         self._started_waiting = perf_counter()  # reset the timer
         self._missing_files = []  # reset list of missing files
         return True
+
+    async def manage_file_downloads(self, overwrite: bool) -> AsyncGenerator[FileInfo]:
+        """Manages file downloads by handling errors, checking for existing files,
+        printing messages to the display, and renaming files after they are downloaded.
+
+        Yields file information.
+        """
+        while not self.finished:
+            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
+                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

@@ -23,18 +23,11 @@
 from pathlib import Path
 
 from ghga_connector import exceptions
-from ghga_connector.core import (
-    CLIMessageDisplay,
-    PartRange,
-    calc_part_ranges,
-)
+from ghga_connector.core import CLIMessageDisplay, PartRange, calc_part_ranges
 from ghga_connector.core.tasks import TaskHandler
 
 from ..progress_bar import DownloadProgressBar
-from .api_calls import (
-    DownloadClient,
-    extract_download_url,
-)
+from .api_calls import DownloadClient, extract_download_url
 from .structs import RetryResponse
 
 logger = logging.getLogger(__name__)

src/ghga_connector/core/downloading/structs.py

@@ -16,10 +16,43 @@
 """Contains additional data structures needed by the download code"""
 
 from dataclasses import dataclass
+from pathlib import Path
+
+from ghga_connector.constants import C4GH
 
 
 @dataclass
 class RetryResponse:
     """Response to download request if file is not yet staged"""
 
     retry_after: int
+
+
+@dataclass
+class FileInfo:
+    """Information about a file to be downloaded"""
+
+    file_id: str
+    file_extension: str
+    file_size: int
+    output_dir: Path
+
+    @property
+    def file_name(self) -> str:
+        """Construct file name with suffix, if given"""
+        file_name = f"{self.file_id}"
+        if self.file_extension:
+            file_name = f"{self.file_id}{self.file_extension}"
+        return file_name
+
+    @property
+    def path_during_download(self) -> Path:
+        """The file path while the file download is still in progress"""
+        # with_suffix() might overwrite existing suffixes, do this instead:
+        output_file = self.path_once_complete
+        return output_file.parent / (output_file.name + ".part")
+
+    @property
+    def path_once_complete(self) -> Path:
+        """The file path once the download is complete"""
+        return self.output_dir / f"{self.file_name}{C4GH}"