cosmotech-acceleration-library 1.1.0__py3-none-any.whl → 2.1.0rc1__py3-none-any.whl

cosmotech/coal/cosmotech_api/apis/runner.py

@@ -0,0 +1,75 @@
+# 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,
+        runner_id: Optional[str] = None,
+        include: Optional[list[str]] = None,
+        exclude: Optional[list[str]] = None,
+    ) -> dict[str, Any]:
+        runner = self.get_runner(
+            self.configuration.cosmotech.organization_id,
+            self.configuration.cosmotech.workspace_id,
+            runner_id or self.configuration.cosmotech.runner_id,
+        )
+
+        return runner.model_dump(by_alias=True, exclude_none=True, include=include, exclude=exclude, mode="json")
+
+    def download_runner_data(
+        self,
+        download_datasets: Optional[str] = None,
+    ):
+        LOGGER.info(T("coal.cosmotech_api.runner.starting_download"))
+
+        # Get runner data
+        runner = self.get_runner(
+            self.configuration.cosmotech.organization_id,
+            self.configuration.cosmotech.workspace_id,
+            self.configuration.cosmotech.runner_id,
+        )
+
+        # Skip if no parameters found
+        if not runner.parameters_values:
+            LOGGER.warning(T("coal.cosmotech_api.runner.no_parameters"))
+        else:
+            LOGGER.info(T("coal.cosmotech_api.runner.loaded_data"))
+            parameters = Parameters(runner)
+            parameters.write_parameters_to_json(self.configuration.cosmotech.parameters_absolute_path)
+
+        if runner.datasets.parameter:
+            ds_api = DatasetApi(self.configuration)
+            ds_api.download_parameter(runner.datasets.parameter)
+
+        # Download datasets if requested
+        if download_datasets:
+            LOGGER.info(T("coal.cosmotech_api.runner.downloading_datasets").format(count=len(runner.datasets.bases)))
+            if runner.datasets.bases:
+                ds_api = DatasetApi(self.configuration)
+                for dataset_id in runner.datasets.bases:
+                    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"))

cosmotech/coal/cosmotech_api/objects/parameters.py

@@ -0,0 +1,127 @@
+# 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 json
+import os
+import pathlib
+from csv import DictWriter
+from typing import Any, Dict, List
+
+from cosmotech.orchestrator.utils.translate import T
+
+from cosmotech.coal.utils.logger import LOGGER
+
+
+class Parameters:
+    values: Dict[str, Any] = dict()
+    parameters_list: List[Dict[str, Any]] = list()
+
+    def __init__(self, runner_data: Any):
+        """
+        Extract parameters from runner data.
+
+        Args:
+            runner_data: Runner data object
+
+        Returns:
+            Dictionary mapping parameter IDs to values
+        """
+        for parameter in runner_data.parameters_values:
+            self.values[parameter.parameter_id] = parameter.value
+        self.parameters_list = self.format_parameters_list(runner_data)
+
+    @staticmethod
+    def format_parameters_list(runner_data: Any) -> List[Dict[str, Any]]:
+        """
+        Format parameters from runner data as a list of dictionaries.
+
+        Args:
+            runner_data: Runner data object
+
+        Returns:
+            List of parameter dictionaries
+        """
+        parameters = []
+
+        if not runner_data.parameters_values:
+            return parameters
+
+        max_name_size = max(map(lambda r: len(r.parameter_id), runner_data.parameters_values))
+        max_type_size = max(map(lambda r: len(r.var_type), runner_data.parameters_values))
+
+        for parameter_data in runner_data.parameters_values:
+            parameter_name = parameter_data.parameter_id
+            value = parameter_data.value
+            var_type = parameter_data.var_type
+            is_inherited = parameter_data.is_inherited
+
+            parameters.append(
+                {
+                    "parameterId": parameter_name,
+                    "value": value,
+                    "varType": var_type,
+                    "isInherited": is_inherited,
+                }
+            )
+
+            LOGGER.debug(
+                T("coal.cosmotech_api.runner.parameter_debug").format(
+                    param_id=parameter_name,
+                    max_name_size=max_name_size,
+                    var_type=var_type,
+                    max_type_size=max_type_size,
+                    value=value,
+                    inherited=" inherited" if is_inherited else "",
+                )
+            )
+
+        return parameters
+
+    def write_parameters_to_json(
+        self,
+        parameter_folder: str,
+    ) -> str:
+        pathlib.Path(parameter_folder).mkdir(exist_ok=True, parents=True)
+        tmp_parameter_file = os.path.join(parameter_folder, "parameters.json")
+
+        LOGGER.info(T("coal.cosmotech_api.runner.generating_file").format(file=tmp_parameter_file))
+
+        with open(tmp_parameter_file, "w") as _file:
+            json.dump(self.parameters_list, _file, indent=2)
+
+        return tmp_parameter_file
+
+    def write_parameters_to_csv(
+        self,
+        parameter_folder: str,
+    ) -> str:
+        pathlib.Path(parameter_folder).mkdir(exist_ok=True, parents=True)
+        tmp_parameter_file = os.path.join(parameter_folder, "parameters.csv")
+
+        LOGGER.info(T("coal.cosmotech_api.runner.generating_file").format(file=tmp_parameter_file))
+
+        with open(tmp_parameter_file, "w") as _file:
+            _w = DictWriter(_file, fieldnames=["parameterId", "value", "varType", "isInherited"])
+            _w.writeheader()
+            _w.writerows(self.parameters_list)
+
+        return tmp_parameter_file
+
+    def write_parameters(
+        self,
+        parameter_folder: str,
+        write_csv: bool = True,
+        write_json: bool = False,
+    ) -> Dict[str, str]:
+        result = {}
+
+        if write_csv:
+            result["csv"] = self.write_parameters_to_csv(parameter_folder)
+
+        if write_json:
+            result["json"] = self.write_parameters_to_json(parameter_folder)
+
+        return result

cosmotech/coal/postgresql/runner.py

@@ -13,83 +13,100 @@ for runner metadata operations.
 """
 
 from adbc_driver_postgresql import dbapi
+from cosmotech.orchestrator.utils.translate import T
 
-from cosmotech.coal.cosmotech_api.connection import get_api_client
-from cosmotech.coal.cosmotech_api.runner.metadata import get_runner_metadata
+from cosmotech.coal.cosmotech_api.apis.runner import RunnerApi
+from cosmotech.coal.postgresql.utils import PostgresUtils
+from cosmotech.coal.utils.configuration import Configuration
 from cosmotech.coal.utils.logger import LOGGER
-from cosmotech.coal.utils.postgresql import generate_postgresql_full_uri
-from cosmotech.orchestrator.utils.translate import T
 
 
 def send_runner_metadata_to_postgresql(
-    organization_id: str,
-    workspace_id: str,
-    runner_id: str,
-    postgres_host: str,
-    postgres_port: int,
-    postgres_db: str,
-    postgres_schema: str,
-    postgres_user: str,
-    postgres_password: str,
-    table_prefix: str = "Cosmotech_",
-    force_encode: bool = False,
-) -> None:
+    configuration: Configuration,
+) -> str:
     """
     Send runner metadata to a PostgreSQL database.
 
     Args:
+        configuration: coal configuration
         organization_id: Organization ID
         workspace_id: Workspace ID
         runner_id: Runner ID
