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

cosmotech/coal/__init__.py

@@ -5,4 +5,4 @@
 # etc., to any person is prohibited unless it has been previously and
 # specifically authorized by written means by Cosmo Tech.
 
-__version__ = "1.1.0"
+__version__ = "2.1.0-rc1"

cosmotech/coal/aws/__init__.py

@@ -12,12 +12,4 @@ 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,
-)
+from cosmotech.coal.aws.s3 import S3

cosmotech/coal/aws/s3.py

@@ -6,230 +6,197 @@
 # specifically authorized by written means by Cosmo Tech.
 
 """
-S3 bucket operations module.
+s3 bucket operations module.
 
-This module provides functions for interacting with S3 buckets, including
+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.orchestrator.utils.translate import T
 
+from cosmotech.coal.utils.configuration import Configuration
 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
+class S3:
+
+    def __init__(self, configuration: Configuration):
+        self._configuration = configuration.s3
+
+    @property
+    def file_prefix(self):
+        if "bucket_prefix" in self._configuration:
+            return self._configuration.bucket_prefix
+        return ""
+
+    @property
+    def use_ssl(self):
+        if "use_ssl" in self._configuration:
+            return self._configuration.use_ssl
+        return True
+
+    @property
+    def ssl_cert_bundle(self):
+        if "ssl_cert_bundle" in self._configuration:
+            return self._configuration.ssl_cert_bundle
+        return None
+
+    @property
+    def access_key_id(self):
+        return self._configuration.access_key_id
+
+    @property
+    def endpoint_url(self):
+        return self._configuration.endpoint_url
+
+    @property
+    def bucket_name(self):
+        return self._configuration.bucket_name
+
+    @property
+    def secret_access_key(self):
+        return self._configuration.secret_access_key
+
+    @property
+    def output_type(self):
+        if "output_type" in self._configuration:
+            return self._configuration.output_type
+        return "csv"
+
+    @property
+    def client(self) -> boto3.client:
+        boto3_parameters = {
+            "use_ssl": self.use_ssl,
+            "endpoint_url": self.endpoint_url,
+            "aws_access_key_id": self.access_key_id,
+            "aws_secret_access_key": self.secret_access_key,
+        }
+        if self.ssl_cert_bundle:
+            boto3_parameters["verify"] = self.ssl_cert_bundle
+
+        return boto3.client("s3", **boto3_parameters)
+
+    @property
+    def resource(self) -> boto3.resource:
+        boto3_parameters = {
+            "use_ssl": self.use_ssl,
+            "endpoint_url": self.endpoint_url,
+            "aws_access_key_id": self.access_key_id,
+            "aws_secret_access_key": self.secret_access_key,
+        }
+        if self.ssl_cert_bundle:
+            boto3_parameters["verify"] = self.ssl_cert_bundle
+
+        return boto3.resource("s3", **boto3_parameters)
+
+    def upload_file(self, file_path: pathlib.Path) -> None:
+        """
+        Upload a single file to an S3 bucket.
+
+        Args:
+            file_path: Path to the file to upload
+        """
+        uploaded_file_name = self.file_prefix + file_path.name
+        LOGGER.info(
+            T("coal.common.data_transfer.file_sent").format(file_path=file_path, uploaded_name=uploaded_file_name)
+        )
+        self.resource.Bucket(self.bucket_name).upload_file(str(file_path), uploaded_file_name)
+
+    def upload_folder(self, source_folder: 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
+            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 = self.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"))
+                    self.resource.Bucket(self.bucket_name).upload_file(str(_file_path), uploaded_file_name)
+        else:
+            self.upload_file(source_path)
+
+    def download_files(self, destination_folder: str) -> None:
+        """
+        Download files from an S3 bucket to a local folder.
+
+        Args:
+            destination_folder: Local folder to download files to
+        """
+        bucket = self.resource.Bucket(self.bucket_name)
+
+        pathlib.Path(destination_folder).mkdir(parents=True, exist_ok=True)
+        remove_prefix = False
+        if self.file_prefix:
+            bucket_files = bucket.objects.filter(Prefix=self.file_prefix)
+            if self.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(self.file_prefix)
+                output_file = f"{destination_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(self, data_stream: BytesIO, file_name: str) -> None:
+        """
+        Upload a data stream to an S3 bucket.
+
+        Args:
+            data_stream: BytesIO stream containing the data to upload
+            file_name: Name of the file to create in the bucket
+        """
+        uploaded_file_name = self.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))
+        self.client.upload_fileobj(data_stream, self.bucket_name, uploaded_file_name)
+
+    def delete_objects(self) -> 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 = self.resource.Bucket(self.bucket_name)
+
+        if self.file_prefix:
+            bucket_files = bucket.objects.filter(Prefix=self.file_prefix)
+        else:
+            bucket_files = bucket.objects.all()
+
+        boto_objects = [{"Key": _file.key} for _file in bucket_files if _file.key != self.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

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

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

cosmotech/coal/azure/adx/auth.py

@@ -6,13 +6,13 @@
 # specifically authorized by written means by Cosmo Tech.
 
 import os
-from typing import Union, Optional, Tuple
+from typing import Optional, Tuple
 
 from azure.kusto.data import KustoClient, KustoConnectionStringBuilder
 from azure.kusto.ingest import QueuedIngestClient
+from cosmotech.orchestrator.utils.translate import T
 
 from cosmotech.coal.utils.logger import LOGGER
-from cosmotech.orchestrator.utils.translate import T
 
 
 def create_kusto_client(

cosmotech/coal/azure/adx/ingestion.py

@@ -5,28 +5,24 @@
 # etc., to any person is prohibited unless it has been previously and
 # specifically authorized by written means by Cosmo Tech.
 
+import os
+import time
 from enum import Enum
-from typing import Dict
-from typing import Iterator
-from typing import List
-from typing import Optional
-from typing import Tuple
+from typing import Dict, Iterator, List, Optional, Tuple
 
-import os
 import pandas as pd
-import time
 import tqdm
 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 azure.kusto.ingest.status import FailureMessage
-from azure.kusto.ingest.status import KustoIngestStatusQueues
-from azure.kusto.ingest.status import SuccessMessage
+from azure.kusto.ingest import IngestionProperties, QueuedIngestClient, ReportLevel
+from azure.kusto.ingest.status import (
+    FailureMessage,
+    KustoIngestStatusQueues,
+    SuccessMessage,
+)
 from cosmotech.orchestrator.utils.translate import T
 
-from cosmotech.coal.azure.adx.tables import create_table, _drop_by_tag
+from cosmotech.coal.azure.adx.tables import _drop_by_tag, create_table
 from cosmotech.coal.azure.adx.utils import type_mapping
 from cosmotech.coal.utils.logger import LOGGER
 

cosmotech/coal/azure/adx/query.py

@@ -7,9 +7,9 @@
 
 from azure.kusto.data import KustoClient
 from azure.kusto.data.response import KustoResponseDataSet
+from cosmotech.orchestrator.utils.translate import T
 
 from cosmotech.coal.utils.logger import LOGGER
-from cosmotech.orchestrator.utils.translate import T
 
 
 def run_query(client: KustoClient, database: str, query: str) -> KustoResponseDataSet:

cosmotech/coal/azure/adx/runner.py

@@ -13,25 +13,24 @@ This module provides functions for ingesting runner data into Azure Data Explore
 
 import pathlib
 import time
-from collections import defaultdict
-from typing import Dict, Any, List, Tuple, Optional
+from typing import Any, Dict
 
 from azure.kusto.data.response import KustoResponseDataSet
-from azure.kusto.ingest import ColumnMapping
-from azure.kusto.ingest import FileDescriptor
-from azure.kusto.ingest import IngestionMappingKind
-from azure.kusto.ingest import IngestionProperties
-from azure.kusto.ingest import IngestionResult
-from azure.kusto.ingest import ReportLevel
-
-from azure.kusto.data import KustoClient
-from azure.kusto.ingest import QueuedIngestClient
+from azure.kusto.ingest import (
+    ColumnMapping,
+    FileDescriptor,
+    IngestionMappingKind,
+    IngestionProperties,
+    IngestionResult,
+    QueuedIngestClient,
+    ReportLevel,
+)
+from cosmotech.orchestrator.utils.translate import T
 
 from cosmotech.coal.azure.adx.auth import initialize_clients
-from cosmotech.coal.azure.adx.query import run_query, run_command_query
-from cosmotech.coal.azure.adx.ingestion import check_ingestion_status, IngestionStatus
+from cosmotech.coal.azure.adx.ingestion import IngestionStatus, check_ingestion_status
+from cosmotech.coal.azure.adx.query import run_query
 from cosmotech.coal.utils.logger import LOGGER
-from cosmotech.orchestrator.utils.translate import T
 
 
 def prepare_csv_content(folder_path: str) -> Dict[str, Dict[str, Any]]: