cosmotech-acceleration-library 1.0.0__py3-none-any.whl

cosmotech/coal/cosmotech_api/run_data.py

@@ -0,0 +1,173 @@
+# 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.
+
+"""
+Run Data Service operations module.
+
+This module provides functions for interacting with the Run Data Service,
+including sending and loading data.
+"""
+
+import json
+import pathlib
+from csv import DictReader, DictWriter
+from typing import Dict, List, Any, Optional, Set
+
+from cosmotech_api import SendRunDataRequest, RunDataQuery
+from cosmotech_api.api.run_api import RunApi
+
+from cosmotech.coal.cosmotech_api.connection import get_api_client
+from cosmotech.coal.store.store import Store
+from cosmotech.coal.store.native_python import convert_table_as_pylist
+from cosmotech.coal.utils.logger import LOGGER
+from cosmotech.orchestrator.utils.translate import T
+
+
+def send_csv_to_run_data(
+    source_folder: str,
+    organization_id: str,
+    workspace_id: str,
+    runner_id: str,
+    run_id: str,
+) -> None:
+    """
+    Send CSV files to the Run Data Service.
+
+    Args:
+        source_folder: Folder containing CSV files
+        organization_id: Organization ID
+        workspace_id: Workspace ID
+        runner_id: Runner ID
+        run_id: Run ID
+    """
+    source_dir = pathlib.Path(source_folder)
+
+    if not source_dir.exists():
+        LOGGER.error(T("coal.common.file_operations.not_found").format(source_folder=source_dir))
+        raise FileNotFoundError(f"{source_dir} does not exist")
+
+    with get_api_client()[0] as api_client:
+        api_run = RunApi(api_client)
+        for csv_path in source_dir.glob("*.csv"):
+            with open(csv_path) as _f:
+                dr = DictReader(_f)
+                table_name = csv_path.name.replace(".csv", "")
+                LOGGER.info(T("coal.cosmotech_api.run_data.sending_to_table").format(table_name=f"CD_{table_name}"))
+                LOGGER.debug(T("coal.services.database.column_list").format(columns=dr.fieldnames))
+                data = []
+
+                for row in dr:
+                    n_row = dict()
+                    for k, v in row.items():
+                        if isinstance(v, str):
+                            try:
+                                n_row[k] = json.loads(v)
+                            except json.decoder.JSONDecodeError:
+                                n_row[k] = v
+                        else:
+                            n_row[k] = v
+                    data.append(n_row)
+
+                LOGGER.info(T("coal.services.database.row_count").format(count=len(data)))
+                api_run.send_run_data(
+                    organization_id,
+                    workspace_id,
+                    runner_id,
+                    run_id,
+                    SendRunDataRequest(id=table_name, data=data),
+                )
+
+
+def send_store_to_run_data(
+    store_folder: str,
+    organization_id: str,
+    workspace_id: str,
+    runner_id: str,
+    run_id: str,
+) -> None:
+    """
+    Send store data to the Run Data Service.
+
+    Args:
+        store_folder: Folder containing the store
+        organization_id: Organization ID
+        workspace_id: Workspace ID
+        runner_id: Runner ID
+        run_id: Run ID
+    """
+    source_dir = pathlib.Path(store_folder)
+
+    if not source_dir.exists():
+        LOGGER.error(T("coal.common.file_operations.not_found").format(source_folder=source_dir))
+        raise FileNotFoundError(f"{source_dir} does not exist")
+
+    with get_api_client()[0] as api_client:
+        api_run = RunApi(api_client)
+        _s = Store()
+        for table_name in _s.list_tables():
+            LOGGER.info(T("coal.cosmotech_api.run_data.sending_to_table").format(table_name=f"CD_{table_name}"))
+            data = convert_table_as_pylist(table_name)
+            if not len(data):
+                LOGGER.info(T("coal.services.database.no_rows"))
+                continue
+            fieldnames = _s.get_table_schema(table_name).names
+            for row in data:
+                for field in fieldnames:
+                    if row[field] is None:
+                        del row[field]
+            LOGGER.debug(T("coal.services.database.column_list").format(columns=fieldnames))
+            LOGGER.info(T("coal.services.database.row_count").format(count=len(data)))
+            api_run.send_run_data(
+                organization_id,
+                workspace_id,
+                runner_id,
+                run_id,
+                SendRunDataRequest(id=table_name, data=data),
+            )
+
+
+def load_csv_from_run_data(
+    target_folder: str,
+    organization_id: str,
+    workspace_id: str,
+    runner_id: str,
+    run_id: str,
+    file_name: str = "results",
+    query: str = "SELECT table_name FROM information_schema.tables WHERE table_schema='public'",
+) -> None:
+    """
+    Load data from the Run Data Service and save it as a CSV file.
+
+    Args:
+        target_folder: Folder to save the CSV file to
+        organization_id: Organization ID
+        workspace_id: Workspace ID
+        runner_id: Runner ID
+        run_id: Run ID
+        file_name: Name of the CSV file to create
+        query: SQL query to execute
+    """
+    target_dir = pathlib.Path(target_folder)
+    target_dir.mkdir(parents=True, exist_ok=True)
+
+    with get_api_client()[0] as api_client:
+        api_run = RunApi(api_client)
+        query_result = api_run.query_run_data(
+            organization_id, workspace_id, runner_id, run_id, RunDataQuery(query=query)
+        )
+        if query_result.result:
+            LOGGER.info(T("coal.services.database.query_results").format(count=len(query_result.result)))
+            with open(target_dir / (file_name + ".csv"), "w") as _f:
+                headers = set()
+                for r in query_result.result:
+                    headers = headers | set(r.keys())
+                dw = DictWriter(_f, fieldnames=sorted(headers))
+                dw.writeheader()
+                dw.writerows(query_result.result)
+            LOGGER.info(T("coal.services.database.saved_results").format(file=f"{target_dir / file_name}.csv"))
+        else:
+            LOGGER.info(T("coal.services.database.no_results"))

