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

cosmotech/coal/cosmotech_api/runner/parameters.py

@@ -0,0 +1,157 @@
+# 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.
+
+"""
+Parameter handling functions.
+"""
+
+import json
+import os
+import pathlib
+from csv import DictWriter
+from typing import List, Dict, Any
+
+from cosmotech.coal.utils.logger import LOGGER
+from cosmotech.orchestrator.utils.translate import T
+
+
+def get_runner_parameters(runner_data: Any) -> Dict[str, Any]:
+    """
+    Extract parameters from runner data.
+
+    Args:
+        runner_data: Runner data object
+
+    Returns:
+        Dictionary mapping parameter IDs to values
+    """
+    content = dict()
+    for parameter in runner_data.parameters_values:
+        content[parameter.parameter_id] = parameter.value
+    return content
+
+
+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(parameter_folder: str, parameters: List[Dict[str, Any]]) -> str:
+    """
+    Write parameters to a JSON file.
+
+    Args:
+        parameter_folder: Folder to write the file to
+        parameters: List of parameter dictionaries
+
+    Returns:
+        Path to the created file
+    """
+    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(parameters, _file, indent=2)
+
+    return tmp_parameter_file
+
+
+def write_parameters_to_csv(parameter_folder: str, parameters: List[Dict[str, Any]]) -> str:
+    """
+    Write parameters to a CSV file.
+
+    Args:
+        parameter_folder: Folder to write the file to
+        parameters: List of parameter dictionaries
+
+    Returns:
+        Path to the created file
+    """
+    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(parameters)
+
+    return tmp_parameter_file
+
+
+def write_parameters(
+    parameter_folder: str,
+    parameters: List[Dict[str, Any]],
+    write_csv: bool = True,
+    write_json: bool = False,
+) -> Dict[str, str]:
+    """
+    Write parameters to files based on specified formats.
+
+    Args:
+        parameter_folder: Folder to write the files to
+        parameters: List of parameter dictionaries
+        write_csv: Whether to write a CSV file
+        write_json: Whether to write a JSON file
+
+    Returns:
+        Dictionary mapping file types to file paths
+    """
+    result = {}
+
+    if write_csv:
+        result["csv"] = write_parameters_to_csv(parameter_folder, parameters)
+
+    if write_json:
+        result["json"] = write_parameters_to_json(parameter_folder, parameters)
+
+    return result

cosmotech/coal/cosmotech_api/twin_data_layer.py

