cosmotech-acceleration-library 1.1.0__py3-none-any.whl → 2.0.0__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.0.0"

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/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/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]]:

cosmotech/coal/azure/adx/store.py

@@ -7,26 +7,22 @@
 
 import os
 import tempfile
+import time
 import uuid
-from typing import Optional, List, Dict, Tuple, Union, Any
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 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 azure.kusto.ingest import IngestionProperties, QueuedIngestClient, 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.azure.adx.ingestion import handle_failures, monitor_ingestion
+from cosmotech.coal.azure.adx.tables import _drop_by_tag, check_and_create_table
 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(
@@ -176,80 +172,3 @@ def send_store_to_adx(
         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

@@ -5,13 +5,13 @@
 # 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
+from typing import Dict
 
 import pyarrow
 from azure.kusto.data import KustoClient
+from cosmotech.orchestrator.utils.translate import T
 
 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:

cosmotech/coal/azure/blob.py

@@ -12,19 +12,16 @@ 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 azure.identity import ClientSecretCredential
+from azure.storage.blob import BlobServiceClient
+from cosmotech.orchestrator.utils.translate import T
 
 from cosmotech.coal.store.store import Store
 from cosmotech.coal.utils.logger import LOGGER
-from cosmotech.orchestrator.utils.translate import T
 
 VALID_TYPES = (
     "sqlite",
@@ -42,6 +39,7 @@ def dump_store_to_azure(
     client_secret: str,
     output_type: str = "sqlite",
     file_prefix: str = "",
+    selected_tables: list[str] = [],
 ) -> None:
     """
     Dump Store data to Azure Blob Storage.
@@ -90,6 +88,8 @@ def dump_store_to_azure(
             container_client.upload_blob(name=_uploaded_file_name, data=data, overwrite=True)
     else:
         tables = list(_s.list_tables())
+        if selected_tables:
+            tables = [t for t in tables if t in selected_tables]
         for table_name in tables:
             _data_stream = BytesIO()
             _file_name = None