-        postgres_host: PostgreSQL host
-        postgres_port: PostgreSQL port
-        postgres_db: PostgreSQL database name
-        postgres_schema: PostgreSQL schema
-        postgres_user: PostgreSQL username
-        postgres_password: PostgreSQL password
-        table_prefix: Table prefix
-        force_encode: force password encoding to percent encoding
     """
-    # Get runner metadata
-    with get_api_client()[0] as api_client:
-        runner = get_runner_metadata(api_client, organization_id, workspace_id, runner_id)
+    _psql = PostgresUtils(configuration)
 
-    # Generate PostgreSQL URI
-    postgresql_full_uri = generate_postgresql_full_uri(
-        postgres_host, str(postgres_port), postgres_db, postgres_user, postgres_password, force_encode=force_encode
+    # Get runner metadata
+    _runner_api = RunnerApi(configuration)
+    runner = _runner_api.get_runner_metadata(
+        configuration.cosmotech.runner_id,
     )
 
     # Connect to PostgreSQL and update runner metadata
-    with dbapi.connect(postgresql_full_uri, autocommit=True) as conn:
+    with dbapi.connect(_psql.full_uri, autocommit=True) as conn:
         with conn.cursor() as curs:
-            schema_table = f"{postgres_schema}.{table_prefix}RunnerMetadata"
+            schema_table = f"{str(_psql.db_schema)}.{str(_psql.table_prefix)}RunnerMetadata"
             sql_create_table = f"""
                 CREATE TABLE IF NOT EXISTS {schema_table}  (
                   id varchar(32) PRIMARY KEY,
                   name varchar(256),
-                  last_run_id varchar(32),
+                  last_csm_run_id varchar(32),
                   run_template_id varchar(32)
                 );
             """
-            sql_upsert = f"""
-                INSERT INTO {schema_table} (id, name, last_run_id, run_template_id)
-                  VALUES(%s, %s, %s, %s)
-                  ON CONFLICT (id)
-                  DO
-                    UPDATE SET name = EXCLUDED.name, last_run_id = EXCLUDED.last_run_id;
-            """
             LOGGER.info(T("coal.services.postgresql.creating_table").format(schema_table=schema_table))
             curs.execute(sql_create_table)
             conn.commit()
             LOGGER.info(T("coal.services.postgresql.metadata"))
+            sql_upsert = f"""
+                INSERT INTO {schema_table} (id, name, last_csm_run_id, run_template_id)
+                  VALUES ($1, $2, $3, $4)
+                  ON CONFLICT (id)
+                  DO
+                    UPDATE SET name = EXCLUDED.name, last_csm_run_id = EXCLUDED.last_csm_run_id;
+            """
+            LOGGER.debug(runner)
             curs.execute(
                 sql_upsert,
                 (
                     runner.get("id"),
                     runner.get("name"),
-                    runner.get("lastRunId"),
+                    runner.get("lastRunInfo").get("lastRunId"),
                     runner.get("runTemplateId"),
                 ),
             )
             conn.commit()
             LOGGER.info(T("coal.services.postgresql.metadata_updated"))
+    return runner.get("lastRunInfo").get("lastRunId")
+
+
+def remove_runner_metadata_from_postgresql(
+    configuration: Configuration,
+) -> str:
+    """
+    Removes run_id from metadata table that trigger cascade delete on other tables
+
+    Args:
+        configuration: coal configuration
+        organization_id: Organization ID
+        workspace_id: Workspace ID
+        runner_id: Runner ID
+    """
+    _psql = PostgresUtils(configuration)
+
+    # Get runner metadata
+    _runner_api = RunnerApi(configuration)
+    runner = _runner_api.get_runner_metadata(
+        configuration.cosmotech.runner_id,
+    )
+
+    # Connect to PostgreSQL and remove runner metadata row
+    with dbapi.connect(_psql.full_uri, autocommit=True) as conn:
+        with conn.cursor() as curs:
+            schema_table = f"{_psql.db_schema}.{_psql.table_prefix}RunnerMetadata"
+            sql_delete_from_metatable = f"""
+                DELETE FROM {schema_table}
+                WHERE last_csm_run_id={runner.get("lastRunId")};
+            """
+            curs.execute(sql_delete_from_metatable)
+            conn.commit()
+    return runner.get("lastRunId")