@@ -0,0 +1,512 @@
+# 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.
+
+"""
+Twin Data Layer operations module.
+
+This module provides functions for interacting with the Twin Data Layer,
+including sending and loading files.
+"""
+
+import json
+import pathlib
+from csv import DictReader, DictWriter
+from io import StringIO
+from typing import Dict, List, Any, Optional, Set, Tuple
+
+import requests
+from cosmotech_api import DatasetApi, RunnerApi, DatasetTwinGraphQuery
+
+from cosmotech.coal.cosmotech_api.connection import get_api_client
+from cosmotech.coal.utils.logger import LOGGER
+from cosmotech.orchestrator.utils.translate import T
+
+ID_COLUMN = "id"
+
+SOURCE_COLUMN = "src"
+
+TARGET_COLUMN = "dest"
+
+BATCH_SIZE_LIMIT = 10000
+
+
+class CSVSourceFile:
+    def __init__(self, file_path: pathlib.Path):
+        self.file_path = file_path
+        if not file_path.name.endswith(".csv"):
+            raise ValueError(T("coal.common.validation.not_csv_file").format(file_path=file_path))
+        with open(file_path) as _file:
+            dr = DictReader(_file)
+            self.fields = list(dr.fieldnames)
+        self.object_type = file_path.name[:-4]
+
+        self.id_column = None
+        self.source_column = None
+        self.target_column = None
+
+        for _c in self.fields:
+            if _c.lower() == ID_COLUMN:
+                self.id_column = _c
+            if _c.lower() == SOURCE_COLUMN:
+                self.source_column = _c
+            if _c.lower() == TARGET_COLUMN:
+                self.target_column = _c
+
+        has_id = self.id_column is not None
+        has_source = self.source_column is not None
+        has_target = self.target_column is not None
+
+        is_relation = all([has_source, has_target])
+
+        if not has_id and not is_relation:
+            LOGGER.error(T("coal.common.validation.invalid_nodes_relations").format(file_path=file_path))
+            LOGGER.error(T("coal.common.validation.node_requirements").format(id_column=ID_COLUMN))
+            LOGGER.error(
+                T("coal.common.validation.relationship_requirements").format(
+                    id_column=ID_COLUMN,
+                    source_column=SOURCE_COLUMN,
+                    target_column=TARGET_COLUMN,
+                )
+            )
+            raise ValueError(T("coal.common.validation.invalid_nodes_relations").format(file_path=file_path))
+
+        self.is_node = has_id and not is_relation
+
+        self.content_fields = {
+            _f: _f for _f in self.fields if _f not in [self.id_column, self.source_column, self.target_column]
+        }
+        if has_id:
+            self.content_fields[ID_COLUMN] = self.id_column
+        if is_relation:
+            self.content_fields[SOURCE_COLUMN] = self.source_column
+            self.content_fields[TARGET_COLUMN] = self.target_column
+
+    def reload(self, inplace: bool = False) -> "CSVSourceFile":
+        if inplace:
+            self.__init__(self.file_path)
+            return self
+        return CSVSourceFile(self.file_path)
+
+    def generate_query_insert(self) -> str:
+        """
+        Read a CSV file headers and generate a CREATE cypher query
+        :return: the Cypher query for CREATE
+        """
+
+        field_names = sorted(self.content_fields.keys(), key=len, reverse=True)
+
+        if self.is_node:
+            query = (
+                "CREATE (:"
+                + self.object_type
+                + ", ".join(f"{property_name}: ${self.content_fields[property_name]}" for property_name in field_names)
+                + "})"
+            )
+            # query = ("UNWIND $params AS params " +
+            #          f"MERGE (n:{self.object_type}) " +
+            #          "SET n += params")
+        else:
+            query = (
+                "MATCH "
+                + "(source {"
+                + ID_COLUMN
+                + ":$"
+                + self.source_column
+                + "}),\n"
+                + "(target {"
+                + ID_COLUMN
+                + ":$"
+                + self.target_column
+                + "})\n"
+                + "CREATE (source)-[rel:"
+                + self.object_type
+                + " {"
+                + ", ".join(f"{property_name}: ${self.content_fields[property_name]}" for property_name in field_names)
+                + "}"
+                + "]->(target)\n"
+            )
+            # query = ("UNWIND $params AS params " +
+            #          "MATCH (source {" + ID_COLUMN + ":params." + self.source_column + "})\n" +
+            #          "MATCH (target {" + ID_COLUMN + ":params." + self.target_column + "})\n" +
+            #          f"CREATE (from) - [rel:{self.object_type}]->(to)" +
+            #          "SET rel += params")
+        return query
+
+
+def get_dataset_id_from_runner(organization_id: str, workspace_id: str, runner_id: str) -> str:
+    """
+    Get the dataset ID from a runner.
+
+    Args:
+        organization_id: Organization ID
+        workspace_id: Workspace ID
+        runner_id: Runner ID
+
+    Returns:
+        Dataset ID
+    """
+    api_client, _ = get_api_client()
+    api_runner = RunnerApi(api_client)
+
+    runner_info = api_runner.get_runner(
+        organization_id,
+        workspace_id,
+        runner_id,
+    )
+
+    if (datasets_len := len(runner_info.dataset_list)) != 1:
+        LOGGER.error(
+            T("coal.cosmotech_api.runner.not_single_dataset").format(runner_id=runner_info.id, count=datasets_len)
+        )
+        LOGGER.debug(T("coal.cosmotech_api.runner.runner_info").format(info=runner_info))
+        raise ValueError(f"Runner {runner_info.id} does not have exactly one dataset")
+
+    return runner_info.dataset_list[0]
+
+
+def send_files_to_tdl(
+    api_url: str,
+    organization_id: str,
+    workspace_id: str,
+    runner_id: str,
+    directory_path: str,
+    clear: bool = True,
+) -> None:
+    """
+    Send CSV files to the Twin Data Layer.
+
+    Args:
+        api_url: API URL
+        organization_id: Organization ID
+        workspace_id: Workspace ID
+        runner_id: Runner ID
+        directory_path: Directory containing CSV files
+        clear: Whether to clear the dataset before sending files
+    """
+    api_client, _ = get_api_client()
+    api_ds = DatasetApi(api_client)
+
+    # Get dataset ID from runner
+    dataset_id = get_dataset_id_from_runner(organization_id, workspace_id, runner_id)
+
+    # Get dataset info
+    dataset_info = api_ds.find_dataset_by_id(organization_id, dataset_id)
+    dataset_info.ingestion_status = "SUCCESS"
+    api_ds.update_dataset(organization_id, dataset_id, dataset_info)
+
+    # Process CSV files
+    entities_queries = {}
+    relation_queries = {}
+
+    content_path = pathlib.Path(directory_path)
+    if not content_path.is_dir():
+        LOGGER.error(T("coal.common.file_operations.not_directory").format(target_dir=directory_path))
+        raise ValueError(f"{directory_path} is not a directory")
+
+    # Process CSV files
+    for file_path in content_path.glob("*.csv"):
+        _csv = CSVSourceFile(file_path)
+        if _csv.is_node:
+            LOGGER.info(T("coal.services.azure_storage.sending_content").format(file=file_path))
+            entities_queries[file_path] = _csv.generate_query_insert()
+        else:
+            LOGGER.info(T("coal.services.azure_storage.sending_content").format(file=file_path))
+            relation_queries[file_path] = _csv.generate_query_insert()
+
+    # Prepare headers
+    header = {
+        "Accept": "application/json",
+        "Content-Type": "text/csv",
+        "User-Agent": "OpenAPI-Generator/1.0.0/python",
+    }
+    header.update(api_client.default_headers)
+
+    for authtype, authinfo in api_ds.api_client.configuration.auth_settings().items():
+        api_ds.api_client._apply_auth_params(header, None, None, None, None, authinfo)
+
+    # Clear dataset if requested
+    if clear:
+        LOGGER.info(T("coal.services.azure_storage.clearing_content"))
+        clear_query = "MATCH (n) DETACH DELETE n"
+        api_ds.twingraph_query(organization_id, dataset_id, DatasetTwinGraphQuery(query=str(clear_query)))
+
+    # Send files
+    for query_dict in [entities_queries, relation_queries]:
+        for file_path, query in query_dict.items():
+            _process_csv_file(
+                file_path=file_path,
+                query=query,
+                api_url=api_url,
+                organization_id=organization_id,
+                dataset_id=dataset_id,
+                header=header,
+            )
+
+    LOGGER.info(T("coal.services.azure_storage.all_data_sent"))
+
+    # Update dataset status
+    dataset_info.ingestion_status = "SUCCESS"
+    dataset_info.twincache_status = "FULL"
+    api_ds.update_dataset(organization_id, dataset_id, dataset_info)
+
+
+def _process_csv_file(
+    file_path: pathlib.Path,
+    query: str,
+    api_url: str,
+    organization_id: str,
+    dataset_id: str,
+    header: Dict[str, str],
+) -> None:
+    """
+    Process a CSV file and send it to the Twin Data Layer.
+
+    Args:
+        file_path: Path to the CSV file
+        query: Query to execute
+        api_url: API URL
+        organization_id: Organization ID
+        dataset_id: Dataset ID
+        header: HTTP headers
+    """
+    content = StringIO()
+    size = 0
+    batch = 1
+    errors = []
+    query_craft = api_url + f"/organizations/{organization_id}/datasets/{dataset_id}/batch?query={query}"
+    LOGGER.info(T("coal.services.azure_storage.sending_content").format(file=file_path))
+
+    with open(file_path, "r") as _f:
+        dr = DictReader(_f)
+        dw = DictWriter(content, fieldnames=sorted(dr.fieldnames, key=len, reverse=True))
+        dw.writeheader()
+        for row in dr:
+            dw.writerow(row)
+            size += 1
+            if size > BATCH_SIZE_LIMIT:
+                LOGGER.info(T("coal.services.azure_storage.row_batch").format(count=batch * BATCH_SIZE_LIMIT))
+                batch += 1
+                content.seek(0)
+                post = requests.post(query_craft, data=content.read(), headers=header)
+                post.raise_for_status()
+                errors.extend(json.loads(post.content)["errors"])
+                content = StringIO()
+                dw = DictWriter(
+                    content,
+                    fieldnames=sorted(dr.fieldnames, key=len, reverse=True),
+                )
+                dw.writeheader()
+                size = 0
+
+    if size > 0:
+        content.seek(0)
+        post = requests.post(query_craft, data=content.read(), headers=header)
+        post.raise_for_status()
+        errors.extend(json.loads(post.content)["errors"])
+
+    if len(errors):
+        LOGGER.error(T("coal.services.azure_storage.import_errors").format(count=len(errors)))
+        for _err in errors:
+            LOGGER.error(T("coal.services.azure_storage.error_detail").format(error=str(_err)))
+        raise ValueError(f"Error importing data from {file_path}")
+
+
+def load_files_from_tdl(
+    organization_id: str,
+    workspace_id: str,
+    directory_path: str,
+    runner_id: str,
+) -> None:
+    """
+    Load files from the Twin Data Layer.
+
+    Args:
+        organization_id: Organization ID
+        workspace_id: Workspace ID
+        directory_path: Directory to save files to
+        runner_id: Runner ID
+    """
+    api_client, _ = get_api_client()
+    api_ds = DatasetApi(api_client)
+
+    # Get dataset ID from runner
+    dataset_id = get_dataset_id_from_runner(organization_id, workspace_id, runner_id)
+
+    # Get dataset info
+    dataset_info = api_ds.find_dataset_by_id(organization_id, dataset_id)
+    if dataset_info.ingestion_status != "SUCCESS":
+        LOGGER.error(
+            T("coal.cosmotech_api.runner.dataset_state").format(
+                dataset_id=dataset_id, status=dataset_info.ingestion_status
+            )
+        )
+        LOGGER.debug(T("coal.cosmotech_api.runner.dataset_info").format(info=dataset_info))
+        raise ValueError(f"Dataset {dataset_id} is not in SUCCESS state")
+
+    # Create directory
+    directory_path = pathlib.Path(directory_path)
+    if directory_path.is_file():
+        LOGGER.error(T("coal.common.file_operations.not_directory").format(target_dir=directory_path))
+        raise ValueError(f"{directory_path} is not a directory")
+
+    directory_path.mkdir(parents=True, exist_ok=True)
+
+    # Get node and relationship properties
+    item_queries = {}
+    properties_nodes = _get_node_properties(api_ds, organization_id, dataset_id)
+    properties_relationships = _get_relationship_properties(api_ds, organization_id, dataset_id)
+
+    # Create queries
+    for label, keys in properties_nodes.items():
+        node_query = f"MATCH (n:{label}) RETURN {', '.join(map(lambda k: f'n.`{k}` as `{k}`', keys))}"
+        item_queries[label] = node_query
+
+    for label, keys in properties_relationships.items():
+        rel_query = f"MATCH ()-[n:{label}]->() RETURN {', '.join(map(lambda k: f'n.`{k}` as `{k}`', keys))}"
+        item_queries[label] = rel_query
+
+    # Execute queries and write files
+    files_content, files_headers = _execute_queries(api_ds, organization_id, dataset_id, item_queries)
+    _write_files(directory_path, files_content, files_headers)
+
+    LOGGER.info(T("coal.services.azure_storage.all_csv_written"))
+
+
+def _get_node_properties(api_ds: DatasetApi, organization_id: str, dataset_id: str) -> Dict[str, Set[str]]:
+    """
+    Get node properties from the Twin Data Layer.
+
+    Args:
+        api_ds: Dataset API
+        organization_id: Organization ID
+        dataset_id: Dataset ID
+
+    Returns:
+        Dictionary of node labels to sets of property keys
+    """
+    get_node_properties_query = "MATCH (n) RETURN distinct labels(n)[0] as label, keys(n) as keys"
+    node_properties_results: List[Dict[str, Any]] = api_ds.twingraph_query(
+        organization_id,
+        dataset_id,
+        DatasetTwinGraphQuery(query=get_node_properties_query),
+    )
+
+    properties_nodes = {}
+    for _r in node_properties_results:
+        label = _r["label"]
+        keys = _r["keys"]
+        if label not in properties_nodes:
+            properties_nodes[label] = set()
+        properties_nodes[label].update(keys)
+
+    return properties_nodes
+
+
+def _get_relationship_properties(api_ds: DatasetApi, organization_id: str, dataset_id: str) -> Dict[str, Set[str]]:
+    """
+    Get relationship properties from the Twin Data Layer.
+
+    Args:
+        api_ds: Dataset API
+        organization_id: Organization ID
+        dataset_id: Dataset ID
+
+    Returns:
+        Dictionary of relationship types to sets of property keys
+    """
+    get_relationship_properties_query = "MATCH ()-[r]->() RETURN distinct type(r) as label, keys(r) as keys"
+    relationship_properties_results: List[Dict[str, Any]] = api_ds.twingraph_query(
+        organization_id,
+        dataset_id,
+        DatasetTwinGraphQuery(query=get_relationship_properties_query),
+    )
+
+    properties_relationships = {}
+    for _r in relationship_properties_results:
+        label = _r["label"]
+        keys = _r["keys"]
+        if label not in properties_relationships:
+            properties_relationships[label] = set()
+        properties_relationships[label].update(keys)
+
+    return properties_relationships
+
+
+def _execute_queries(
+    api_ds: DatasetApi, organization_id: str, dataset_id: str, item_queries: Dict[str, str]
+) -> Tuple[Dict[str, List[Dict[str, Any]]], Dict[str, Set[str]]]:
+    """
+    Execute queries against the Twin Data Layer.
+
+    Args:
+        api_ds: Dataset API
+        organization_id: Organization ID
+        dataset_id: Dataset ID
+        item_queries: Dictionary of element types to queries
+
+    Returns:
+        Tuple of (files_content, files_headers)
+    """
+    files_content = {}
+    files_headers = {}
+
+    for element_type, query in item_queries.items():
+        element_query: List[Dict[str, Any]] = api_ds.twingraph_query(
+            organization_id, dataset_id, DatasetTwinGraphQuery(query=query)
+        )
+        for element in element_query:
+            if element_type not in files_content:
+                files_content[element_type] = []
+                files_headers[element_type] = set()
+            files_content[element_type].append(element)
+            files_headers[element_type].update(element.keys())
+
+    return files_content, files_headers
+
+
+def _write_files(
+    directory_path: pathlib.Path,
+    files_content: Dict[str, List[Dict[str, Any]]],
+    files_headers: Dict[str, Set[str]],
+) -> None:
+    """
+    Write files to disk.
+
+    Args:
+        directory_path: Directory to write files to
+        files_content: Dictionary of file names to lists of rows
+        files_headers: Dictionary of file names to sets of headers
+    """
+    for file_name in files_content.keys():
+        file_path = directory_path / (file_name + ".csv")
+        LOGGER.info(
+            T("coal.services.azure_storage.writing_lines").format(count=len(files_content[file_name]), file=file_path)
+        )
+        with file_path.open("w") as _f:
+            headers = files_headers[file_name]
+            has_id = "id" in headers
+            is_relation = "src" in headers
+            new_headers = []
+            if has_id:
+                headers.remove("id")
+                new_headers.append("id")
+            if is_relation:
+                headers.remove("src")
+                headers.remove("dest")
+                new_headers.append("src")
+                new_headers.append("dest")
+            headers = new_headers + sorted(headers)
+
+            dw = DictWriter(_f, fieldnames=headers)
+            dw.writeheader()
+            for row in sorted(files_content[file_name], key=lambda r: r.get("id", "")):
+                dw.writerow(
+                    {
+                        key: (json.dumps(value) if isinstance(value, (bool, dict, list)) else value)
+                        for key, value in row.items()
+                    }
+                )