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

cosmotech/coal/__init__.py

@@ -0,0 +1,8 @@
+# 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.
+
+__version__ = "1.0.0"

cosmotech/coal/aws/__init__.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.
+
+"""
+AWS services integration module.
+
+This module provides functions for interacting with AWS services like S3.
+"""
+
+# Re-export S3 functions for easier importing
+from cosmotech.coal.aws.s3 import (
+    create_s3_client,
+    create_s3_resource,
+    upload_file,
+    upload_folder,
+    download_files,
+    upload_data_stream,
+    delete_objects,
+)

cosmotech/coal/aws/s3.py

@@ -0,0 +1,235 @@
+# 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.
+
+"""
+S3 bucket operations module.
+
+This module provides functions for interacting with S3 buckets, including
+uploading, downloading, and deleting files.
+"""
+
+import pathlib
+from io import BytesIO
+from typing import Optional, Dict, Any, List, Iterator
+
+import boto3
+
+from cosmotech.coal.utils.logger import LOGGER
+from cosmotech.orchestrator.utils.translate import T
+
+
+def create_s3_client(
+    endpoint_url: str,
+    access_id: str,
+    secret_key: str,
+    use_ssl: bool = True,
+    ssl_cert_bundle: Optional[str] = None,
+) -> boto3.client:
+    """
+    Create an S3 client with the given credentials and configuration.
+
+    Args:
+        endpoint_url: The S3 endpoint URL
+        access_id: The AWS access key ID
+        secret_key: The AWS secret access key
+        use_ssl: Whether to use SSL for the connection
+        ssl_cert_bundle: Path to the SSL certificate bundle
+
+    Returns:
+        An S3 client object
+    """
+    boto3_parameters = {
+        "use_ssl": use_ssl,
+        "endpoint_url": endpoint_url,
+        "aws_access_key_id": access_id,
+        "aws_secret_access_key": secret_key,
+    }
+    if ssl_cert_bundle:
+        boto3_parameters["verify"] = ssl_cert_bundle
+
+    return boto3.client("s3", **boto3_parameters)
+
+
+def create_s3_resource(
+    endpoint_url: str,
+    access_id: str,
+    secret_key: str,
+    use_ssl: bool = True,
+    ssl_cert_bundle: Optional[str] = None,
+) -> boto3.resource:
+    """
+    Create an S3 resource with the given credentials and configuration.
+
+    Args:
+        endpoint_url: The S3 endpoint URL
+        access_id: The AWS access key ID
+        secret_key: The AWS secret access key
+        use_ssl: Whether to use SSL for the connection
+        ssl_cert_bundle: Path to the SSL certificate bundle
+
+    Returns:
+        An S3 resource object
+    """
+    boto3_parameters = {
+        "use_ssl": use_ssl,
+        "endpoint_url": endpoint_url,
+        "aws_access_key_id": access_id,
+        "aws_secret_access_key": secret_key,
+    }
+    if ssl_cert_bundle:
+        boto3_parameters["verify"] = ssl_cert_bundle
+
+    return boto3.resource("s3", **boto3_parameters)
+
+
+def upload_file(
+    file_path: pathlib.Path,
+    bucket_name: str,
+    s3_resource: boto3.resource,
+    file_prefix: str = "",
+) -> None:
+    """
+    Upload a single file to an S3 bucket.
+
+    Args:
+        file_path: Path to the file to upload
+        bucket_name: Name of the S3 bucket
+        s3_resource: S3 resource object
+        file_prefix: Prefix to add to the file name in the bucket
+    """
+    uploaded_file_name = file_prefix + file_path.name
+    LOGGER.info(T("coal.common.data_transfer.file_sent").format(file_path=file_path, uploaded_name=uploaded_file_name))
+    s3_resource.Bucket(bucket_name).upload_file(str(file_path), uploaded_file_name)
+
+
+def upload_folder(
+    source_folder: str,
+    bucket_name: str,
+    s3_resource: boto3.resource,
+    file_prefix: str = "",
+    recursive: bool = False,
+) -> None:
+    """
+    Upload files from a folder to an S3 bucket.
+
+    Args:
+        source_folder: Path to the folder containing files to upload
+        bucket_name: Name of the S3 bucket
+        s3_resource: S3 resource object
+        file_prefix: Prefix to add to the file names in the bucket
+        recursive: Whether to recursively upload files from subdirectories
+    """
+    source_path = pathlib.Path(source_folder)
+    if not source_path.exists():
+        LOGGER.error(T("coal.common.file_operations.not_found").format(source_folder=source_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)
+        for _file_path in source_path.glob("**/*" if recursive else "*"):
+            if _file_path.is_file():
+                _file_name = str(_file_path).removeprefix(_source_name).removeprefix("/")
+                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
+                    )
+                )
+                s3_resource.Bucket(bucket_name).upload_file(str(_file_path), uploaded_file_name)
+    else:
+        upload_file(source_path, bucket_name, s3_resource, file_prefix)
+
+
+def download_files(
+    target_folder: str,
+    bucket_name: str,
+    s3_resource: boto3.resource,
+    file_prefix: Optional[str] = None,
+) -> None:
+    """
+    Download files from an S3 bucket to a local folder.
+
+    Args:
+        target_folder: Local folder to download files to
+        bucket_name: Name of the S3 bucket
+        s3_resource: S3 resource object
+        file_prefix: Optional prefix to filter objects to download
+    """
+    bucket = s3_resource.Bucket(bucket_name)
+
+    pathlib.Path(target_folder).mkdir(parents=True, exist_ok=True)
+    remove_prefix = False
+    if file_prefix:
+        bucket_files = bucket.objects.filter(Prefix=file_prefix)
+        if file_prefix.endswith("/"):
+            remove_prefix = True
+    else:
+        bucket_files = bucket.objects.all()
+    for _file in bucket_files:
+        if not (path_name := str(_file.key)).endswith("/"):
+            target_file = path_name
+            if remove_prefix:
+                target_file = target_file.removeprefix(file_prefix)
+            output_file = f"{target_folder}/{target_file}"
+            pathlib.Path(output_file).parent.mkdir(parents=True, exist_ok=True)
+            LOGGER.info(T("coal.services.azure_storage.downloading").format(path=path_name, output=output_file))
+            bucket.download_file(_file.key, output_file)
+
+
+def upload_data_stream(
+    data_stream: BytesIO,
+    bucket_name: str,
+    s3_client: boto3.client,
+    file_name: str,
+    file_prefix: str = "",
+) -> None:
+    """
+    Upload a data stream to an S3 bucket.
+
+    Args:
+        data_stream: BytesIO stream containing the data to upload
+        bucket_name: Name of the S3 bucket
+        s3_client: S3 client object
+        file_name: Name of the file to create in the bucket
+        file_prefix: Prefix to add to the file name in the bucket
+    """
+    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))
+    s3_client.upload_fileobj(data_stream, bucket_name, uploaded_file_name)
+
+
+def delete_objects(
+    bucket_name: str,
+    s3_resource: boto3.resource,
+    file_prefix: Optional[str] = None,
+) -> None:
+    """
+    Delete objects from an S3 bucket, optionally filtered by prefix.
+
+    Args:
+        bucket_name: Name of the S3 bucket
+        s3_resource: S3 resource object
+        file_prefix: Optional prefix to filter objects to delete
+    """
+    bucket = s3_resource.Bucket(bucket_name)
+
+    if file_prefix:
+        bucket_files = bucket.objects.filter(Prefix=file_prefix)
+    else:
+        bucket_files = bucket.objects.all()
+
+    boto_objects = [{"Key": _file.key} for _file in bucket_files if _file.key != file_prefix]
+    if boto_objects:
+        LOGGER.info(T("coal.services.azure_storage.deleting_objects").format(objects=boto_objects))
+        boto_delete_request = {"Objects": boto_objects}
+        bucket.delete_objects(Delete=boto_delete_request)
+    else:
+        LOGGER.info(T("coal.services.azure_storage.no_objects"))

