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

cosmotech/coal/azure/adx/store.py

@@ -0,0 +1,255 @@
+# 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 tempfile
+import uuid
+from typing import Optional, List, Dict, Tuple, Union, Any
+
+import pyarrow
+import pyarrow.csv as pc
+import time
+from azure.kusto.data import KustoClient
+from azure.kusto.data.data_format import DataFormat
+from azure.kusto.ingest import IngestionProperties
+from azure.kusto.ingest import QueuedIngestClient
+from azure.kusto.ingest import ReportLevel
+from cosmotech.orchestrator.utils.translate import T
+from time import perf_counter
+
+from cosmotech.coal.azure.adx.tables import check_and_create_table, _drop_by_tag
+from cosmotech.coal.azure.adx.auth import initialize_clients
+from cosmotech.coal.azure.adx.ingestion import monitor_ingestion, handle_failures
+from cosmotech.coal.store.store import Store
+from cosmotech.coal.utils.logger import LOGGER
+from cosmotech.coal.utils.postgresql import send_pyarrow_table_to_postgresql
+
+
+def send_table_data(
+    ingest_client: QueuedIngestClient, database: str, table_name: str, data: pyarrow.Table, operation_tag: str
+) -> Tuple[str, str]:
+    """
+    Send a PyArrow table to ADX.
+
+    Args:
+        ingest_client: The ingest client
+        database: The database name
+        table_name: The table name
+        data: The PyArrow table data
+        operation_tag: The operation tag for tracking
+
+    Returns:
+        tuple: (source_id, table_name)
+    """
+    LOGGER.debug(T("coal.services.adx.sending_data").format(table_name=table_name))
+    result = send_pyarrow_table_to_adx(ingest_client, database, table_name, data, operation_tag)
+    return result.source_id, table_name
+
+
+def process_tables(
+    store: Store, kusto_client: KustoClient, ingest_client: QueuedIngestClient, database: str, operation_tag: str
+) -> Tuple[List[str], Dict[str, str]]:
+    """
+    Process all tables in the store.
+
+    Args:
+        store: The data store
+        kusto_client: The Kusto client
+        ingest_client: The ingest client
+        database: The database name
+        operation_tag: The operation tag for tracking
+
+    Returns:
+        tuple: (source_ids, table_ingestion_id_mapping)
+    """
+    source_ids = []
+    table_ingestion_id_mapping = dict()
+
+    LOGGER.debug(T("coal.services.adx.listing_tables"))
+    table_list = list(store.list_tables())
+
+    for target_table_name in table_list:
+        LOGGER.info(T("coal.services.adx.working_on_table").format(table_name=target_table_name))
+        data = store.get_table(target_table_name)
+
+        if data.num_rows < 1:
+            LOGGER.warning(T("coal.services.adx.table_empty").format(table_name=target_table_name))
+            continue
+
+        check_and_create_table(kusto_client, database, target_table_name, data)
+
+        source_id, _ = send_table_data(ingest_client, database, target_table_name, data, operation_tag)
+        source_ids.append(source_id)
+        table_ingestion_id_mapping[source_id] = target_table_name
+
+    return source_ids, table_ingestion_id_mapping
+
+
+def send_pyarrow_table_to_adx(
+    client: QueuedIngestClient,
+    database: str,
+    table_name: str,
+    table_data: pyarrow.Table,
+    drop_by_tag: Optional[str] = None,
+):
+    drop_by_tags = [drop_by_tag] if (drop_by_tag is not None) else None
+
+    properties = IngestionProperties(
+        database=database,
+        table=table_name,
+        data_format=DataFormat.CSV,
+        drop_by_tags=drop_by_tags,
+        report_level=ReportLevel.FailuresAndSuccesses,
+        flush_immediately=True,
+    )
+
+    file_name = f"adx_{database}_{table_name}_{int(time.time())}_{uuid.uuid4()}.csv"
+    temp_file_path = os.path.join(os.environ.get("CSM_TEMP_ABSOLUTE_PATH", tempfile.gettempdir()), file_name)
+    pc.write_csv(table_data, temp_file_path, pc.WriteOptions(include_header=False))
+    try:
+        return client.ingest_from_file(temp_file_path, properties)
+    finally:
+        os.unlink(temp_file_path)
+
+
+def send_store_to_adx(
+    adx_uri: str,
+    adx_ingest_uri: str,
+    database_name: str,
+    wait: bool = False,
+    tag: Optional[str] = None,
+    store_location: Optional[str] = None,
+) -> Union[bool, Any]:
+    """
+    Send data from the store to Azure Data Explorer.
+
+    Args:
+        adx_uri: The Azure Data Explorer resource URI
+        adx_ingest_uri: The Azure Data Explorer resource ingest URI
+        database_name: The database name
+        wait: Whether to wait for ingestion to complete
+        tag: The operation tag for tracking (will generate a unique one if not provided)
+        store_location: Optional store location (uses default if not provided)
+
+    Returns:
+        bool: True if successful, False otherwise
+    """
+    # Generate a unique operation tag if none provided
+    operation_tag = tag or f"op-{str(uuid.uuid4())}"
+    LOGGER.debug(T("coal.services.adx.starting_ingestion").format(operation_tag=operation_tag))
+
+    # Initialize clients
+    kusto_client, ingest_client = initialize_clients(adx_uri, adx_ingest_uri)
+    database = database_name
+
+    # Load datastore
+    LOGGER.debug(T("coal.services.adx.loading_datastore"))
+    store = Store(store_location=store_location)
+
+    try:
+        # Process tables
+        source_ids, table_ingestion_id_mapping = process_tables(
+            store, kusto_client, ingest_client, database, operation_tag
+        )
+
+        LOGGER.info(T("coal.services.adx.data_sent"))
+
+        # Monitor ingestion if wait is True
+        has_failures = False
+        if wait and source_ids:
+            has_failures = monitor_ingestion(ingest_client, source_ids, table_ingestion_id_mapping)
+
+        # Handle failures
+        should_abort = handle_failures(kusto_client, database, operation_tag, has_failures)
+        if should_abort:
+            return False
+
+        return True
+
+    except Exception as e:
+        LOGGER.exception(T("coal.services.adx.ingestion_error"))
+        # Perform rollback using the tag
+        LOGGER.warning(T("coal.services.adx.dropping_data").format(operation_tag=operation_tag))
+        _drop_by_tag(kusto_client, database, operation_tag)
+        raise e
+
+
+def dump_store_to_adx(
+    store_folder: str,
+    postgres_host: str,
+    postgres_port: int,
+    postgres_db: str,
+    postgres_schema: str,
+    postgres_user: str,
+    postgres_password: str,
+    table_prefix: str = "Cosmotech_",
+    replace: bool = True,
+) -> None:
+    """
+    Dump Store data to an Azure Data Explorer database.
+
+    Args:
+        store_folder: Folder containing the Store
+        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
+        replace: Whether to replace existing tables
+    """
+    _s = Store(store_location=store_folder)
+
+    tables = list(_s.list_tables())
+    if len(tables):
+        LOGGER.info(T("coal.services.database.sending_data").format(table=f"{postgres_db}.{postgres_schema}"))
+        total_rows = 0
+        _process_start = perf_counter()
+        for table_name in tables:
+            _s_time = perf_counter()
+            target_table_name = f"{table_prefix}{table_name}"
+            LOGGER.info(T("coal.services.database.table_entry").format(table=target_table_name))
+            data = _s.get_table(table_name)
+            if not len(data):
+                LOGGER.info(T("coal.services.database.no_rows"))
+                continue
+            _dl_time = perf_counter()
+            rows = send_pyarrow_table_to_postgresql(
+                data,
+                target_table_name,
+                postgres_host,
+                postgres_port,
+                postgres_db,
+                postgres_schema,
+                postgres_user,
+                postgres_password,
+                replace,
+            )
+            total_rows += rows
+            _up_time = perf_counter()
+            LOGGER.info(T("coal.services.database.row_count").format(count=rows))
+            LOGGER.debug(
+                T("coal.common.timing.operation_completed").format(
+                    operation="Load from datastore", time=f"{_dl_time - _s_time:0.3}"
+                )
+            )
+            LOGGER.debug(
+                T("coal.common.timing.operation_completed").format(
+                    operation="Send to postgresql", time=f"{_up_time - _dl_time:0.3}"
+                )
+            )
+        _process_end = perf_counter()
+        LOGGER.info(
+            T("coal.services.database.rows_fetched").format(
+                table="all tables",
+                count=total_rows,
+                time=f"{_process_end - _process_start:0.3}",
+            )
+        )
+    else:
+        LOGGER.info(T("coal.services.database.store_empty"))

