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

cosmotech/coal/postgresql/store.py

@@ -13,12 +13,13 @@ for store operations.
 """
 
 from time import perf_counter
-import pyarrow
 
+from cosmotech.orchestrator.utils.translate import T
+
+from cosmotech.coal.postgresql.utils import PostgresUtils
 from cosmotech.coal.store.store import Store
+from cosmotech.coal.utils.configuration import Configuration
 from cosmotech.coal.utils.logger import LOGGER
-from cosmotech.coal.utils.postgresql import send_pyarrow_table_to_postgresql
-from cosmotech.orchestrator.utils.translate import T
 
 
 def dump_store_to_postgresql(
@@ -32,6 +33,8 @@ def dump_store_to_postgresql(
     table_prefix: str = "Cosmotech_",
     replace: bool = True,
     force_encode: bool = False,
+    selected_tables: list[str] = [],
+    fk_id: str = None,
 ) -> None:
     """
     Dump Store data to a PostgreSQL database.
@@ -46,36 +49,74 @@ def dump_store_to_postgresql(
         postgres_password: PostgreSQL password
         table_prefix: Table prefix
         replace: Whether to replace existing tables
-        force_encode: force password encoding
+        force_encode: force password encoding to percent encoding
+        selected_tables: list of tables to send
+        fk_id: foreign key id to add to all table on all rows
     """
-    _s = Store(store_location=store_folder)
+    _c = Configuration(
+        {
+            "coal": {"store": store_folder},
+            "postgres": {
+                "host": postgres_host,
+                "port": postgres_port,
+                "db_name": postgres_db,
+                "db_schema": postgres_schema,
+                "user_name": postgres_user,
+                "user_password": postgres_password,
+                "password_encoding": force_encode,
+                "table_prefix": table_prefix,
+            },
+        }
+    )
+
+    dump_store_to_postgresql_from_conf(configuration=_c, replace=replace, selected_tables=selected_tables, fk_id=fk_id)
+
+
+def dump_store_to_postgresql_from_conf(
+    configuration: Configuration,
+    replace: bool = True,
+    selected_tables: list[str] = [],
+    fk_id: str = None,
+) -> None:
+    """
+    Dump Store data to a PostgreSQL database.
+
+    Args:
+        configuration: coal Configuration
+        replace: Whether to replace existing tables
+        selected_tables: list of tables to send
+        fk_id: foreign key id to add to all table on all rows
+    """
+    _psql = PostgresUtils(configuration)
+    _s = Store(configuration=configuration)
 
     tables = list(_s.list_tables())
+    if selected_tables:
+        tables = [t for t in tables if t in selected_tables]
     if len(tables):
-        LOGGER.info(T("coal.services.database.sending_data").format(table=f"{postgres_db}.{postgres_schema}"))
+        LOGGER.info(T("coal.services.database.sending_data").format(table=f"{_psql.db_name}.{_psql.db_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}"
+            target_table_name = f"{_psql.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
+            if fk_id:
+                data = data.append_column("csm_run_id", [[fk_id] * data.num_rows])
             _dl_time = perf_counter()
-            rows = send_pyarrow_table_to_postgresql(
+            rows = _psql.send_pyarrow_table_to_postgresql(
                 data,
                 target_table_name,
-                postgres_host,
-                postgres_port,
-                postgres_db,
-                postgres_schema,
-                postgres_user,
-                postgres_password,
                 replace,
-                force_encode,
             )
+            if fk_id and _psql.is_metadata_exists():
+                metadata_table = f"{_psql.table_prefix}RunnerMetadata"
+                _psql.add_fk_constraint(table_name, "csm_run_id", metadata_table, "last_csm_run_id")
+
             total_rows += rows
             _up_time = perf_counter()
             LOGGER.info(T("coal.services.database.row_count").format(count=rows))

cosmotech/coal/postgresql/utils.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.
+
+from typing import Optional
+from urllib.parse import quote
+
+import adbc_driver_manager
+import pyarrow as pa
+from adbc_driver_postgresql import dbapi
+from cosmotech.orchestrator.utils.translate import T
+from pyarrow import Table
+
+from cosmotech.coal.utils.configuration import Configuration
+from cosmotech.coal.utils.logger import LOGGER
+
+
+class PostgresUtils:
+
+    def __init__(self, configuration: Configuration):
+        self._configuration = configuration.postgres
+
+    @property
+    def table_prefix(self):
+        if "table_prefix" in self._configuration:
+            return self._configuration.table_prefix
+        return "Cosmotech_"
+
+    @property
+    def db_name(self):
+        return self._configuration.db_name
+
+    @property
+    def db_schema(self):
+        return self._configuration.db_schema
+
+    @property
+    def host_uri(self):
+        return self._configuration.host
+
+    @property
+    def host_port(self):
+        return self._configuration.port
+
+    @property
+    def user_name(self):
+        return self._configuration.user_name
+
+    @property
+    def user_password(self):
+        return self._configuration.user_password
+
+    @property
+    def password_encoding(self):
+        if "password_encoding" in self._configuration:
+            return self._configuration.password_encoding
+        return False
+
+    @property
+    def full_uri(self) -> str:
+        # Check if password needs percent encoding (contains special characters)
+        # We don't log anything about the password for security
+        encoded_password = self.user_password
+        if self.password_encoding:
+            encoded_password = quote(self.user_password, safe="")
+
+        return (
+            "postgresql://" + f"{self.user_name}"
+            f":{encoded_password}"
+            f"@{self.host_uri}"
+            f":{self.host_port}"
+            f"/{self.db_name}"
+        )
+
+    @property
+    def metadata_table_name(self) -> str:
+        return f"{self.table_prefix}RunnerMetadata"
+
+    def get_postgresql_table_schema(self, target_table_name: str) -> Optional[pa.Schema]:
+        """
+        Get the schema of an existing PostgreSQL table using SQL queries.
+
+        Args:
+            target_table_name: Name of the table
+
+        Returns:
+            PyArrow Schema if table exists, None otherwise
+        """
+        LOGGER.debug(
+            T("coal.services.postgresql.getting_schema").format(
+                postgres_schema=self.db_schema, target_table_name=target_table_name
+            )
+        )
+
+        with dbapi.connect(self.full_uri) as conn:
+            try:
+                return conn.adbc_get_table_schema(
+                    target_table_name,
+                    db_schema_filter=self.db_schema,
+                )
+            except adbc_driver_manager.ProgrammingError:
+                LOGGER.warning(
+                    T("coal.services.postgresql.table_not_found").format(
+                        postgres_schema=self.db_schema, target_table_name=target_table_name
+                    )
+                )
+            return None
+
+    def send_pyarrow_table_to_postgresql(
+        self,
+        data: Table,
+        target_table_name: str,
+        replace: bool,
+    ) -> int:
+        LOGGER.debug(
+            T("coal.services.postgresql.preparing_send").format(
+                postgres_schema=self.db_schema, target_table_name=target_table_name
+            )
+        )
+        LOGGER.debug(T("coal.services.postgresql.input_rows").format(rows=len(data)))
+
+        # Get existing schema if table exists
+        existing_schema = self.get_postgresql_table_schema(target_table_name)
+
+        if existing_schema is not None:
+            LOGGER.debug(T("coal.services.postgresql.found_existing_table").format(schema=existing_schema))
+            if not replace:
+                LOGGER.debug(T("coal.services.postgresql.adapting_data"))
+                data = adapt_table_to_schema(data, existing_schema)
+            else:
+                LOGGER.debug(T("coal.services.postgresql.replace_mode"))
+        else:
+            LOGGER.debug(T("coal.services.postgresql.no_existing_table"))
+
+        # Proceed with ingestion
+        total = 0
+
+        LOGGER.debug(T("coal.services.postgresql.connecting"))
+        with dbapi.connect(self.full_uri, autocommit=True) as conn:
+            with conn.cursor() as curs:
+                mode = "replace" if replace else "create_append"
+                LOGGER.debug(T("coal.services.postgresql.ingesting_data").format(mode=mode))
+                total += curs.adbc_ingest(target_table_name, data, mode, db_schema_name=self.db_schema)
+
+        LOGGER.debug(T("coal.services.postgresql.ingestion_success").format(rows=total))
+        return total
+
+    def add_fk_constraint(
+        self,
+        from_table: str,
+        from_col: str,
+        to_table: str,
+        to_col: str,
+    ) -> None:
+        # Connect to PostgreSQL and remove runner metadata row
+        with dbapi.connect(self.full_uri, autocommit=True) as conn:
+            with conn.cursor() as curs:
+                sql_add_fk = f"""
+                    ALTER TABLE {self.db_schema}.{from_table}
+                    CONSTRAINT metadata FOREIGN KEY ({from_col}) REFERENCES {to_table}({to_col})
+                """
+                curs.execute(sql_add_fk)
+                conn.commit()
+
+    def is_metadata_exists(self) -> None:
+        with dbapi.connect(self.full_uri, autocommit=True) as conn:
+            try:
+                conn.adbc_get_table_schema(
+                    self.metadata_table_name,
+                    db_schema_filter=self.db_schema,
+                )
+                return True
+            except adbc_driver_manager.ProgrammingError:
+                return False
+
+
+def adapt_table_to_schema(data: pa.Table, target_schema: pa.Schema) -> pa.Table:
+    """
+    Adapt a PyArrow table to match a target schema with detailed logging.
+    """
+    LOGGER.debug(T("coal.services.postgresql.schema_adaptation_start").format(rows=len(data)))
+    LOGGER.debug(T("coal.services.postgresql.original_schema").format(schema=data.schema))
+    LOGGER.debug(T("coal.services.postgresql.target_schema").format(schema=target_schema))
+
+    target_fields = {field.name: field.type for field in target_schema}
+    new_columns = []
+
+    # Track adaptations for summary
+    added_columns = []
+    dropped_columns = []
+    type_conversions = []
+    failed_conversions = []
+
+    # Process each field in target schema
+    for field_name, target_type in target_fields.items():
+        if field_name in data.column_names:
+            # Column exists - try to cast to target type
+            col = data[field_name]
+            original_type = col.type
+
+            if original_type != target_type:
+                LOGGER.debug(
+                    T("coal.services.postgresql.casting_column").format(
+                        field_name=field_name,
+                        original_type=original_type,
+                        target_type=target_type,
+                    )
+                )
+                try:
+                    new_col = pa.compute.cast(col, target_type)
+                    new_columns.append(new_col)
+                    type_conversions.append(f"{field_name}: {original_type} -> {target_type}")
+                except pa.ArrowInvalid as e:
+                    LOGGER.warning(
+                        T("coal.services.postgresql.cast_failed").format(
+                            field_name=field_name,
+                            original_type=original_type,
+                            target_type=target_type,
+                            error=str(e),
+                        )
+                    )
+                    new_columns.append(pa.nulls(len(data), type=target_type))
+                    failed_conversions.append(f"{field_name}: {original_type} -> {target_type}")
+            else:
+                new_columns.append(col)
+        else:
+            # Column doesn't exist - add nulls
+            LOGGER.debug(T("coal.services.postgresql.adding_missing_column").format(field_name=field_name))
+            new_columns.append(pa.nulls(len(data), type=target_type))
+            added_columns.append(field_name)
+
+    # Log columns that will be dropped
+    dropped_columns = [name for name in data.column_names if name not in target_fields]
+    if dropped_columns:
+        LOGGER.debug(T("coal.services.postgresql.dropping_columns").format(columns=dropped_columns))
+
+    # Create new table
+    adapted_table = pa.Table.from_arrays(new_columns, schema=target_schema)
+
+    # Log summary of adaptations
+    LOGGER.debug(T("coal.services.postgresql.adaptation_summary"))
+    if added_columns:
+        LOGGER.debug(T("coal.services.postgresql.added_columns").format(columns=added_columns))
+    if dropped_columns:
+        LOGGER.debug(T("coal.services.postgresql.dropped_columns").format(columns=dropped_columns))
+    if type_conversions:
+        LOGGER.debug(T("coal.services.postgresql.successful_conversions").format(conversions=type_conversions))
+    if failed_conversions:
+        LOGGER.debug(T("coal.services.postgresql.failed_conversions").format(conversions=failed_conversions))
+
+    LOGGER.debug(T("coal.services.postgresql.final_schema").format(schema=adapted_table.schema))
+    return adapted_table

cosmotech/coal/singlestore/store.py

@@ -12,15 +12,16 @@ This module provides functions for interacting with SingleStore databases
 for store operations.
 """
 
+import csv
 import pathlib
 import time
-import csv
+
 import singlestoredb as s2
+from cosmotech.orchestrator.utils.translate import T
 
 from cosmotech.coal.store.csv import store_csv_file
 from cosmotech.coal.store.store import Store
 from cosmotech.coal.utils.logger import LOGGER
-from cosmotech.orchestrator.utils.translate import T
 
 
 def _get_data(table_name: str, output_directory: str, cursor) -> None:

cosmotech/coal/store/__init__.py

@@ -12,31 +12,34 @@ This module provides functions for working with the Store,
 including loading and converting data.
 """
 
-# Re-export the Store class
-from cosmotech.coal.store.store import Store
-
 # Re-export functions from the csv module
 from cosmotech.coal.store.csv import (
-    store_csv_file,
     convert_store_table_to_csv,
+    store_csv_file,
 )
 
 # Re-export functions from the native_python module
 from cosmotech.coal.store.native_python import (
-    store_pylist,
     convert_table_as_pylist,
+    store_pylist,
 )
-
-# Re-export functions from the pandas module (if available)
-
 from cosmotech.coal.store.pandas import (
-    store_dataframe,
     convert_store_table_to_dataframe as convert_store_table_to_pandas_dataframe,
 )
-
-# Re-export functions from the pyarrow module (if available)
-
+from cosmotech.coal.store.pandas import (
+    store_dataframe,
+)
 from cosmotech.coal.store.pyarrow import (
-    store_table,
     convert_store_table_to_dataframe as convert_store_table_to_pyarrow_table,
 )
+from cosmotech.coal.store.pyarrow import (
+    store_table,
+)
+
+# Re-export the Store class
+from cosmotech.coal.store.store import Store
+
+# Re-export functions from the pandas module (if available)
+
+
+# Re-export functions from the pyarrow module (if available)

cosmotech/coal/store/output/aws_channel.py

@@ -0,0 +1,74 @@
+from io import BytesIO
+from typing import Optional
+
+import pyarrow.csv as pc
+import pyarrow.parquet as pq
+from cosmotech.orchestrator.utils.translate import T
+
+from cosmotech.coal.aws import S3
+from cosmotech.coal.store.output.channel_interface import (
+    ChannelInterface,
+    MissingChannelConfigError,
+)
+from cosmotech.coal.store.store import Store
+from cosmotech.coal.utils.configuration import Configuration, Dotdict
+from cosmotech.coal.utils.logger import LOGGER
+
+
+class AwsChannel(ChannelInterface):
+    required_keys = {
+        "coal": ["store"],
+        "s3": ["access_key_id", "endpoint_url", "secret_access_key"],
+    }
+    requirement_string = required_keys
+
+    def __init__(self, dct: Dotdict = None):
+        super().__init__(dct)
+        self._s3 = S3(self.configuration)
+
+    def send(self, filter: Optional[list[str]] = None) -> bool:
+
+        _s = Store(configuration=self.configuration)
+
+        if self._s3.output_type not in ("sqlite", "csv", "parquet"):
+            LOGGER.error(T("coal.common.errors.data_invalid_output_type").format(output_type=self._s3.output_type))
+            raise ValueError(T("coal.common.errors.data_invalid_output_type").format(output_type=self._s3.output_type))
+
+        if self._s3.output_type == "sqlite":
+            _file_path = _s._database_path
+            _file_name = "db.sqlite"
+            _uploaded_file_name = self.configuration.s3.bucket_prefix + _file_name
+            LOGGER.info(
+                T("coal.common.data_transfer.file_sent").format(file_path=_file_path, uploaded_name=_uploaded_file_name)
+            )
+            self._s3.upload_file(_file_path, _uploaded_file_name)
+        else:
+            tables = list(_s.list_tables())
+            if filter:
+                tables = [t for t in tables if t in filter]
+
+            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 self._s3.output_type == "csv":
+                    _file_name = table_name + ".csv"
+                    pc.write_csv(_data, _data_stream)
+                elif self._s3.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=self._s3.output_type
+                    )
+                )
+                self._s3.upload_data_stream(
+                    data_stream=_data_stream,
+                    file_name=_file_name,
+                )
+
+    def delete(self):
+        self._s3.delete_objects()

cosmotech/coal/store/output/az_storage_channel.py

@@ -0,0 +1,33 @@
+from typing import Optional
+
+from cosmotech.coal.azure.blob import dump_store_to_azure
+from cosmotech.coal.store.output.channel_interface import (
+    ChannelInterface,
+    MissingChannelConfigError,
+)
+from cosmotech.coal.utils.configuration import Configuration, Dotdict
+
+
+class AzureStorageChannel(ChannelInterface):
+    required_keys = {
+        "coal": ["store"],
+        "azure": [
+            "account_name",
+            "container_name",
+            "tenant_id",
+            "client_id",
+            "client_secret",
+            "output_type",
+            "file_prefix",
+        ],
+    }
+    requirement_string = required_keys
+
+    def send(self, filter: Optional[list[str]] = None) -> bool:
+        dump_store_to_azure(
+            self.configuration,
+            selected_tables=filter,
+        )
+
+    def delete(self):
+        pass

cosmotech/coal/store/output/channel_interface.py

@@ -0,0 +1,38 @@
+from typing import Optional
+
+from cosmotech.orchestrator.utils.translate import T
+
+from cosmotech.coal.utils.configuration import Configuration, Dotdict
+
+
+class ChannelInterface:
+    required_keys = {}
+    requirement_string: str = T("coal.store.output.data_interface.requirements")
+
+    def __init__(self, dct: Dotdict = None):
+        self.configuration = Configuration(dct)
+        if not self.is_available():
+            raise MissingChannelConfigError(self)
+
+    def send(self, filter: Optional[list[str]] = None) -> bool:
+        raise NotImplementedError()
+
+    def delete(self) -> bool:
+        raise NotImplementedError()
+
+    def is_available(self) -> bool:
+        try:
+            return all(
+                all(key in self.configuration[section] for key in self.required_keys[section])
+                for section in self.required_keys.keys()
+            )
+        except KeyError:
+            return False
+
+
+class MissingChannelConfigError(Exception):
+    def __init__(self, interface_class):
+        self.message = T("coal.store.output.split.requirements").format(
+            interface_name=interface_class.__class__.__name__, requirements=interface_class.requirement_string
+        )
+        super().__init__(self.message)

cosmotech/coal/store/output/channel_spliter.py

@@ -0,0 +1,61 @@
+from typing import Optional
+
+from cosmotech.orchestrator.utils.translate import T
+
+from cosmotech.coal.store.output.aws_channel import AwsChannel
+from cosmotech.coal.store.output.az_storage_channel import AzureStorageChannel
+from cosmotech.coal.store.output.channel_interface import ChannelInterface
+from cosmotech.coal.store.output.postgres_channel import PostgresChannel
+from cosmotech.coal.utils.configuration import Dotdict
+from cosmotech.coal.utils.logger import LOGGER
+
+
+class ChannelSpliter(ChannelInterface):
+    requirement_string: str = "(Requires any working interface)"
+    targets = list()
+    available_interfaces: dict[str, ChannelInterface] = {
+        "s3": AwsChannel,
+        "az_storage": AzureStorageChannel,
+        "postgres": PostgresChannel,
+    }
+
+    def __init__(self, dct: Dotdict = None):
+        super().__init__(dct)
+        self.targets = list()
+        if "outputs" not in self.configuration:
+            raise AttributeError(T("coal.store.output.split.no_targets"))
+        for output in self.configuration.outputs:
+            channel = self.available_interfaces[output.type]
+            _i = channel(output.conf)
+            if _i.is_available():
+                self.targets.append(_i)
+            else:
+                LOGGER.warning(
+                    T("coal.store.output.split.requirements").format(
+                        interface_name=channel.__class__.__name__, requirements=channel.requirement_string
+                    )
+                )
+        if not self.targets:
+            raise AttributeError(T("coal.store.output.split.no_targets"))
+
+    def send(self, filter: Optional[list[str]] = None) -> bool:
+        any_ok = False
+        for i in self.targets:
+            try:
+                any_ok = i.send(filter=filter) or any_ok
+            except Exception:
+                LOGGER.error(T("coal.store.output.split.send.error").format(interface_name=i.__class__.__name__))
+                if len(self.targets) < 2:
+                    raise
+        return any_ok
+
+    def delete(self, filter: Optional[list[str]] = None) -> bool:
+        any_ok = False
+        for i in self.targets:
+            try:
+                any_ok = i.delete() or any_ok
+            except Exception:
+                LOGGER.error(T("coal.store.output.split.delete.error").format(interface_name=i.__class__.__name__))
+                if len(self.targets) < 2:
+                    raise
+        return any_ok

cosmotech/coal/store/output/postgres_channel.py

@@ -0,0 +1,37 @@
+from typing import Optional
+
+from cosmotech.coal.postgresql.runner import (
+    remove_runner_metadata_from_postgresql,
+    send_runner_metadata_to_postgresql,
+)
+from cosmotech.coal.postgresql.store import dump_store_to_postgresql_from_conf
+from cosmotech.coal.store.output.channel_interface import ChannelInterface
+
+
+class PostgresChannel(ChannelInterface):
+    required_keys = {
+        "coal": ["store"],
+        "cosmotech": ["organization_id", "workspace_id", "runner_id"],
+        "postgres": [
+            "host",
+            "port",
+            "db_name",
+            "db_schema",
+            "user_name",
+            "user_password",
+        ],
+    }
+    requirement_string = required_keys
+
+    def send(self, filter: Optional[list[str]] = None) -> bool:
+        run_id = send_runner_metadata_to_postgresql(self.configuration)
+        dump_store_to_postgresql_from_conf(
+            configuration=self.configuration,
+            selected_tables=filter,
+            fk_id=run_id,
+            replace=False,
+        )
+
+    def delete(self):
+        # removing metadata will trigger cascade delete on real data
+        remove_runner_metadata_from_postgresql(self.configuration)