cosmotech/coal/azure/__init__.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.
+
+"""
+Azure services integration module.
+
+This module provides functions for interacting with Azure services like Storage and ADX.
+"""
+
+# Re-export storage functions for easier importing
+from cosmotech.coal.azure.storage import (
+    upload_file,
+    upload_folder,
+)
+
+# Re-export blob functions for easier importing
+from cosmotech.coal.azure.blob import (
+    dump_store_to_azure,
+)

cosmotech/coal/azure/adx/__init__.py

@@ -0,0 +1,26 @@
+# 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.azure.adx.auth import create_kusto_client, create_ingest_client, initialize_clients
+from cosmotech.coal.azure.adx.query import run_query, run_command_query
+from cosmotech.coal.azure.adx.ingestion import (
+    ingest_dataframe,
+    send_to_adx,
+    check_ingestion_status,
+    monitor_ingestion,
+    handle_failures,
+    IngestionStatus,
+)
+from cosmotech.coal.azure.adx.tables import table_exists, create_table, check_and_create_table, _drop_by_tag
+from cosmotech.coal.azure.adx.utils import type_mapping, create_column_mapping
+from cosmotech.coal.azure.adx.store import send_pyarrow_table_to_adx, send_table_data, process_tables, send_store_to_adx
+from cosmotech.coal.azure.adx.runner import (
+    prepare_csv_content,
+    construct_create_query,
+    insert_csv_files,
+    send_runner_data,
+)

cosmotech/coal/azure/adx/auth.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.
+
+import os
+from typing import Union, Optional, Tuple
+
+from azure.kusto.data import KustoClient, KustoConnectionStringBuilder
+from azure.kusto.ingest import QueuedIngestClient
+
+from cosmotech.coal.utils.logger import LOGGER
+from cosmotech.orchestrator.utils.translate import T
+
+
+def create_kusto_client(
+    cluster_url: str,
+    client_id: Optional[str] = None,
+    client_secret: Optional[str] = None,
+    tenant_id: Optional[str] = None,
+) -> KustoClient:
+    """
+    Create a KustoClient for querying ADX.
+
+    Args:
+        cluster_url: The URL of the ADX cluster
+        client_id: Azure client ID (optional, will use environment variable if not provided)
+        client_secret: Azure client secret (optional, will use environment variable if not provided)
+        tenant_id: Azure tenant ID (optional, will use environment variable if not provided)
+
+    Returns:
+        KustoClient: A client for querying ADX
+    """
+    LOGGER.debug(T("coal.services.adx.creating_kusto_client").format(cluster_url=cluster_url))
+
+    try:
+        az_client_id = client_id or os.environ["AZURE_CLIENT_ID"]
+        az_client_secret = client_secret or os.environ["AZURE_CLIENT_SECRET"]
+        az_tenant_id = tenant_id or os.environ["AZURE_TENANT_ID"]
+
+        kcsb = KustoConnectionStringBuilder.with_aad_application_key_authentication(
+            cluster_url, az_client_id, az_client_secret, az_tenant_id
+        )
+        LOGGER.debug(T("coal.services.adx.using_app_auth"))
+    except KeyError:
+        LOGGER.debug(T("coal.services.adx.using_cli_auth"))
+        kcsb = KustoConnectionStringBuilder.with_az_cli_authentication(cluster_url)
+
+    return KustoClient(kcsb)
+
+
+def create_ingest_client(
+    ingest_url: str,
+    client_id: Optional[str] = None,
+    client_secret: Optional[str] = None,
+    tenant_id: Optional[str] = None,
+) -> QueuedIngestClient:
+    """
+    Create a QueuedIngestClient for ingesting data to ADX.
+
+    Args:
+        ingest_url: The ingestion URL of the ADX cluster
+        client_id: Azure client ID (optional, will use environment variable if not provided)
+        client_secret: Azure client secret (optional, will use environment variable if not provided)
+        tenant_id: Azure tenant ID (optional, will use environment variable if not provided)
+
+    Returns:
+        QueuedIngestClient: A client for ingesting data to ADX
+    """
+    LOGGER.debug(T("coal.services.adx.creating_ingest_client").format(ingest_url=ingest_url))
+
+    try:
+        az_client_id = client_id or os.environ["AZURE_CLIENT_ID"]
+        az_client_secret = client_secret or os.environ["AZURE_CLIENT_SECRET"]
+        az_tenant_id = tenant_id or os.environ["AZURE_TENANT_ID"]
+
+        kcsb = KustoConnectionStringBuilder.with_aad_application_key_authentication(
+            ingest_url, az_client_id, az_client_secret, az_tenant_id
+        )
+        LOGGER.debug(T("coal.services.adx.using_app_auth"))
+    except KeyError:
+        LOGGER.debug(T("coal.services.adx.using_cli_auth"))
+        kcsb = KustoConnectionStringBuilder.with_az_cli_authentication(ingest_url)
+
+    return QueuedIngestClient(kcsb)
+
+
+def initialize_clients(adx_uri: str, adx_ingest_uri: str) -> Tuple[KustoClient, QueuedIngestClient]:
+    """
+    Initialize and return the Kusto and ingest clients.
+
+    Args:
+        adx_uri: The Azure Data Explorer resource URI
+        adx_ingest_uri: The Azure Data Explorer resource ingest URI
+
+    Returns:
+        tuple: (kusto_client, ingest_client)
+    """
+    LOGGER.debug(T("coal.services.adx.initializing_clients"))
+    kusto_client = create_kusto_client(adx_uri)
+    ingest_client = create_ingest_client(adx_ingest_uri)
+    return kusto_client, ingest_client
+
+
+def get_cluster_urls(cluster_name: str, cluster_region: str) -> Tuple[str, str]:
+    """
+    Generate cluster and ingest URLs from cluster name and region.
+
+    Args:
+        cluster_name: The name of the ADX cluster
+        cluster_region: The region of the ADX cluster
+
+    Returns:
+        tuple: (cluster_url, ingest_url)
+    """
+    LOGGER.debug(
+        T("coal.services.adx.generating_urls").format(cluster_name=cluster_name, cluster_region=cluster_region)
+    )
+
+    cluster_url = f"https://{cluster_name}.{cluster_region}.kusto.windows.net"
+    ingest_url = f"https://ingest-{cluster_name}.{cluster_region}.kusto.windows.net"
+
+    return cluster_url, ingest_url