cosmotech/coal/cosmotech_api/run_template.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.
+
+"""
+Run Template operations module.
+
+This module provides functions for interacting with Run Templates,
+including downloading and extracting handlers.
+"""
+
+import pathlib
+from io import BytesIO
+from zipfile import BadZipfile, ZipFile
+from typing import List
+
+from cosmotech_api.api.solution_api import SolutionApi
+from cosmotech_api.api.workspace_api import Workspace, WorkspaceApi
+from cosmotech_api.exceptions import ServiceException
+
+from cosmotech.coal.cosmotech_api.connection import get_api_client
+from cosmotech.coal.utils.logger import LOGGER
+from cosmotech.orchestrator.utils.translate import T
+
+
+def load_run_template_handlers(
+    organization_id: str,
+    workspace_id: str,
+    run_template_id: str,
+    handler_list: str,
+) -> bool:
+    """
+    Download and extract run template handlers.
+
+    Args:
+        organization_id: Organization ID
+        workspace_id: Workspace ID
+        run_template_id: Run Template ID
+        handler_list: Comma-separated list of handlers to download
+
+    Returns:
+        True if all handlers were downloaded successfully, False otherwise
+
+    Raises:
+        ValueError: If the workspace or solution is not found
+    """
+    has_errors = False
+    with get_api_client()[0] as api_client:
+        api_w = WorkspaceApi(api_client)
+
+        LOGGER.info(T("coal.cosmotech_api.run_template.loading_solution"))
+        try:
+            r_data: Workspace = api_w.find_workspace_by_id(organization_id=organization_id, workspace_id=workspace_id)
+        except ServiceException as e:
+            LOGGER.error(
+                T("coal.cosmotech_api.workspace.not_found").format(
+                    workspace_id=workspace_id, organization_id=organization_id
+                )
+            )
+            LOGGER.debug(T("coal.cosmotech_api.run_template.error_details").format(details=e.body))
+            raise ValueError(f"Workspace {workspace_id} not found in organization {organization_id}")
+        solution_id = r_data.solution.solution_id
+
+        api_sol = SolutionApi(api_client)
+        handler_list = handler_list.replace("handle-parameters", "parameters_handler")
+        root_path = pathlib.Path("../csm_orc_port")
+        template_path = root_path / run_template_id
+        for handler_id in handler_list.split(","):
+            handler_path: pathlib.Path = template_path / handler_id
+            LOGGER.info(
+                T("coal.cosmotech_api.run_template.querying_handler").format(
+                    handler=handler_id, template=run_template_id
+                )
+            )
+            try:
+                rt_data = api_sol.download_run_template_handler(
+                    organization_id=organization_id,
+                    solution_id=solution_id,
+                    run_template_id=run_template_id,
+                    handler_id=handler_id,
+                )
+            except ServiceException as e:
+                LOGGER.error(
+                    T("coal.cosmotech_api.run_template.handler_not_found").format(
+                        handler=handler_id,
+                        template=run_template_id,
+                        solution=solution_id,
+                    )
+                )
+                LOGGER.debug(T("coal.cosmotech_api.run_template.error_details").format(details=e.body))
+                has_errors = True
+                continue
+            LOGGER.info(T("coal.cosmotech_api.run_template.extracting_handler").format(path=handler_path.absolute()))
+            handler_path.mkdir(parents=True, exist_ok=True)
+
+            try:
+                with ZipFile(BytesIO(rt_data)) as _zip:
+                    _zip.extractall(handler_path)
+            except BadZipfile:
+                LOGGER.error(T("coal.cosmotech_api.run_template.handler_not_zip").format(handler=handler_id))
+                has_errors = True
+        if has_errors:
+            LOGGER.error(T("coal.cosmotech_api.run_template.run_issues"))
+            return False
+        return True

