Move runtime config to contextvars to reduce value relay

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,7 +296,6 @@ 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],

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,58 @@ 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
+    """
+    ghga_pubkey = await _get_wkvs_value(client, value_name="crypt4gh_public_key")
+    wps_api_url = (await _get_wkvs_value(client, value_name="wps_api_url")).rstrip("/")
+    dcs_api_url = (await _get_wkvs_value(client, value_name="dcs_api_url")).rstrip("/")
+    ucs_api_url = (await _get_wkvs_value(client, value_name="ucs_api_url")).rstrip("/")
+
+    async with (
+        set_context_var(ghga_pubkey_var, ghga_pubkey),
+        set_context_var(wps_api_url_var, wps_api_url),
+        set_context_var(dcs_api_url_var, dcs_api_url),
+        set_context_var(ucs_api_url_var, ucs_api_url),
+    ):
+        yield
+
+
+async def _get_wkvs_value(client: httpx.AsyncClient, *, value_name: str) -> Any:
+    """Retrieve a value from the well-known-value-service.
+
+    Args:
+        value_name (str): the name of the value to be retrieved
+
+    Raises:
+        WellKnownValueNotFound: when a 404 response is received from the WKVS
+        KeyError: when a successful response is received but doesn't contain the expected value
+    """
+    url = f"{CONFIG.wkvs_api_url}/values/{value_name}"
+
+    try:
+        response = await client.get(url)  # verify is True by default
+    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 == 404:
+        raise exceptions.WellKnownValueNotFound(value_name=value_name)
+
+    try:
+        value = response.json()[value_name]
+    except KeyError as err:
+        raise KeyError(
+            "Response from well-known-value-service did not include expected field"
+            + f" '{value_name}'"
+        ) from err
+    return value

src/ghga_connector/core/api_calls/well_knowns.py

@@ -20,7 +20,7 @@
 
 import httpx
 
-from ghga_connector.core import exceptions
+from ghga_connector import exceptions
 
 
 @dataclass

src/ghga_connector/core/crypt/encryption.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/downloading/api_calls.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,

src/ghga_connector/core/downloading/batch_processing.py

@@ -20,11 +20,11 @@
 
 import httpx
 
-from ghga_connector.config import Config
+from ghga_connector import exceptions
+from ghga_connector.config import Config, get_dcs_api_url
 from ghga_connector.core import (
     CLIMessageDisplay,
     WorkPackageAccessor,
-    exceptions,
 )
 from ghga_connector.core.api_calls import is_service_healthy
 
@@ -70,11 +70,10 @@ def handle_response(self, *, response: str):
 class FileStager:
     """Utility class to deal with file staging in batch processing."""
 
-    def __init__(  # noqa: PLR0913
+    def __init__(
         self,
         *,
         wanted_file_ids: list[str],
-        dcs_api_url: str,
         output_dir: Path,
         work_package_accessor: WorkPackageAccessor,
         client: httpx.AsyncClient,
@@ -83,9 +82,9 @@ def __init__(  # noqa: PLR0913
         """Initialize the FileStager."""
         self.io_handler = CliIoHandler()
         existing_file_ids = set(self.io_handler.check_output(location=output_dir))
-        if not is_service_healthy(dcs_api_url):
-            raise exceptions.ApiNotReachableError(api_url=dcs_api_url)
-        self.api_url = dcs_api_url
+        self.api_url = get_dcs_api_url()
+        if not is_service_healthy(self.api_url):
+            raise exceptions.ApiNotReachableError(api_url=self.api_url)
         self.work_package_accessor = work_package_accessor
         self.max_wait_time = config.max_wait_time
         self.client = client

src/ghga_connector/core/downloading/downloader.py

@@ -26,14 +26,14 @@
 import httpx
 from tenacity import RetryError
 
+from ghga_connector import exceptions
 from ghga_connector.core import (
     CLIMessageDisplay,
     PartRange,
     ResponseExceptionTranslator,
     RetryHandler,
     WorkPackageAccessor,
     calc_part_ranges,
-    exceptions,
 )
 from ghga_connector.core.tasks import TaskHandler