Utility Module Updates: Update API and PostgreSQL utilities with tests

Author: neomatamune

cosmotech/coal/utils/__init__.py

@@ -1,4 +1,11 @@
-from CosmoTech_Acceleration_Library import __version__
+# Copyright (C) - 2023 - 2025 - Cosmo Tech
+# This document and all information contained herein is the exclusive property -
+# including all intellectual property rights pertaining thereto - of Cosmo Tech.
+# Any use, reproduction, translation, broadcasting, transmission, distribution,
+# etc., to any person is prohibited unless it has been previously and
+# specifically authorized by written means by Cosmo Tech.
+
+from cosmotech.coal import __version__
 
 WEB_DOCUMENTATION_ROOT = f"https://cosmo-tech.github.io/CosmoTech-Acceleration-Library/{__version__}/"
 

cosmotech/coal/utils/api.py

@@ -19,6 +19,7 @@
 
 from cosmotech.coal.cosmotech_api.connection import get_api_client
 from cosmotech.coal.utils.logger import LOGGER
+from cosmotech.orchestrator.utils.translate import T
 
 
 def read_solution_file(solution_file) -> Optional[Solution]:
@@ -28,29 +29,32 @@ def read_solution_file(solution_file) -> Optional[Solution]:
     elif solution_path.suffix == ".json":
         open_function = json.load
     else:
-        LOGGER.error(f"{solution_file} is not a `.yaml` or `.json` file")
+        LOGGER.error(T("coal.errors.solution.invalid_file").format(file=solution_file))
         return None
     with solution_path.open() as _sf:
         solution_content = open_function(_sf)
-    LOGGER.info(f"Loaded {solution_path.absolute()}")
-    _solution = Solution(_configuration=cosmotech_api.Configuration(),
-                         _spec_property_naming=True,
-                         **solution_content)
+    LOGGER.info(T("coal.solution.loaded").format(path=solution_path.absolute()))
+    _solution = Solution(
+        _configuration=cosmotech_api.Configuration(),
+        _spec_property_naming=True,
+        **solution_content,
+    )
     LOGGER.debug(json.dumps(_solution.to_dict(), indent=2, default=str))
     return _solution
 
 
 def get_solution(organization_id, workspace_id) -> Optional[Solution]:
-    LOGGER.info("Configuration to the api set")
+    LOGGER.info(T("coal.solution.api_configured"))
     with get_api_client()[0] as api_client:
         api_w = WorkspaceApi(api_client)
 
-        LOGGER.info("Loading Workspace information to get Solution ID")
+        LOGGER.info(T("coal.solution.loading_workspace"))
         try:
             r_data: Workspace = api_w.find_workspace_by_id(organization_id=organization_id, workspace_id=workspace_id)
         except ServiceException as e:
-            LOGGER.error(f"Workspace {workspace_id} was not found "
-                         f"in Organization {organization_id}")
+            LOGGER.error(
+                T("coal.errors.workspace.not_found").format(workspace_id=workspace_id, organization_id=organization_id)
+            )
             LOGGER.debug(e.body)
             return None
         solution_id = r_data.solution.solution_id

cosmotech/coal/utils/postgresql.py

@@ -8,8 +8,10 @@
 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.logger import LOGGER
@@ -20,16 +22,22 @@ def generate_postgresql_full_uri(
     postgres_port: str,
     postgres_db: str,
     postgres_user: str,
-    postgres_password: str, ) -> str:
+    postgres_password: str,
+    force_encode: bool = False,
+) -> str:
     # Check if password needs percent encoding (contains special characters)
     # We don't log anything about the password for security