cosmotech/coal/cosmotech_api/runner/__init__.py

@@ -0,0 +1,28 @@
+# 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.
+
+"""
+Runner and Run data handling module.
+"""
+
+# Re-export functions from submodules
+from cosmotech.coal.cosmotech_api.runner.data import get_runner_data
+from cosmotech.coal.cosmotech_api.runner.parameters import (
+    get_runner_parameters,
+    format_parameters_list,
+    write_parameters,
+    write_parameters_to_json,
+    write_parameters_to_csv,
+)
+from cosmotech.coal.cosmotech_api.runner.datasets import (
+    get_dataset_ids_from_runner,
+    download_dataset,
+    download_datasets,
+    dataset_to_file,
+)
+from cosmotech.coal.cosmotech_api.runner.download import download_runner_data
+from cosmotech.coal.cosmotech_api.runner.metadata import get_runner_metadata

cosmotech/coal/cosmotech_api/runner/data.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.
+
+"""
+Core runner data retrieval functions.
+"""
+
+from cosmotech_api.api.runner_api import RunnerApi
+from cosmotech.coal.cosmotech_api.connection import get_api_client
+from cosmotech.coal.utils.logger import LOGGER
+from cosmotech.orchestrator.utils.translate import T
+
+
+def get_runner_data(organization_id: str, workspace_id: str, runner_id: str):
+    """
+    Get runner data from the API.
+
+    Args:
+        organization_id: The ID of the organization
+        workspace_id: The ID of the workspace
+        runner_id: The ID of the runner
+
+    Returns:
+        Runner data object
+    """
+    LOGGER.info(T("coal.cosmotech_api.runner.loading_data"))
+    with get_api_client()[0] as api_client:
+        api_instance = RunnerApi(api_client)
+        runner_data = api_instance.get_runner(
+            organization_id=organization_id,
+            workspace_id=workspace_id,
+            runner_id=runner_id,
+        )
+    return runner_data