cosmotech-acceleration-library 1.0.1__py3-none-any.whl → 2.0.0__py3-none-any.whl

cosmotech/coal/cosmotech_api/objects/parameters.py

@@ -0,0 +1,127 @@
+# 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 json
+import os
+import pathlib
+from csv import DictWriter
+from typing import Any, Dict, List
+
+from cosmotech.orchestrator.utils.translate import T
+
+from cosmotech.coal.utils.logger import LOGGER
+
+
+class Parameters:
+    values: Dict[str, Any] = dict()
+    parameters_list: List[Dict[str, Any]] = list()
+
+    def __init__(self, runner_data: Any):
+        """
+        Extract parameters from runner data.
+
+        Args:
+            runner_data: Runner data object
+
+        Returns:
+            Dictionary mapping parameter IDs to values
+        """
+        for parameter in runner_data.parameters_values:
+            self.values[parameter.parameter_id] = parameter.value
+        self.parameters_list = self.format_parameters_list(runner_data)
+
+    @staticmethod
+    def format_parameters_list(runner_data: Any) -> List[Dict[str, Any]]:
+        """
+        Format parameters from runner data as a list of dictionaries.
+
+        Args:
+            runner_data: Runner data object
+
+        Returns:
+            List of parameter dictionaries
+        """
+        parameters = []
+
+        if not runner_data.parameters_values:
+            return parameters
+
+        max_name_size = max(map(lambda r: len(r.parameter_id), runner_data.parameters_values))
+        max_type_size = max(map(lambda r: len(r.var_type), runner_data.parameters_values))
+
+        for parameter_data in runner_data.parameters_values:
+            parameter_name = parameter_data.parameter_id
+            value = parameter_data.value
+            var_type = parameter_data.var_type
+            is_inherited = parameter_data.is_inherited
+
+            parameters.append(
+                {
+                    "parameterId": parameter_name,
+                    "value": value,
+                    "varType": var_type,
+                    "isInherited": is_inherited,
+                }
+            )
+
+            LOGGER.debug(
+                T("coal.cosmotech_api.runner.parameter_debug").format(
+                    param_id=parameter_name,
+                    max_name_size=max_name_size,
+                    var_type=var_type,
+                    max_type_size=max_type_size,
+                    value=value,
+                    inherited=" inherited" if is_inherited else "",
+                )
+            )
+
+        return parameters
+
+    def write_parameters_to_json(
+        self,
+        parameter_folder: str,
+    ) -> str:
+        pathlib.Path(parameter_folder).mkdir(exist_ok=True, parents=True)
+        tmp_parameter_file = os.path.join(parameter_folder, "parameters.json")
+
+        LOGGER.info(T("coal.cosmotech_api.runner.generating_file").format(file=tmp_parameter_file))
+
+        with open(tmp_parameter_file, "w") as _file:
+            json.dump(self.parameters_list, _file, indent=2)
+
+        return tmp_parameter_file
+
+    def write_parameters_to_csv(
+        self,
+        parameter_folder: str,
+    ) -> str:
+        pathlib.Path(parameter_folder).mkdir(exist_ok=True, parents=True)
+        tmp_parameter_file = os.path.join(parameter_folder, "parameters.csv")
+
+        LOGGER.info(T("coal.cosmotech_api.runner.generating_file").format(file=tmp_parameter_file))
+
+        with open(tmp_parameter_file, "w") as _file:
+            _w = DictWriter(_file, fieldnames=["parameterId", "value", "varType", "isInherited"])
+            _w.writeheader()
+            _w.writerows(self.parameters_list)
+
+        return tmp_parameter_file
+
+    def write_parameters(
+        self,
+        parameter_folder: str,
+        write_csv: bool = True,
+        write_json: bool = False,
+    ) -> Dict[str, str]:
+        result = {}
+
+        if write_csv:
+            result["csv"] = self.write_parameters_to_csv(parameter_folder)
+
+        if write_json:
+            result["json"] = self.write_parameters_to_json(parameter_folder)
+
+        return result

cosmotech/coal/postgresql/runner.py

@@ -13,56 +13,40 @@ for runner metadata operations.
 """
 
 from adbc_driver_postgresql import dbapi
+from cosmotech.orchestrator.utils.translate import T
 
-from cosmotech.coal.cosmotech_api.connection import get_api_client
-from cosmotech.coal.cosmotech_api.runner.metadata import get_runner_metadata
+from cosmotech.coal.cosmotech_api.apis.runner import RunnerApi
+from cosmotech.coal.postgresql.utils import PostgresUtils
+from cosmotech.coal.utils.configuration import Configuration
 from cosmotech.coal.utils.logger import LOGGER
-from cosmotech.coal.utils.postgresql import generate_postgresql_full_uri
-from cosmotech.orchestrator.utils.translate import T
 
 
 def send_runner_metadata_to_postgresql(
-    organization_id: str,
-    workspace_id: str,
-    runner_id: str,
-    postgres_host: str,
-    postgres_port: int,
-    postgres_db: str,
-    postgres_schema: str,
-    postgres_user: str,
-    postgres_password: str,
-    table_prefix: str = "Cosmotech_",
-    force_encode: bool = False,
-) -> None:
+    configuration: Configuration,
+) -> str:
     """
     Send runner metadata to a PostgreSQL database.
 
     Args:
+        configuration: coal configuration
         organization_id: Organization ID
         workspace_id: Workspace ID
         runner_id: Runner ID
