Big refactor: flatten CLI (#119)

* Flatten CLI and further refactor top of download path

* Fix test problems

* Decouple download error handling again

Author: TheByronHimes

src/ghga_connector/cli.py

@@ -16,85 +16,17 @@
 """CLI-specific wrappers around core functions."""
 
 import asyncio
-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
-
-
-def strtobool(value: str) -> bool:
-    """Inplace replacement for distutils.utils"""
-    return value.lower() in ("y", "yes", "on", "1", "true", "t")
-
-
-def exception_hook(
-    type_: BaseException,
-    value: BaseException,
-    traceback: TracebackType | None,
-):
-    """When debug mode is NOT enabled, gets called to perform final error handling
-    before program exits
-    """
-    message = (
-        "An error occurred. Rerun command"
-        + " with --debug at the end to see more information."
-    )
-
-    if value.args:
-        message += f"\n{value.args[0]}"
-
-    CLIMessageDisplay.failure(message)
-
-
-def modify_for_debug(debug: bool):
-    """Enable debug logging and configure exception printing if debug=True"""
-    if debug:
-        # enable debug logging
-        logging.basicConfig(level=logging.DEBUG)
-        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
-    )
-
+from ghga_connector.core import CLIMessageDisplay, async_client
+from ghga_connector.core.main import async_download, decrypt_file, upload_file
+from ghga_connector.core.utils import modify_for_debug, strtobool
 
 cli = typer.Typer(no_args_is_help=True)
 
@@ -203,113 +135,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,20 @@ 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 checking for existing files before download,
+        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)
+                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."
+                )
+            staged_files.clear()

src/ghga_connector/core/downloading/downloader.py

@@ -19,27 +19,39 @@
 import gc
 import logging
 from asyncio import PriorityQueue, Queue, Semaphore, Task
+from contextlib import contextmanager
 from io import BufferedWriter
 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 .structs import RetryResponse
+from .api_calls import DownloadClient, extract_download_url
+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.