-    encoded_password = quote(postgres_password, safe='')
-    return ('postgresql://' +
-            f'{postgres_user}'
-            f':{encoded_password}'
-            f'@{postgres_host}'
-            f':{postgres_port}'
-            f'/{postgres_db}')
+    encoded_password = postgres_password
+    if force_encode:
+        encoded_password = quote(postgres_password, safe="")
+
+    return (
+        "postgresql://" + f"{postgres_user}"
+        f":{encoded_password}"
+        f"@{postgres_host}"
+        f":{postgres_port}"
+        f"/{postgres_db}"
+    )
 
 
 def get_postgresql_table_schema(
@@ -40,10 +48,11 @@ def get_postgresql_table_schema(
     postgres_schema: str,
     postgres_user: str,
     postgres_password: str,
+    force_encode: bool = False,
 ) -> Optional[pa.Schema]:
     """
     Get the schema of an existing PostgreSQL table using SQL queries.
-    
+
     Args:
         target_table_name: Name of the table
         postgres_host: PostgreSQL host
@@ -52,46 +61,47 @@ def get_postgresql_table_schema(
         postgres_schema: PostgreSQL schema name
         postgres_user: PostgreSQL username
         postgres_password: PostgreSQL password
-        
+
     Returns:
         PyArrow Schema if table exists, None otherwise
     """
-    LOGGER.debug(f"Getting schema for table {postgres_schema}.{target_table_name}")
+    LOGGER.debug(
+        T("coal.logs.postgresql.getting_schema").format(
+            postgres_schema=postgres_schema, target_table_name=target_table_name
+        )
+    )
 
-    postgresql_full_uri = generate_postgresql_full_uri(postgres_host,
-                                                       postgres_port,
-                                                       postgres_db,
-                                                       postgres_user,
-                                                       postgres_password)
+    postgresql_full_uri = generate_postgresql_full_uri(
+        postgres_host,
+        postgres_port,
+        postgres_db,
+        postgres_user,
+        postgres_password,
+        force_encode,
+    )
 
-    with (dbapi.connect(postgresql_full_uri) as conn):
+    with dbapi.connect(postgresql_full_uri) as conn:
         try:
-            catalog = conn.adbc_get_objects(depth="tables",
-                                            catalog_filter=postgres_db,
-                                            db_schema_filter=postgres_schema,
-                                            table_name_filter=target_table_name).read_all().to_pylist()[0]
-            schema = catalog["catalog_db_schemas"][0]
-            table = schema["db_schema_tables"][0]
-            if table["table_name"] == target_table_name:
-                return conn.adbc_get_table_schema(
-                    target_table_name,
-                    db_schema_filter=postgres_schema,
+            return conn.adbc_get_table_schema(
+                target_table_name,
+                db_schema_filter=postgres_schema,
+            )
+        except adbc_driver_manager.ProgrammingError:
+            LOGGER.warning(
+                T("coal.logs.postgresql.table_not_found").format(
+                    postgres_schema=postgres_schema, target_table_name=target_table_name
                 )
-        except IndexError:
-            LOGGER.warning(f"Table {postgres_schema}.{target_table_name} not found")
+            )
         return None
 
 
-def adapt_table_to_schema(
-    data: pa.Table,
-    target_schema: pa.Schema
-) -> pa.Table:
+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(f"Starting schema adaptation for table with {len(data)} rows")
-    LOGGER.debug(f"Original schema: {data.schema}")
-    LOGGER.debug(f"Target schema: {target_schema}")
+    LOGGER.debug(T("coal.logs.postgresql.schema_adaptation_start").format(rows=len(data)))
+    LOGGER.debug(T("coal.logs.postgresql.original_schema").format(schema=data.schema))
+    LOGGER.debug(T("coal.logs.postgresql.target_schema").format(schema=target_schema))
 
     target_fields = {field.name: field.type for field in target_schema}
     new_columns = []
@@ -111,63 +121,55 @@ def adapt_table_to_schema(
 
             if original_type != target_type:
                 LOGGER.debug(
-                    f"Attempting to cast column '{field_name}' "
-                    f"from {original_type} to {target_type}"
+                    T("coal.logs.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}"
-                    )
+                    type_conversions.append(f"{field_name}: {original_type} -> {target_type}")
                 except pa.ArrowInvalid as e:
                     LOGGER.warning(
-                        f"Failed to cast column '{field_name}' "
-                        f"from {original_type} to {target_type}. "
-                        f"Filling with nulls. Error: {str(e)}"
+                        T("coal.logs.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}"
-                    )
+                    failed_conversions.append(f"{field_name}: {original_type} -> {target_type}")
             else:
                 new_columns.append(col)
         else:
             # Column doesn't exist - add nulls
-            LOGGER.debug(f"Adding missing column '{field_name}' with null values")
+            LOGGER.debug(T("coal.logs.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
-    ]
+    dropped_columns = [name for name in data.column_names if name not in target_fields]
     if dropped_columns:
-        LOGGER.debug(
-            f"Dropping extra columns not in target schema: {dropped_columns}"
-        )
+        LOGGER.debug(T("coal.logs.postgresql.dropping_columns").format(columns=dropped_columns))
 
     # Create new table
-    adapted_table = pa.Table.from_arrays(
-        new_columns,
-        schema=target_schema
-    )
+    adapted_table = pa.Table.from_arrays(new_columns, schema=target_schema)
 
     # Log summary of adaptations
-    LOGGER.debug("Schema adaptation summary:")
+    LOGGER.debug(T("coal.logs.postgresql.adaptation_summary"))
     if added_columns:
-        LOGGER.debug(f"- Added columns (filled with nulls): {added_columns}")
+        LOGGER.debug(T("coal.logs.postgresql.added_columns").format(columns=added_columns))
     if dropped_columns:
-        LOGGER.debug(f"- Dropped columns: {dropped_columns}")
+        LOGGER.debug(T("coal.logs.postgresql.dropped_columns").format(columns=dropped_columns))
     if type_conversions:
-        LOGGER.debug(f"- Successful type conversions: {type_conversions}")
+        LOGGER.debug(T("coal.logs.postgresql.successful_conversions").format(conversions=type_conversions))
     if failed_conversions:
-        LOGGER.debug(
-            f"- Failed conversions (filled with nulls): {failed_conversions}"
-        )
+        LOGGER.debug(T("coal.logs.postgresql.failed_conversions").format(conversions=failed_conversions))
 
-    LOGGER.debug(f"Final adapted table schema: {adapted_table.schema}")
+    LOGGER.debug(T("coal.logs.postgresql.final_schema").format(schema=adapted_table.schema))
     return adapted_table
 
 
@@ -180,12 +182,15 @@ def send_pyarrow_table_to_postgresql(
     postgres_schema: str,
     postgres_user: str,
     postgres_password: str,
-    replace: bool
+    replace: bool,
+    force_encode: bool = False,
 ) -> int:
     LOGGER.debug(
-        f"Preparing to send data to PostgreSQL table '{postgres_schema}.{target_table_name}'"
+        T("coal.logs.postgresql.preparing_send").format(
+            postgres_schema=postgres_schema, target_table_name=target_table_name
+        )
     )
-    LOGGER.debug(f"Input table has {len(data)} rows")
+    LOGGER.debug(T("coal.logs.postgresql.input_rows").format(rows=len(data)))
 
     # Get existing schema if table exists
     existing_schema = get_postgresql_table_schema(
@@ -195,18 +200,19 @@ def send_pyarrow_table_to_postgresql(
         postgres_db,
         postgres_schema,
         postgres_user,
-        postgres_password
+        postgres_password,
+        force_encode,
     )
 
     if existing_schema is not None:
-        LOGGER.debug(f"Found existing table with schema: {existing_schema}")
+        LOGGER.debug(T("coal.logs.postgresql.found_existing_table").format(schema=existing_schema))
         if not replace:
-            LOGGER.debug("Adapting incoming data to match existing schema")
+            LOGGER.debug(T("coal.logs.postgresql.adapting_data"))
             data = adapt_table_to_schema(data, existing_schema)
         else:
-            LOGGER.debug("Replace mode enabled - skipping schema adaptation")
+            LOGGER.debug(T("coal.logs.postgresql.replace_mode"))
     else:
-        LOGGER.debug("No existing table found - will create new table")
+        LOGGER.debug(T("coal.logs.postgresql.no_existing_table"))
 
     # Proceed with ingestion
     total = 0
@@ -215,21 +221,16 @@ def send_pyarrow_table_to_postgresql(
         postgres_port,
         postgres_db,
         postgres_user,
-        postgres_password
+        postgres_password,
+        force_encode,
     )
 
-    LOGGER.debug("Connecting to PostgreSQL database")
+    LOGGER.debug(T("coal.logs.postgresql.connecting"))
     with dbapi.connect(postgresql_full_uri, autocommit=True) as conn:
         with conn.cursor() as curs:
-            LOGGER.debug(
-                f"Ingesting data with mode: {'replace' if replace else 'create_append'}"
-            )
-            total += curs.adbc_ingest(
-                target_table_name,
-                data,
-                "replace" if replace else "create_append",
-                db_schema_name=postgres_schema
-            )
+            mode = "replace" if replace else "create_append"
+            LOGGER.debug(T("coal.logs.postgresql.ingesting_data").format(mode=mode))
+            total += curs.adbc_ingest(target_table_name, data, mode, db_schema_name=postgres_schema)
 
-    LOGGER.debug(f"Successfully ingested {total} rows")
+    LOGGER.debug(T("coal.logs.postgresql.ingestion_success").format(rows=total))
     return total