Refactor Download Path

Author: TheByronHimes

src/ghga_connector/cli.py

@@ -25,34 +25,27 @@
 from types import TracebackType
 
 import crypt4gh.keys
-import httpx
 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,
-    WorkPackageAccessor,
+    WorkPackageClient,
     async_client,
 )
-from ghga_connector.core.downloading.batch_processing import FileStager
+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_wps_token,
+    get_work_package_token,
     upload_file,
 )
 
 
-@dataclass
-class DownloadParameters:
-    """Contains parameters returned by API calls to prepare information needed for download"""
-
-    file_ids_with_extension: dict[str, str]
-    work_package_accessor: WorkPackageAccessor
-
-
 @dataclass
 class WorkPackageInformation:
     """Wraps decrypted work package token and id to pass to other functions"""
@@ -93,33 +86,10 @@ def modify_for_debug(debug: bool):
         sys.excepthook = partial(exception_hook)
 
 
-async def retrieve_download_parameters(
-    *,
-    client: httpx.AsyncClient,
-    my_private_key: bytes,
-    my_public_key: bytes,
-    work_package_information: WorkPackageInformation,
-) -> DownloadParameters:
-    """Run necessary API calls to configure file download"""
-    work_package_accessor = WorkPackageAccessor(
-        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_accessor.get_package_files()
-
-    return DownloadParameters(
-        file_ids_with_extension=file_ids_with_extension,
-        work_package_accessor=work_package_accessor,
-    )
-
-
 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_wps_token(max_tries=3)
+    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
@@ -276,38 +246,70 @@ async def async_download(
 
     async with async_client() as client, set_runtime_config(client=client):
         CLIMessageDisplay.display("Retrieving API configuration information...")
-        parameters = await retrieve_download_parameters(
+        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,
-            work_package_information=work_package_information,
+        )
+
+        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_file_ids=list(parameters.file_ids_with_extension),
+            wanted_files=file_ids_with_extension,
             output_dir=output_dir,
-            work_package_accessor=parameters.work_package_accessor,
-            client=client,
+            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_id in staged_files:
-                CLIMessageDisplay.display(f"Downloading file with id '{file_id}'...")
+            for file_info in staged_files:
+                check_for_existing_file(file_info=file_info, overwrite=overwrite)
                 await download_file(
-                    client=client,
-                    file_id=file_id,
-                    file_extension=parameters.file_ids_with_extension[file_id],
-                    output_dir=output_dir,
+                    download_client=download_client,
+                    file_info=file_info,
                     max_concurrent_downloads=CONFIG.max_concurrent_downloads,
                     part_size=CONFIG.part_size,
-                    work_package_accessor=parameters.work_package_accessor,
-                    overwrite=overwrite,
                 )
+                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
     *,
@@ -356,7 +358,7 @@ def decrypt(  # noqa: PLR0912, C901
     skipped_files = []
     file_count = 0
     for input_file in input_dir.iterdir():
-        if not input_file.is_file() or input_file.suffix != ".c4gh":
+        if not input_file.is_file() or input_file.suffix != C4GH:
             skipped_files.append(str(input_file))
             continue
 
@@ -396,7 +398,7 @@ def decrypt(  # noqa: PLR0912, C901
 
     if skipped_files:
         CLIMessageDisplay.display(
-            "The following files were skipped as they are not .c4gh files:"
+            f"The following files were skipped as they are not {C4GH} files:"
         )
         for file in skipped_files:
             CLIMessageDisplay.display(f"- {file}")

src/ghga_connector/config.py

@@ -32,16 +32,18 @@
 __all__ = [
     "CONFIG",
     "Config",
-    "get_dcs_api_url",
+    "get_download_api_url",
     "get_ghga_pubkey",
-    "get_ucs_api_url",
-    "get_wps_api_url",
+    "get_upload_api_url",
+    "get_work_package_api_url",
     "set_runtime_config",
 ]
 
-ucs_api_url_var: ContextVar[str] = ContextVar("ucs_api_url", default="")
-dcs_api_url_var: ContextVar[str] = ContextVar("dcs_api_url", default="")
-wps_api_url_var: ContextVar[str] = ContextVar("wps_api_url", default="")
+upload_api_url_var: ContextVar[str] = ContextVar("upload_api_url_var", default="")
+download_api_url_var: ContextVar[str] = ContextVar("download_api_url_var", default="")
+work_package_api_url_var: ContextVar[str] = ContextVar(
+    "work_package_api_url_var", default=""
+)
 ghga_pubkey_var: ContextVar[str] = ContextVar("ghga_pubkey", default="")
 
 
@@ -52,19 +54,19 @@ def _get_context_var(context_var: ContextVar) -> Any:
     return value
 
 
-def get_ucs_api_url() -> str:
-    """Get the UCS API URL."""
-    return _get_context_var(ucs_api_url_var)
+def get_upload_api_url() -> str:
+    """Get the Upload API URL."""
+    return _get_context_var(upload_api_url_var)
 
 
-def get_dcs_api_url() -> str:
-    """Get the DCS API URL."""
-    return _get_context_var(dcs_api_url_var)
+def get_download_api_url() -> str:
+    """Get the Download API URL."""
+    return _get_context_var(download_api_url_var)
 
 
-def get_wps_api_url() -> str:
-    """Get the WPS API URL."""
-    return _get_context_var(wps_api_url_var)
+def get_work_package_api_url() -> str:
+    """Get the Work Package API URL."""
+    return _get_context_var(work_package_api_url_var)
 
 
 def get_ghga_pubkey() -> str:
@@ -112,9 +114,9 @@ async def set_runtime_config(client: httpx.AsyncClient):
 
     This sets the following values:
     - ghga_pubkey
-    - wps_api_url
-    - dcs_api_url
-    - ucs_api_url
+    - work_package_api_url
+    - download_api_url
+    - upload_api_url
 
     Raises:
         WellKnownValueNotFound: If one of the well-known values is not found in the
@@ -135,9 +137,9 @@ async def set_runtime_config(client: httpx.AsyncClient):
 
     async with (
         set_context_var(ghga_pubkey_var, values["crypt4gh_public_key"]),
-        set_context_var(wps_api_url_var, values["wps_api_url"].rstrip("/")),
-        set_context_var(dcs_api_url_var, values["dcs_api_url"].rstrip("/")),
-        set_context_var(ucs_api_url_var, values["ucs_api_url"].rstrip("/")),
+        set_context_var(work_package_api_url_var, values["wps_api_url"].rstrip("/")),
+        set_context_var(download_api_url_var, values["dcs_api_url"].rstrip("/")),
+        set_context_var(upload_api_url_var, values["ucs_api_url"].rstrip("/")),
     ):
         yield
 

src/ghga_connector/constants.py

@@ -23,3 +23,4 @@
 MAX_RETRIES = 5
 MAX_WAIT_TIME = 60 * 60
 CACHE_MIN_FRESH = 3
+C4GH = ".c4gh"

src/ghga_connector/core/__init__.py

@@ -28,4 +28,4 @@
 from .http_translation import ResponseExceptionTranslator  # noqa: F401
 from .message_display import CLIMessageDisplay, MessageColors  # noqa: F401
 from .structs import PartRange  # noqa: F401
-from .work_package import WorkPackageAccessor  # noqa: F401
+from .work_package import WorkPackageClient  # noqa: F401

src/ghga_connector/core/api_calls/utils.py

@@ -42,3 +42,16 @@ def check_url(api_url: str, *, timeout_in_seconds: int = 5) -> bool:
 
     content = response.json()
     return "status" in content and content["status"].lower() == "ok"
+
+
+def modify_headers_for_cache_refresh(headers: httpx.Headers) -> None:
+    """Update cache-control headers to get fresh response from source.
+
+    Headers object is modified in place, hence no return value.
+    """
+    cache_control_headers = headers.get("Cache-Control")
+    if not cache_control_headers:
+        cache_control_headers = ["max-age=0"]
+    else:
+        cache_control_headers = [cache_control_headers, "max-age=0"]
+    headers["Cache-Control"] = ",".join(cache_control_headers)

src/ghga_connector/core/downloading/__init__.py

@@ -15,5 +15,7 @@
 #
 """This subpackage contains functionality needed to download files from GHGA"""
 
-from ..progress_bar import DownloadProgressBar  # noqa: F401
-from .structs import RetryResponse, UrlAndHeaders, URLResponse  # noqa: F401
+from ..progress_bar import DownloadProgressBar
+from .structs import RetryResponse, UrlAndHeaders
+
+__all__ = ["DownloadProgressBar", "RetryResponse", "UrlAndHeaders"]