cosmotech/coal/azure/adx/tables.py

@@ -0,0 +1,118 @@
+# 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 Dict, Any
+
+import pyarrow
+from azure.kusto.data import KustoClient
+
+from cosmotech.coal.utils.logger import LOGGER
+from cosmotech.orchestrator.utils.translate import T
+
+
+def table_exists(client: KustoClient, database: str, table_name: str) -> bool:
+    """
+    Check if a table exists in the database.
+
+    Args:
+        client: The KustoClient to use
+        database: The name of the database
+        table_name: The name of the table to check
+
+    Returns:
+        bool: True if the table exists, False otherwise
+    """
+    LOGGER.debug(T("coal.services.adx.checking_table").format(database=database, table_name=table_name))
+
+    get_tables_query = f".show database ['{database}'] schema| distinct TableName"
+    tables = client.execute(database, get_tables_query)
+
+    for r in tables.primary_results[0]:
+        if table_name == r[0]:
+            LOGGER.debug(T("coal.services.adx.table_exists").format(table_name=table_name))
+            return True
+
+    LOGGER.debug(T("coal.services.adx.table_not_exists").format(table_name=table_name))
+    return False
+
+
+def check_and_create_table(kusto_client: KustoClient, database: str, table_name: str, data: pyarrow.Table) -> bool:
+    """
+    Check if a table exists and create it if it doesn't.
+
+    Args:
+        kusto_client: The Kusto client
+        database: The database name
+        table_name: The table name
+        data: The PyArrow table data
+
+    Returns:
+        bool: True if the table was created, False if it already existed
+    """
+    LOGGER.debug(T("coal.services.adx.checking_table_exists"))
+    if not table_exists(kusto_client, database, table_name):
+        from cosmotech.coal.azure.adx.utils import create_column_mapping
+
+        mapping = create_column_mapping(data)
+        LOGGER.debug(T("coal.services.adx.creating_nonexistent_table"))
+        create_table(kusto_client, database, table_name, mapping)
+        return True
+    return False
+
+
+def _drop_by_tag(kusto_client: KustoClient, database: str, tag: str) -> None:
+    """
+    Drop all data with the specified tag.
+
+    Args:
+        kusto_client: The Kusto client
+        database: The database name
+        tag: The tag to drop data by
+    """
+    LOGGER.info(T("coal.services.adx.dropping_data_by_tag").format(tag=tag))
+
+    try:
+        # Execute the drop by tag command
+        drop_command = f'.drop extents <| .show database extents where tags has "drop-by:{tag}"'
+        kusto_client.execute_mgmt(database, drop_command)
+        LOGGER.info(T("coal.services.adx.drop_completed"))
+    except Exception as e:
+        LOGGER.error(T("coal.services.adx.drop_error").format(error=str(e)))
+        LOGGER.exception(T("coal.services.adx.drop_details"))
+
+
+def create_table(client: KustoClient, database: str, table_name: str, schema: Dict[str, str]) -> bool:
+    """
+    Create a table in the database.
+
+    Args:
+        client: The KustoClient to use
+        database: The name of the database
+        table_name: The name of the table to create
+        schema: Dictionary mapping column names to ADX types
+
+    Returns:
+        bool: True if the table was created successfully, False otherwise
+    """
+    LOGGER.debug(T("coal.services.adx.creating_table").format(database=database, table_name=table_name))
+
+    create_query = f".create-merge table {table_name}("
+
+    for column_name, column_type in schema.items():
+        create_query += f"{column_name}:{column_type},"
+
+    create_query = create_query[:-1] + ")"
+
+    LOGGER.debug(T("coal.services.adx.create_query").format(query=create_query))
+
+    try:
+        client.execute(database, create_query)
+        LOGGER.info(T("coal.services.adx.table_created").format(table_name=table_name))
+        return True
+    except Exception as e:
+        LOGGER.error(T("coal.services.adx.table_creation_error").format(table_name=table_name, error=str(e)))
+        return False