-        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
-        force_encode: force password encoding to percent encoding
     """
-    # Get runner metadata
-    with get_api_client()[0] as api_client:
-        runner = get_runner_metadata(api_client, organization_id, workspace_id, runner_id)
+    _psql = PostgresUtils(configuration)
 
-    # Generate PostgreSQL URI
-    postgresql_full_uri = generate_postgresql_full_uri(
-        postgres_host, str(postgres_port), postgres_db, postgres_user, postgres_password, force_encode=force_encode
+    # Get runner metadata
+    _runner_api = RunnerApi(configuration)
+    runner = _runner_api.get_runner_metadata(
+        configuration.cosmotech.organization_id,
+        configuration.cosmotech.workspace_id,
+        configuration.cosmotech.runner_id,
     )
 
     # Connect to PostgreSQL and update runner metadata
-    with dbapi.connect(postgresql_full_uri, autocommit=True) as conn:
+    with dbapi.connect(_psql.full_uri, autocommit=True) as conn:
         with conn.cursor() as curs:
-            schema_table = f"{postgres_schema}.{table_prefix}RunnerMetadata"
+            schema_table = f"{str(_psql.db_schema)}.{str(_psql.table_prefix)}RunnerMetadata"
             sql_create_table = f"""
                 CREATE TABLE IF NOT EXISTS {schema_table}  (
                   id varchar(32) PRIMARY KEY,
@@ -71,6 +55,10 @@ def send_runner_metadata_to_postgresql(
                   run_template_id varchar(32)
                 );
             """
+            LOGGER.info(T("coal.services.postgresql.creating_table").format(schema_table=schema_table))
+            curs.execute(sql_create_table)
+            conn.commit()
+            LOGGER.info(T("coal.services.postgresql.metadata"))
             sql_upsert = f"""
                 INSERT INTO {schema_table} (id, name, last_run_id, run_template_id)
                   VALUES(%s, %s, %s, %s)
@@ -78,10 +66,6 @@ def send_runner_metadata_to_postgresql(
                   DO
                     UPDATE SET name = EXCLUDED.name, last_run_id = EXCLUDED.last_run_id;
             """
-            LOGGER.info(T("coal.services.postgresql.creating_table").format(schema_table=schema_table))
-            curs.execute(sql_create_table)
-            conn.commit()
-            LOGGER.info(T("coal.services.postgresql.metadata"))
             curs.execute(
                 sql_upsert,
                 (
@@ -93,3 +77,39 @@ def send_runner_metadata_to_postgresql(
             )
             conn.commit()
             LOGGER.info(T("coal.services.postgresql.metadata_updated"))
+    return runner.get("lastRunId")
+
+
+def remove_runner_metadata_from_postgresql(
+    configuration: Configuration,
+) -> str:
+    """
+    Removes run_id from metadata table that trigger cascade delete on other tables
+
+    Args:
+        configuration: coal configuration
+        organization_id: Organization ID
+        workspace_id: Workspace ID
+        runner_id: Runner ID
+    """
+    _psql = PostgresUtils(configuration)
+
+    # Get runner metadata
+    _runner_api = RunnerApi(configuration)
+    runner = _runner_api.get_runner_metadata(
+        configuration.cosmotech.organization_id,
+        configuration.cosmotech.workspace_id,
+        configuration.cosmotech.runner_id,
+    )
+
+    # Connect to PostgreSQL and remove runner metadata row
+    with dbapi.connect(_psql.full_uri, autocommit=True) as conn:
+        with conn.cursor() as curs:
+            schema_table = f"{_psql.db_schema}.{_psql.table_prefix}RunnerMetadata"
+            sql_delete_from_metatable = f"""
+                DELETE FROM {schema_table}
+                WHERE last_run_id={runner.get("lastRunId")};
+            """
+            curs.execute(sql_delete_from_metatable)
+            conn.commit()
+    return runner.get("lastRunId")

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,79 @@ 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
+    """
+    _c = Configuration()
+    _c.postgres.host = postgres_host
+    _c.postgres.port = postgres_port
+    _c.postgres.db_name = postgres_db
+    _c.postgres.db_schema = postgres_schema
+    _c.postgres.user_name = postgres_user
+    _c.postgres.user_password = postgres_password
+    _c.postgres.password_encoding = force_encode
+    _c.postgres.table_prefix = table_prefix
+
+    dump_store_to_postgresql_from_conf(
+        configuration=_c, store_folder=store_folder, replace=replace, selected_tables=selected_tables, fk_id=fk_id
+    )
+
+
+def dump_store_to_postgresql_from_conf(
+    configuration: Configuration,
+    store_folder: str,
+    replace: bool = True,
+    selected_tables: list[str] = [],
+    fk_id: str = None,
+) -> None:
+    """
+    Dump Store data to a PostgreSQL database.
+
+    Args:
+        configuration: coal Configuration
+        store_folder: Folder containing the Store
+        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)
+    print(_psql.send_pyarrow_table_to_postgresql)
     _s = Store(store_location=store_folder)
 
     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))
+            if fk_id:
+                _s.execute_query(
+                    f"""
+                    ALTER TABLE {_psql.table_prefix}{table_name}
+                    ADD run_id TEXT NOT NULL
+                    DEFAULT ({fk_id})
+                    """
+                )
             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(
+            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, "run_id", metadata_table, "last_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,254 @@
+# 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}"
+        )
+
+    def metadata_table_name(self):
+        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