Refactor config (#117)

* Remove unused max_wait_time parameter

* Move runtime config to contextvars to reduce value relay

* Delete WKVSCaller module and update code

* Fix CLI tests by changing config patch location

* Patch better

* Consolidate WKVS calls

* Apply test config patch in api_calls tests

* Remove trailing slash in wkvs route

Author: TheByronHimes

src/ghga_connector/cli.py

@@ -29,14 +29,13 @@
 import typer
 from ghga_service_commons.utils import crypt
 
-from ghga_connector.config import CONFIG
+from ghga_connector import exceptions
+from ghga_connector.config import CONFIG, set_runtime_config
 from ghga_connector.core import (
     CLIMessageDisplay,
     WorkPackageAccessor,
     async_client,
-    exceptions,
 )
-from ghga_connector.core.api_calls import WKVSCaller
 from ghga_connector.core.downloading.batch_processing import FileStager
 from ghga_connector.core.main import (
     decrypt_file,
@@ -50,19 +49,10 @@
 class DownloadParameters:
     """Contains parameters returned by API calls to prepare information needed for download"""
 
-    dcs_api_url: str
     file_ids_with_extension: dict[str, str]
     work_package_accessor: WorkPackageAccessor
 
 
-@dataclass
-class UploadParameters:
-    """Contains parameters returned by API calls to prepare information needed for upload"""
-
-    ucs_api_url: str
-    server_pubkey: str
-
-
 @dataclass
 class WorkPackageInformation:
     """Wraps decrypted work package token and id to pass to other functions"""
@@ -103,15 +93,6 @@ def modify_for_debug(debug: bool):
         sys.excepthook = partial(exception_hook)
 
 
-async def retrieve_upload_parameters(client: httpx.AsyncClient) -> UploadParameters:
-    """Configure httpx client and retrieve necessary parameters from WKVS"""
-    wkvs_caller = WKVSCaller(client=client, wkvs_url=CONFIG.wkvs_api_url)
-    ucs_api_url = await wkvs_caller.get_ucs_api_url()
-    server_pubkey = await wkvs_caller.get_server_pubkey()
-
-    return UploadParameters(server_pubkey=server_pubkey, ucs_api_url=ucs_api_url)
-
-
 async def retrieve_download_parameters(
     *,
     client: httpx.AsyncClient,
@@ -120,23 +101,16 @@ async def retrieve_download_parameters(
     work_package_information: WorkPackageInformation,
 ) -> DownloadParameters:
     """Run necessary API calls to configure file download"""
-    wkvs_caller = WKVSCaller(client=client, wkvs_url=CONFIG.wkvs_api_url)
-    dcs_api_url = await wkvs_caller.get_dcs_api_url()
-    wps_api_url = await wkvs_caller.get_wps_api_url()
-
     work_package_accessor = WorkPackageAccessor(
         access_token=work_package_information.decrypted_token,
-        api_url=wps_api_url,
         client=client,
-        dcs_api_url=dcs_api_url,
         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(
-        dcs_api_url=dcs_api_url,
         file_ids_with_extension=file_ids_with_extension,
         work_package_accessor=work_package_accessor,
     )
@@ -199,14 +173,11 @@ async def async_upload(
     passphrase: str | None = None,
 ):
     """Upload a file asynchronously"""
-    async with async_client() as client:
-        parameters = await retrieve_upload_parameters(client)
+    async with async_client() as client, set_runtime_config(client=client):
         await upload_file(
-            api_url=parameters.ucs_api_url,
             client=client,
             file_id=file_id,
             file_path=file_path,
-            server_public_key=parameters.server_pubkey,
             my_public_key_path=my_public_key_path,
             my_private_key_path=my_private_key_path,
             passphrase=passphrase,
@@ -303,7 +274,7 @@ async def async_download(
         my_private_key=my_private_key
     )
 
-    async with async_client() as client:
+    async with async_client() as client, set_runtime_config(client=client):
         CLIMessageDisplay.display("Retrieving API configuration information...")
         parameters = await retrieve_download_parameters(
             client=client,
@@ -315,7 +286,6 @@ async def async_download(
         CLIMessageDisplay.display("Preparing files for download...")
         stager = FileStager(
             wanted_file_ids=list(parameters.file_ids_with_extension),
-            dcs_api_url=parameters.dcs_api_url,
             output_dir=output_dir,
             work_package_accessor=parameters.work_package_accessor,
             client=client,
@@ -326,13 +296,11 @@ async def async_download(
             for file_id in staged_files:
                 CLIMessageDisplay.display(f"Downloading file with id '{file_id}'...")
                 await download_file(
-                    api_url=parameters.dcs_api_url,
                     client=client,
                     file_id=file_id,
                     file_extension=parameters.file_ids_with_extension[file_id],
                     output_dir=output_dir,
                     max_concurrent_downloads=CONFIG.max_concurrent_downloads,
-                    max_wait_time=CONFIG.max_wait_time,
                     part_size=CONFIG.part_size,
                     work_package_accessor=parameters.work_package_accessor,
                     overwrite=overwrite,

src/ghga_connector/config.py

@@ -16,12 +16,61 @@
 
 """Global Config Parameters"""
 
+from contextlib import asynccontextmanager
+from contextvars import ContextVar
+from typing import Any
+
+import httpx
 from hexkit.config import config_from_yaml
+from hexkit.utils import set_context_var
 from pydantic import Field, NonNegativeInt, PositiveInt
 from pydantic_settings import BaseSettings
 
+from ghga_connector import exceptions
 from ghga_connector.constants import DEFAULT_PART_SIZE, MAX_RETRIES, MAX_WAIT_TIME
 
+__all__ = [
+    "CONFIG",
+    "Config",
+    "get_dcs_api_url",
+    "get_ghga_pubkey",
+    "get_ucs_api_url",
+    "get_wps_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="")
+ghga_pubkey_var: ContextVar[str] = ContextVar("ghga_pubkey", default="")
+
+
+def _get_context_var(context_var: ContextVar) -> Any:
+    value = context_var.get()
+    if not value:
+        raise ValueError(f"{context_var.name} is not set")
+    return value
+
+
+def get_ucs_api_url() -> str:
+    """Get the UCS API URL."""
+    return _get_context_var(ucs_api_url_var)
+
+
+def get_dcs_api_url() -> str:
+    """Get the DCS API URL."""
+    return _get_context_var(dcs_api_url_var)
+
+
+def get_wps_api_url() -> str:
+    """Get the WPS API URL."""
+    return _get_context_var(wps_api_url_var)
+
+
+def get_ghga_pubkey() -> str:
+    """Get the GHGA crypt4gh public key."""
+    return _get_context_var(ghga_pubkey_var)
+
 
 @config_from_yaml(prefix="ghga_connector")
 class Config(BaseSettings):
@@ -55,3 +104,63 @@ class Config(BaseSettings):
 
 
 CONFIG = Config()
+
+
+@asynccontextmanager
+async def set_runtime_config(client: httpx.AsyncClient):
+    """Set runtime config as context vars to be accessed within a context manager.
+
+    This sets the following values:
+    - ghga_pubkey
+    - wps_api_url
+    - dcs_api_url
+    - ucs_api_url
+
+    Raises:
+        WellKnownValueNotFound: If one of the well-known values is not found in the
+            response from the WKVS.
+        BadResponseCodeError: If the status code returned is not 200
+        ConnectionFailedError: If the request fails due to a timeout/connection problem
+        RequestFailedError: If the request fails for any other reason
+    """
+    values = await _get_wkvs_values(client)
+    for value_name in [
+        "crypt4gh_public_key",
+        "wps_api_url",
+        "dcs_api_url",
+        "ucs_api_url",
+    ]:
+        if value_name not in values:
+            raise exceptions.WellKnownValueNotFound(value_name=value_name)
+
+    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("/")),
+    ):
+        yield
+
+
+async def _get_wkvs_values(client: httpx.AsyncClient) -> dict[str, Any]:
+    """Retrieve a value from the well-known-value-service using the supplied client.
+
+    Raises:
+        BadResponseCodeError: If the status code returned is not 200
+        ConnectionFailedError: If the request fails due to a timeout/connection problem
+        RequestFailedError: If the request fails for any other reason
+    """
+    url = f"{CONFIG.wkvs_api_url}/values"
+
+    try:
+        response = await client.get(url)
+    except httpx.RequestError as request_error:
+        exceptions.raise_if_connection_failed(request_error=request_error, url=url)
+        raise exceptions.RequestFailedError(url=url) from request_error
+
+    if response.status_code != 200:
+        raise exceptions.BadResponseCodeError(
+            url=url, response_code=response.status_code
+        )
+
+    return response.json()

src/ghga_connector/core/api_calls/__init__.py

@@ -18,4 +18,3 @@
 """
 
 from .utils import check_url, is_service_healthy  # noqa: F401
-from .well_knowns import WKVSCaller  # noqa: F401

src/ghga_connector/core/api_calls/well_knowns.py

@@ -26,6 +26,7 @@
 import crypt4gh.lib
 from nacl.bindings import crypto_aead_chacha20poly1305_ietf_encrypt
 
+from ghga_connector.config import get_ghga_pubkey
 from ghga_connector.core import get_segments, read_file_parts
 
 from .abstract_bases import Encryptor
@@ -35,11 +36,10 @@
 class Crypt4GHEncryptor(Encryptor):
     """Handles on the fly encryption and checksum calculation"""
 
-    def __init__(  # noqa: PLR0913
+    def __init__(
         self,
         part_size: int,
         private_key_path: Path,
-        server_public_key: str,
         passphrase: str | None,
         checksums: Checksums = Checksums(),
         file_secret: bytes | None = None,
@@ -48,7 +48,7 @@ def __init__(  # noqa: PLR0913
         self._checksums = checksums
         self._part_size = part_size
         self._private_key_path = private_key_path
-        self._server_public_key = base64.b64decode(server_public_key)
+        self._server_public_key = base64.b64decode(get_ghga_pubkey())
         self._passphrase = passphrase
         if file_secret is None:
             file_secret = os.urandom(32)

src/ghga_connector/core/crypt/encryption.py

@@ -18,8 +18,10 @@
 import httpx
 from tenacity import RetryError
 
+from ghga_connector import exceptions
+from ghga_connector.config import get_dcs_api_url
 from ghga_connector.constants import CACHE_MIN_FRESH, TIMEOUT_LONG
-from ghga_connector.core import RetryHandler, WorkPackageAccessor, exceptions
+from ghga_connector.core import RetryHandler, WorkPackageAccessor
 
 from .structs import (
     RetryResponse,
@@ -63,7 +65,8 @@ async def get_envelope_authorization(
     a Crypt4GH envelope for file identified by `file_id`
     """
     # build url
-    url = f"{work_package_accessor.dcs_api_url}/objects/{file_id}/envelopes"
+    dcs_api_url = get_dcs_api_url()
+    url = f"{dcs_api_url}/objects/{file_id}/envelopes"
     headers = await _get_authorization(
         file_id=file_id, work_package_accessor=work_package_accessor
     )
@@ -78,7 +81,8 @@ async def get_file_authorization(
     object storage URL for file download
     """
     # build URL
-    url = f"{work_package_accessor.dcs_api_url}/objects/{file_id}"
+    dcs_api_url = get_dcs_api_url()
+    url = f"{dcs_api_url}/objects/{file_id}"
     headers = await _get_authorization(
         file_id=file_id,
         work_package_accessor=work_package_accessor,