cosmotech/coal/azure/adx/utils.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.
+
+import dateutil.parser
+from typing import Any, Dict
+
+import pyarrow
+
+from cosmotech.coal.utils.logger import LOGGER
+from cosmotech.orchestrator.utils.translate import T
+
+
+def create_column_mapping(data: pyarrow.Table) -> Dict[str, str]:
+    """
+    Create a column mapping for a PyArrow table.
+
+    Args:
+        data: The PyArrow table data
+
+    Returns:
+        dict: A mapping of column names to their ADX types
+    """
+    mapping = dict()
+    for column_name in data.column_names:
+        column = data.column(column_name)
+        try:
+            ex = next(v for v in column.to_pylist() if v is not None)
+        except StopIteration:
+            LOGGER.error(T("coal.services.adx.empty_column").format(column_name=column_name))
+            mapping[column_name] = type_mapping(column_name, "string")
+            continue
+        else:
+            mapping[column_name] = type_mapping(column_name, ex)
+    return mapping
+
+
+def type_mapping(key: str, key_example_value: Any) -> str:
+    """
+    Map Python types to ADX types.
+
+    Args:
+        key: The name of the key
+        key_example_value: A possible value of the key
+
+    Returns:
+        str: The name of the type used in ADX
+    """
+    LOGGER.debug(T("coal.services.adx.mapping_type").format(key=key, value_type=type(key_example_value).__name__))
+
+    if key == "SimulationRun":
+        return "guid"
+
+    try:
+        # Use dateutil parser to test if the value could be a date, in case of error it is not
+        dateutil.parser.parse(key_example_value, fuzzy=False)
+        return "datetime"
+    except (ValueError, TypeError):
+        pass
+
+    if isinstance(key_example_value, float):
+        return "real"
+
+    if isinstance(key_example_value, int):
+        return "long"
+
+    # Default case to string
+    return "string"

