cosmotech-acceleration-library 1.0.1__py3-none-any.whl → 2.0.0__py3-none-any.whl

cosmotech/coal/azure/storage.py

@@ -15,9 +15,9 @@ uploading files to blob storage.
 import pathlib
 
 from azure.storage.blob import ContainerClient
+from cosmotech.orchestrator.utils.translate import T
 
 from cosmotech.coal.utils.logger import LOGGER
-from cosmotech.orchestrator.utils.translate import T
 
 
 def upload_file(
@@ -65,10 +65,10 @@ def upload_folder(
         raise FileNotFoundError(T("coal.common.file_operations.not_found").format(source_folder=source_folder))
 
     if source_path.is_dir():
-        _source_name = str(source_path)
+        # _source_name = str(source_path)
         for _file_path in source_path.glob("**/*" if recursive else "*"):
             if _file_path.is_file():
-                _file_name = str(_file_path).removeprefix(_source_name).removeprefix("/")
+                # _file_name = str(_file_path).removeprefix(_source_name).removeprefix("/")
                 upload_file(_file_path, blob_name, az_storage_sas_url, file_prefix)
     else:
         upload_file(source_path, blob_name, az_storage_sas_url, file_prefix)

cosmotech/coal/cosmotech_api/__init__.py

@@ -10,27 +10,3 @@ Cosmotech API integration module.
 
 This module provides functions for interacting with the Cosmotech API.
 """
-
-# Re-export functions from the parameters module
-from cosmotech.coal.cosmotech_api.parameters import (
-    write_parameters,
-)
-
-# Re-export functions from the twin_data_layer module
-from cosmotech.coal.cosmotech_api.twin_data_layer import (
-    get_dataset_id_from_runner,
-    send_files_to_tdl,
-    load_files_from_tdl,
-)
-
-# Re-export functions from the run_data module
-from cosmotech.coal.cosmotech_api.run_data import (
-    send_csv_to_run_data,
-    send_store_to_run_data,
-    load_csv_from_run_data,
-)
-
-# Re-export functions from the run_template module
-from cosmotech.coal.cosmotech_api.run_template import (
-    load_run_template_handlers,
-)

cosmotech/coal/cosmotech_api/apis/__init__.py

@@ -0,0 +1,14 @@
+# Copyright (C) - 2023 - 2025 - Cosmo Tech
+# This document and all information contained herein is the exclusive property -
+# including all intellectual property rights pertaining thereto - of Cosmo Tech.
+# Any use, reproduction, translation, broadcasting, transmission, distribution,
+# etc., to any person is prohibited unless it has been previously and
+# specifically authorized by written means by Cosmo Tech.
+
+from cosmotech.coal.cosmotech_api.apis.dataset import DatasetApi
+from cosmotech.coal.cosmotech_api.apis.meta import MetaApi
+from cosmotech.coal.cosmotech_api.apis.organization import OrganizationApi
+from cosmotech.coal.cosmotech_api.apis.run import RunApi
+from cosmotech.coal.cosmotech_api.apis.runner import RunnerApi
+from cosmotech.coal.cosmotech_api.apis.solution import SolutionApi
+from cosmotech.coal.cosmotech_api.apis.workspace import WorkspaceApi

cosmotech/coal/cosmotech_api/apis/dataset.py

@@ -0,0 +1,103 @@
+# Copyright (C) - 2023 - 2025 - Cosmo Tech
+# This document and all information contained herein is the exclusive property -
+# including all intellectual property rights pertaining thereto - of Cosmo Tech.
+# Any use, reproduction, translation, broadcasting, transmission, distribution,
+# etc., to any person is prohibited unless it has been previously and
+# specifically authorized by written means by Cosmo Tech.
+from pathlib import Path
+from typing import Optional, Union
+
+from cosmotech.orchestrator.utils.translate import T
+from cosmotech_api import (
+    Dataset,
+)
+from cosmotech_api import DatasetApi as BaseDatasetApi
+from cosmotech_api import (
+    DatasetCreateRequest,
+    DatasetPartCreateRequest,
+    DatasetPartTypeEnum,
+)
+
+from cosmotech.coal.cosmotech_api.objects.connection import Connection
+from cosmotech.coal.utils.configuration import ENVIRONMENT_CONFIGURATION, Configuration
+from cosmotech.coal.utils.logger import LOGGER
+
+
+class DatasetApi(BaseDatasetApi, Connection):
+
+    def __init__(
+        self,
+        configuration: Configuration = ENVIRONMENT_CONFIGURATION,
+    ):
+        Connection.__init__(self, configuration)
+        BaseDatasetApi.__init__(self, self.api_client)
+
+        LOGGER.debug(T("coal.cosmotech_api.initialization.dataset_api_initialized"))
+
+    def download_dataset(self, dataset_id) -> Dataset:
+        dataset = self.get_dataset(
+            organization_id=self.configuration.cosmotech.organization_id,
+            workspace_id=self.configuration.cosmotech.workspace_id,
+            dataset_id=dataset_id,
+        )
+
+        dataset_dir = self.configuration.cosmotech.dataset_absolute_path
+        dataset_dir_path = Path(dataset_dir) / dataset_id
+        for part in dataset.parts:
+            part_file_path = dataset_dir_path / part.source_name
+            part_file_path.parent.mkdir(parents=True, exist_ok=True)
+            data_part = self.download_dataset_part(
+                organization_id=self.configuration.cosmotech.organization_id,
+                workspace_id=self.configuration.cosmotech.workspace_id,
+                dataset_id=dataset_id,
+                dataset_part_id=part.id,
+            )
+            with open(part_file_path, "wb") as binary_file:
+                binary_file.write(data_part)
+            LOGGER.debug(
+                T("coal.services.dataset.part_downloaded").format(part_name=part.source_name, file_path=part_file_path)
+            )
+        return dataset
+
+    @staticmethod
+    def path_to_parts(_path, part_type) -> list[tuple[str, Path, DatasetPartTypeEnum]]:
+        if (_path := Path(_path)).is_dir():
+            return list((str(_p.relative_to(_path)), _p, part_type) for _p in _path.rglob("*") if _p.is_file())
+        return list(((_path.name, _path, part_type),))
+
+    def upload_dataset(
+        self,
+        dataset_name: str,
+        as_files: Optional[list[Union[Path, str]]] = (),
+        as_db: Optional[list[Union[Path, str]]] = (),
+    ) -> Dataset:
+        _parts = list()
+
+        for _f in as_files:
+            _parts.extend(self.path_to_parts(_f, DatasetPartTypeEnum.FILE))
+
+        for _db in as_db:
+            _parts.extend(self.path_to_parts(_db, DatasetPartTypeEnum.DB))
+
+        d_request = DatasetCreateRequest(
+            name=dataset_name,
+            parts=list(
+                DatasetPartCreateRequest(
+                    name=_p_name,
+                    description=_p_name,
+                    sourceName=_p_name,
+                    type=_type,
+                )
+                for _p_name, _, _type in _parts
+            ),
+        )
+
+        d_ret = self.create_dataset(
+            self.configuration.cosmotech.organization_id,
+            self.configuration.cosmotech.workspace_id,
+            d_request,
+            files=list((_p[0], _p[1].open("rb").read()) for _p in _parts),
+        )
+
+        LOGGER.info(T("coal.services.dataset.dataset_created").format(dataset_id=d_ret.id))
+        return d_ret

cosmotech/coal/cosmotech_api/apis/meta.py

@@ -0,0 +1,25 @@
+# Copyright (C) - 2023 - 2025 - Cosmo Tech
+# This document and all information contained herein is the exclusive property -
+# including all intellectual property rights pertaining thereto - of Cosmo Tech.
+# Any use, reproduction, translation, broadcasting, transmission, distribution,
+# etc., to any person is prohibited unless it has been previously and
+# specifically authorized by written means by Cosmo Tech.
+from cosmotech.orchestrator.utils.translate import T
+from cosmotech_api import MetaApi as BaseMetaApi
+
+from cosmotech.coal.cosmotech_api.objects.connection import Connection
+from cosmotech.coal.utils.configuration import ENVIRONMENT_CONFIGURATION, Configuration
+from cosmotech.coal.utils.logger import LOGGER
+
+
+class MetaApi(BaseMetaApi, Connection):
+
+    def __init__(
+        self,
+        configuration: Configuration = ENVIRONMENT_CONFIGURATION,
+    ):
+        Connection.__init__(self, configuration)
+
+        BaseMetaApi.__init__(self, self.api_client)
+
+        LOGGER.debug(T("coal.cosmotech_api.initialization.meta_api_initialized"))

cosmotech/coal/cosmotech_api/apis/organization.py

@@ -0,0 +1,24 @@
+# Copyright (C) - 2023 - 2025 - Cosmo Tech
+# This document and all information contained herein is the exclusive property -
+# including all intellectual property rights pertaining thereto - of Cosmo Tech.
+# Any use, reproduction, translation, broadcasting, transmission, distribution,
+# etc., to any person is prohibited unless it has been previously and
+# specifically authorized by written means by Cosmo Tech.
+from cosmotech.orchestrator.utils.translate import T
+from cosmotech_api import OrganizationApi as BaseOrganizationApi
+
+from cosmotech.coal.cosmotech_api.objects.connection import Connection
+from cosmotech.coal.utils.configuration import ENVIRONMENT_CONFIGURATION, Configuration
+from cosmotech.coal.utils.logger import LOGGER
+
+
+class OrganizationApi(BaseOrganizationApi, Connection):
+
+    def __init__(
+        self,
+        configuration: Configuration = ENVIRONMENT_CONFIGURATION,
+    ):
+        Connection.__init__(self, configuration)
+        BaseOrganizationApi.__init__(self, self.api_client)
+
+        LOGGER.debug(T("coal.cosmotech_api.initialization.organization_api_initialized"))

cosmotech/coal/cosmotech_api/apis/run.py

@@ -0,0 +1,38 @@
+# Copyright (C) - 2023 - 2025 - Cosmo Tech
+# This document and all information contained herein is the exclusive property -
+# including all intellectual property rights pertaining thereto - of Cosmo Tech.
+# Any use, reproduction, translation, broadcasting, transmission, distribution,
+# etc., to any person is prohibited unless it has been previously and
+# specifically authorized by written means by Cosmo Tech.
+from typing import Any, Optional
+
+from cosmotech.orchestrator.utils.translate import T
+from cosmotech_api import RunApi as BaseRunApi
+
+from cosmotech.coal.cosmotech_api.objects.connection import Connection
+from cosmotech.coal.utils.configuration import ENVIRONMENT_CONFIGURATION, Configuration
+from cosmotech.coal.utils.logger import LOGGER
+
+
+class RunApi(BaseRunApi, Connection):
+
+    def __init__(
+        self,
+        configuration: Configuration = ENVIRONMENT_CONFIGURATION,
+    ):
+        Connection.__init__(self, configuration)
+        BaseRunApi.__init__(self, self.api_client)
+
+        LOGGER.debug(T("coal.cosmotech_api.initialization.run_api_initialized"))
+
+    def get_run_metadata(
+        self,
+        organization_id: str,
+        workspace_id: str,
+        runner_id: str,
+        run_id: str,
+        include: Optional[list[str]] = None,
+        exclude: Optional[list[str]] = None,
+    ) -> dict[str, Any]:
+        run = self.get_run(organization_id, workspace_id, runner_id, run_id)
+        return run.model_dump(by_alias=True, exclude_none=True, include=include, exclude=exclude, mode="json")

cosmotech/coal/cosmotech_api/apis/runner.py

@@ -0,0 +1,71 @@
+# Copyright (C) - 2023 - 2025 - Cosmo Tech
+# This document and all information contained herein is the exclusive property -
+# including all intellectual property rights pertaining thereto - of Cosmo Tech.
+# Any use, reproduction, translation, broadcasting, transmission, distribution,
+# etc., to any person is prohibited unless it has been previously and
+# specifically authorized by written means by Cosmo Tech.
+from typing import Any, Optional
+
+from cosmotech.orchestrator.utils.translate import T
+from cosmotech_api import RunnerApi as BaseRunnerApi
+
+from cosmotech.coal.cosmotech_api.apis.dataset import DatasetApi
+from cosmotech.coal.cosmotech_api.objects.connection import Connection
+from cosmotech.coal.cosmotech_api.objects.parameters import Parameters
+from cosmotech.coal.utils.configuration import ENVIRONMENT_CONFIGURATION, Configuration
+from cosmotech.coal.utils.logger import LOGGER
+
+
+class RunnerApi(BaseRunnerApi, Connection):
+
+    def __init__(
+        self,
+        configuration: Configuration = ENVIRONMENT_CONFIGURATION,
+    ):
+        Connection.__init__(self, configuration)
+        BaseRunnerApi.__init__(self, self.api_client)
+
+        LOGGER.debug(T("coal.cosmotech_api.initialization.runner_api_initialized"))
+
+    def get_runner_metadata(
+        self,
+        organization_id: str,
+        workspace_id: str,
+        runner_id: str,
+        include: Optional[list[str]] = None,
+        exclude: Optional[list[str]] = None,
+    ) -> dict[str, Any]:
+        runner = self.get_runner(organization_id, workspace_id, runner_id)
+
+        return runner.model_dump(by_alias=True, exclude_none=True, include=include, exclude=exclude, mode="json")
+
+    def download_runner_data(
+        self,
+        organization_id: str,
+        workspace_id: str,
+        runner_id: str,
+        parameter_folder: str,
+        dataset_folder: Optional[str] = None,
+    ):
+        LOGGER.info(T("coal.cosmotech_api.runner.starting_download"))
+
+        # Get runner data
+        runner_data = self.get_runner(organization_id, workspace_id, runner_id)
+
+        # Skip if no parameters found
+        if not runner_data.parameters_values:
+            LOGGER.warning(T("coal.cosmotech_api.runner.no_parameters"))
+        else:
+            LOGGER.info(T("coal.cosmotech_api.runner.loaded_data"))
+            parameters = Parameters(runner_data)
+            parameters.write_parameters_to_json(parameter_folder)
+
+        # Download datasets if requested
+        if dataset_folder:
+            datasets_ids = runner_data.datasets.bases
+
+            if datasets_ids:
+                LOGGER.info(T("coal.cosmotech_api.runner.downloading_datasets").format(count=len(datasets_ids)))
+                ds_api = DatasetApi(self.configuration)
+                for dataset_id in datasets_ids:
+                    ds_api.download_dataset(dataset_id)

cosmotech/coal/cosmotech_api/apis/solution.py

@@ -0,0 +1,23 @@
+# Copyright (C) - 2023 - 2025 - Cosmo Tech
+# This document and all information contained herein is the exclusive property -
+# including all intellectual property rights pertaining thereto - of Cosmo Tech.
+# Any use, reproduction, translation, broadcasting, transmission, distribution,
+# etc., to any person is prohibited unless it has been previously and
+# specifically authorized by written means by Cosmo Tech.
+from cosmotech.orchestrator.utils.translate import T
+from cosmotech_api import SolutionApi as BaseSolutionApi
+
+from cosmotech.coal.cosmotech_api.objects.connection import Connection
+from cosmotech.coal.utils.configuration import ENVIRONMENT_CONFIGURATION, Configuration
+from cosmotech.coal.utils.logger import LOGGER
+
+
+class SolutionApi(BaseSolutionApi, Connection):
+    def __init__(
+        self,
+        configuration: Configuration = ENVIRONMENT_CONFIGURATION,
+    ):
+        Connection.__init__(self, configuration)
+        BaseSolutionApi.__init__(self, self.api_client)
+
+        LOGGER.debug(T("coal.cosmotech_api.initialization.solution_api_initialized"))

cosmotech/coal/cosmotech_api/apis/workspace.py

@@ -0,0 +1,108 @@
+# Copyright (C) - 2023 - 2025 - Cosmo Tech
+# This document and all information contained herein is the exclusive property -
+# including all intellectual property rights pertaining thereto - of Cosmo Tech.
+# Any use, reproduction, translation, broadcasting, transmission, distribution,
+# etc., to any person is prohibited unless it has been previously and
+# specifically authorized by written means by Cosmo Tech.
+from pathlib import Path
+
+from cosmotech.orchestrator.utils.translate import T
+from cosmotech_api import ApiException
+from cosmotech_api import WorkspaceApi as BaseWorkspaceApi
+
+from cosmotech.coal.cosmotech_api.objects.connection import Connection
+from cosmotech.coal.utils.configuration import ENVIRONMENT_CONFIGURATION, Configuration
+from cosmotech.coal.utils.logger import LOGGER
+
+
+class WorkspaceApi(BaseWorkspaceApi, Connection):
+
+    def __init__(
+        self,
+        configuration: Configuration = ENVIRONMENT_CONFIGURATION,
+    ):
+        Connection.__init__(self, configuration)
+        BaseWorkspaceApi.__init__(self, self.api_client)
+
+        LOGGER.debug(T("coal.cosmotech_api.initialization.workspace_api_initialized"))
+
+    def list_filtered_workspace_files(
+        self,
+        organization_id: str,
+        workspace_id: str,
+        file_prefix: str,
+    ) -> list[str]:
+        target_list = []
+        LOGGER.info(T("coal.cosmotech_api.workspace.target_is_folder"))
+        wsf = self.list_workspace_files(organization_id, workspace_id)
+        for workspace_file in wsf:
+            if workspace_file.file_name.startswith(file_prefix):
+                target_list.append(workspace_file.file_name)
+
+        if not target_list:
+            LOGGER.error(
+                T("coal.common.errors.data_no_workspace_files").format(
+                    file_prefix=file_prefix, workspace_id=workspace_id
+                )
+            )
+            raise ValueError(
+                T("coal.common.errors.data_no_workspace_files").format(
+                    file_prefix=file_prefix, workspace_id=workspace_id
+                )
+            )
+
+        return target_list
+
+    def download_workspace_file(
+        self,
+        organization_id: str,
+        workspace_id: str,
+        file_name: str,
+        target_dir: Path,
+    ) -> Path:
+        if target_dir.is_file():
+            raise ValueError(T("coal.common.file_operations.not_directory").format(target_dir=target_dir))
+
+        LOGGER.info(T("coal.cosmotech_api.workspace.loading_file").format(file_name=file_name))
+
+        _file_content = self.get_workspace_file(organization_id, workspace_id, file_name)
+
+        local_target_file = target_dir / file_name
+        local_target_file.parent.mkdir(parents=True, exist_ok=True)
+
+        with open(local_target_file, "wb") as _file:
+            _file.write(_file_content)
+
+        LOGGER.info(T("coal.cosmotech_api.workspace.file_loaded").format(file=local_target_file))
+
+        return local_target_file
+
+    def upload_workspace_file(
+        self,
+        organization_id: str,
+        workspace_id: str,
+        file_path: str,
+        workspace_path: str,
+        overwrite: bool = True,
+    ) -> str:
+        target_file = Path(file_path)
+        if not target_file.exists():
+            LOGGER.error(T("coal.common.file_operations.not_exists").format(file_path=file_path))
+            raise ValueError(T("coal.common.file_operations.not_exists").format(file_path=file_path))
+        if not target_file.is_file():
+            LOGGER.error(T("coal.common.file_operations.not_single_file").format(file_path=file_path))
+            raise ValueError(T("coal.common.file_operations.not_single_file").format(file_path=file_path))
+
+        destination = workspace_path + target_file.name if workspace_path.endswith("/") else workspace_path
+
+        LOGGER.info(T("coal.cosmotech_api.workspace.sending_to_api").format(destination=destination))
+        try:
+            _file = self.create_workspace_file(
+                organization_id, workspace_id, file_path, overwrite, destination=destination
+            )
+        except ApiException as e:
+            LOGGER.error(T("coal.common.file_operations.already_exists").format(csv_path=destination))
+            raise e
+
+        LOGGER.info(T("coal.cosmotech_api.workspace.file_sent").format(file=_file.file_name))
+        return _file.file_name

cosmotech/coal/cosmotech_api/objects/__init__.py

@@ -0,0 +1,9 @@
+# Copyright (C) - 2023 - 2025 - Cosmo Tech
+# This document and all information contained herein is the exclusive property -
+# including all intellectual property rights pertaining thereto - of Cosmo Tech.
+# Any use, reproduction, translation, broadcasting, transmission, distribution,
+# etc., to any person is prohibited unless it has been previously and
+# specifically authorized by written means by Cosmo Tech.
+
+from cosmotech.coal.cosmotech_api.objects.connection import Connection
+from cosmotech.coal.cosmotech_api.objects.parameters import Parameters

cosmotech/coal/cosmotech_api/objects/connection.py

@@ -0,0 +1,125 @@
+# Copyright (C) - 2023 - 2025 - Cosmo Tech
+# This document and all information contained herein is the exclusive property -
+# including all intellectual property rights pertaining thereto - of Cosmo Tech.
+# Any use, reproduction, translation, broadcasting, transmission, distribution,
+# etc., to any person is prohibited unless it has been previously and
+# specifically authorized by written means by Cosmo Tech.
+
+# Copyright (C) - 2023 - 2025 - Cosmo Tech
+# This document and all information contained herein is the exclusive property -
+# including all intellectual property rights pertaining thereto - of Cosmo Tech.
+# Any use, reproduction, translation, broadcasting, transmission, distribution,
+# etc., to any person is prohibited unless it has been previously and
+# specifically authorized by written means by Cosmo Tech.
+import os
+import pathlib
+
+import cosmotech_api
+from cosmotech.orchestrator.utils.translate import T
+from cosmotech_api import ApiClient
+
+from cosmotech.coal.utils.configuration import ENVIRONMENT_CONFIGURATION, Configuration
+from cosmotech.coal.utils.logger import LOGGER
+
+
+class Connection:
+    configuration: Configuration
+    api_client: ApiClient
+    api_type: str
+
+    # TODO :
+    # Replace use of environment variable by configuration usage
+    __api_env_keys = {"CSM_API_KEY", "CSM_API_URL"}
+    __azure_env_keys = {
+        "AZURE_CLIENT_ID",
+        "AZURE_CLIENT_SECRET",
+        "AZURE_TENANT_ID",
+        "CSM_API_URL",
+        "CSM_API_SCOPE",
+    }
+    __keycloak_env_keys = {
+        "IDP_TENANT_ID",
+        "IDP_CLIENT_ID",
+        "IDP_CLIENT_SECRET",
+        "IDP_BASE_URL",
+        "CSM_API_URL",
+    }
+
+    def __init__(
+        self,
+        configuration: Configuration = ENVIRONMENT_CONFIGURATION,
+    ):
+        self.configuration = configuration
+        self.api_client, self.api_type = self.get_api_client()
+
+    def get_api_client(self) -> (cosmotech_api.ApiClient, str):
+        existing_keys = set(os.environ.keys())
+        missing_azure_keys = self.__azure_env_keys - existing_keys
+        missing_api_keys = self.__api_env_keys - existing_keys
+        missing_keycloak_keys = self.__keycloak_env_keys - existing_keys
+
+        if all((missing_api_keys, missing_azure_keys, missing_keycloak_keys)):
+            LOGGER.error(T("coal.common.errors.no_env_vars"))
+            LOGGER.error(T("coal.cosmotech_api.connection.existing_sets"))
+            LOGGER.error(
+                T("coal.cosmotech_api.connection.azure_connection").format(keys=", ".join(self.__azure_env_keys))
+            )
+            LOGGER.error(
+                T("coal.cosmotech_api.connection.api_key_connection").format(keys=", ".join(self.__api_env_keys))
+            )
+            LOGGER.error(
+                T("coal.cosmotech_api.connection.keycloak_connection").format(keys=", ".join(self.__keycloak_env_keys))
+            )
+            raise EnvironmentError(T("coal.common.errors.no_env_vars"))
+
+        if not missing_keycloak_keys:
+            LOGGER.debug(T("coal.cosmotech_api.connection.found_keycloak"))
+            from keycloak import KeycloakOpenID
+
+            server_url = os.environ.get("IDP_BASE_URL")
+            if server_url[-1] != "/":
+                server_url = server_url + "/"
+            keycloack_parameters = dict(
+                server_url=server_url,
+                client_id=os.environ.get("IDP_CLIENT_ID"),
+                realm_name=os.environ.get("IDP_TENANT_ID"),
+                client_secret_key=os.environ.get("IDP_CLIENT_SECRET"),
+            )
+            if (ca_cert_path := os.environ.get("IDP_CA_CERT")) and pathlib.Path(ca_cert_path).exists():
+                LOGGER.info(T("coal.cosmotech_api.connection.found_cert_authority"))
+                keycloack_parameters["verify"] = ca_cert_path
+            keycloak_openid = KeycloakOpenID(**keycloack_parameters)
+
+            access_token = keycloak_openid.token(grant_type="client_credentials")
+
+            configuration = cosmotech_api.Configuration(
+                host=os.environ.get("CSM_API_URL"),
+                access_token=access_token["access_token"],
+            )
+            return cosmotech_api.ApiClient(configuration), "Keycloak Connection"
+
+        if not missing_api_keys:
+            LOGGER.debug(T("coal.cosmotech_api.connection.found_api_key"))
+            configuration = cosmotech_api.Configuration(
+                host=os.environ.get("CSM_API_URL"),
+            )
+            return (
+                cosmotech_api.ApiClient(
+                    configuration,
+                    os.environ.get("CSM_API_KEY_HEADER", "X-CSM-API-KEY"),
+                    os.environ.get("CSM_API_KEY"),
+                ),
+                "Cosmo Tech API Key",
+            )
+
+        if not missing_azure_keys:
+            LOGGER.debug(T("coal.cosmotech_api.connection.found_azure"))
+            from azure.identity import EnvironmentCredential
+
+            credentials = EnvironmentCredential()
+            token = credentials.get_token(os.environ.get("CSM_API_SCOPE"))
+
+            configuration = cosmotech_api.Configuration(host=os.environ.get("CSM_API_URL"), access_token=token.token)
+            return cosmotech_api.ApiClient(configuration), "Azure Entra Connection"
+
+        raise EnvironmentError(T("coal.common.errors.no_valid_connection"))