cosmotech/coal/azure/blob.py

@@ -0,0 +1,109 @@
+# 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.
+
+"""
+Azure Blob Storage operations module.
+
+This module provides functions for interacting with Azure Blob Storage,
+including uploading data from the Store.
+"""
+
+import pathlib
+from io import BytesIO
+from typing import List, Optional
+
+from azure.identity import ClientSecretCredential
+from azure.storage.blob import BlobServiceClient
+
+import pyarrow.csv as pc
+import pyarrow.parquet as pq
+
+from cosmotech.coal.store.store import Store
+from cosmotech.coal.utils.logger import LOGGER
+from cosmotech.orchestrator.utils.translate import T
+
+VALID_TYPES = (
+    "sqlite",
+    "csv",
+    "parquet",
+)
+
+
+def dump_store_to_azure(
+    store_folder: str,
+    account_name: str,
+    container_name: str,
+    tenant_id: str,
+    client_id: str,
+    client_secret: str,
+    output_type: str = "sqlite",
+    file_prefix: str = "",
+) -> None:
+    """
+    Dump Store data to Azure Blob Storage.
+
+    Args:
+        store_folder: Folder containing the Store
+        account_name: Azure Storage account name
+        container_name: Azure Storage container name
+        tenant_id: Azure tenant ID
+        client_id: Azure client ID
+        client_secret: Azure client secret
+        output_type: Output file type (sqlite, csv, or parquet)
+        file_prefix: Prefix for uploaded files
+
+    Raises:
+        ValueError: If the output type is invalid
+    """
+    _s = Store(store_location=store_folder)
+
+    if output_type not in VALID_TYPES:
+        LOGGER.error(T("coal.common.validation.invalid_output_type").format(output_type=output_type))
+        raise ValueError(T("coal.common.validation.invalid_output_type").format(output_type=output_type))
+
+    container_client = BlobServiceClient(
+        account_url=f"https://{account_name}.blob.core.windows.net/",
+        credential=ClientSecretCredential(tenant_id=tenant_id, client_id=client_id, client_secret=client_secret),
+    ).get_container_client(container_name)
+
+    def data_upload(data_stream: BytesIO, file_name: str):
+        uploaded_file_name = file_prefix + file_name
+        data_stream.seek(0)
+        size = len(data_stream.read())
+        data_stream.seek(0)
+
+        LOGGER.info(T("coal.common.data_transfer.sending_data").format(size=size))
+        container_client.upload_blob(name=uploaded_file_name, data=data_stream, length=size, overwrite=True)
+
+    if output_type == "sqlite":
+        _file_path = _s._database_path
+        _file_name = "db.sqlite"
+        _uploaded_file_name = file_prefix + _file_name
+        LOGGER.info(
+            T("coal.common.data_transfer.file_sent").format(file_path=_file_path, uploaded_name=_uploaded_file_name)
+        )
+        with open(_file_path, "rb") as data:
+            container_client.upload_blob(name=_uploaded_file_name, data=data, overwrite=True)
+    else:
+        tables = list(_s.list_tables())
+        for table_name in tables:
+            _data_stream = BytesIO()
+            _file_name = None
+            _data = _s.get_table(table_name)
+            if not len(_data):
+                LOGGER.info(T("coal.common.data_transfer.table_empty").format(table_name=table_name))
+                continue
+            if output_type == "csv":
+                _file_name = table_name + ".csv"
+                pc.write_csv(_data, _data_stream)
+            elif output_type == "parquet":
+                _file_name = table_name + ".parquet"
+                pq.write_table(_data, _data_stream)
+            LOGGER.info(
+                T("coal.common.data_transfer.sending_table").format(table_name=table_name, output_type=output_type)
+            )
+            data_upload(_data_stream, _file_name)

cosmotech/coal/azure/functions.py

@@ -0,0 +1,72 @@
+# 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 azure.functions as func
+from cosmotech.coal.cosmotech_api.runner.download import download_runner_data
+from cosmotech_api.api.runner_api import RunnerApi
+
+import json
+import http
+import traceback
+
+
+def generate_main(apply_update, parallel=True):
+    def main(req: func.HttpRequest) -> func.HttpResponse:
+        try:
+            runner_id = req.params.get("scenario-id")  # Keep parameter name for backward compatibility
+            organization_id = req.params.get("organization-id")
+            workspace_id = req.params.get("workspace-id")
+            access_token: str = req.headers.get("authorization", None)
+            if access_token:
+                access_token = access_token.split(" ")[1]
+
+            if runner_id is None or organization_id is None or workspace_id is None:
+                return func.HttpResponse(
+                    body=f"Invalid request: organization-id={organization_id}, workspace-id={workspace_id}, scenario-id={runner_id}",
+                    status_code=http.HTTPStatus.BAD_REQUEST,
+                )
+
+            # Get runner data
+            result = download_runner_data(
+                organization_id=organization_id,
+                workspace_id=workspace_id,
+                runner_id=runner_id,
+                parameter_folder=None,  # We don't need to save to files
+                read_files=True,
+                parallel=parallel,
+                write_json=False,
+                write_csv=False,
+                fetch_dataset=True,
+            )
+
+            content = {
+                "datasets": result["datasets"],
+                "parameters": result["parameters"],
+            }
+
+            runner_data = result["runner_data"]
+
+            updated_content = apply_update(
+                content=content, scenario_data=runner_data
+            )  # Keep parameter name for backward compatibility
+
+            return func.HttpResponse(
+                body=json.dumps(updated_content),
+                headers={"Content-Type": "application/json"},
+            )
+        except Exception as e:
+            response = {
+                "error": getattr(e, "message", str(e)),
+                "type": type(e).__name__,
+                "trace": traceback.format_exc(),
+            }
+            return func.HttpResponse(
+                status_code=http.HTTPStatus.INTERNAL_SERVER_ERROR,
+                body=json.dumps(response),
+                headers={"Content-Type": "application/json"},
+            )
+
+    return main