PostgreSQL Integration: Add PostgreSQL runner and store functionality with tests

neomatamune · neomatamune · commit 89f2205e6816 · 2025-03-24T15:40:52.000+01:00
diff --git a/cosmotech/coal/postgresql/__init__.py b/cosmotech/coal/postgresql/__init__.py
@@ -0,0 +1,22 @@
+# 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.
+
+"""
+PostgreSQL integration module.
+
+This module provides functions for interacting with PostgreSQL databases.
+"""
+
+# Re-export functions from the runner module
+from cosmotech.coal.postgresql.runner import (
+    send_runner_metadata_to_postgresql,
+)
+
+# Re-export functions from the store module
+from cosmotech.coal.postgresql.store import (
+    dump_store_to_postgresql,
+)
diff --git a/cosmotech/coal/postgresql/runner.py b/cosmotech/coal/postgresql/runner.py
@@ -0,0 +1,93 @@
+# 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.
+
+"""
+PostgreSQL runner operations module.
+
+This module provides functions for interacting with PostgreSQL databases
+for runner metadata operations.
+"""
+
+from adbc_driver_postgresql import dbapi
+
+from cosmotech.coal.cosmotech_api.connection import get_api_client
+from cosmotech.coal.cosmotech_api.runner.metadata import get_runner_metadata
+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_",
+) -> None:
+    """
+    Send runner metadata to a PostgreSQL database.
+
+    Args:
+        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
+    """
+    # Get runner metadata
+    with get_api_client()[0] as api_client:
+        runner = get_runner_metadata(api_client, organization_id, workspace_id, runner_id)
+
+    # Generate PostgreSQL URI
+    postgresql_full_uri = generate_postgresql_full_uri(
+        postgres_host, postgres_port, postgres_db, postgres_user, postgres_password
+    )
+
+    # Connect to PostgreSQL and update runner metadata
+    with dbapi.connect(postgresql_full_uri, autocommit=True) as conn:
+        with conn.cursor() as curs:
+            schema_table = f"{postgres_schema}.{table_prefix}RunnerMetadata"
+            sql_create_table = f"""
+                CREATE TABLE IF NOT EXISTS {schema_table}  (
+                  id varchar(32) PRIMARY KEY,
+                  name varchar(256),
+                  last_run_id varchar(32),
+                  run_template_id varchar(32)
+                );
+            """
+            sql_upsert = f"""
+                INSERT INTO {schema_table} (id, name, last_run_id, run_template_id)
+                  VALUES(%s, %s, %s, %s)
+                  ON CONFLICT (id)
+                  DO
+                    UPDATE SET name = EXCLUDED.name, last_run_id = EXCLUDED.last_run_id;
+            """
+            LOGGER.info(f"creating table {schema_table}")
+            curs.execute(sql_create_table)
+            conn.commit()
+            LOGGER.info(f"adding/updating runner metadata")
+            curs.execute(
+                sql_upsert,
+                (
+                    runner.get("id"),
+                    runner.get("name"),
+                    runner.get("lastRunId"),
+                    runner.get("runTemplateId"),
+                ),
+            )
+            conn.commit()
+            LOGGER.info("Runner metadata table has been updated")
diff --git a/cosmotech/coal/postgresql/store.py b/cosmotech/coal/postgresql/store.py
@@ -0,0 +1,98 @@
+# 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.
+
+"""
+PostgreSQL store operations module.
+
+This module provides functions for interacting with PostgreSQL databases
+for store operations.
+"""
+
+from time import perf_counter
+import pyarrow
+
+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
+from cosmotech.orchestrator.utils.translate import T
+
+
+def dump_store_to_postgresql(
+    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 a PostgreSQL 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.logs.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.logs.database.table_entry").format(table=target_table_name))
+            data = _s.get_table(table_name)
+            if not len(data):
+                LOGGER.info(T("coal.logs.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.logs.database.row_count").format(count=rows))
+            LOGGER.debug(
+                T("coal.logs.progress.operation_timing").format(
+                    operation="Load from datastore", time=f"{_dl_time - _s_time:0.3}"
+                )
+            )
+            LOGGER.debug(
+                T("coal.logs.progress.operation_timing").format(
+                    operation="Send to postgresql", time=f"{_up_time - _dl_time:0.3}"
+                )
+            )
+        _process_end = perf_counter()
+        LOGGER.info(
+            T("coal.logs.database.rows_fetched").format(
+                table="all tables",
+                count=total_rows,
+                time=f"{_process_end - _process_start:0.3}",
+            )
+        )
+    else:
+        LOGGER.info(T("coal.logs.database.store_empty"))
diff --git a/tests/unit/coal/test_postgresql/test_postgresql_runner.py b/tests/unit/coal/test_postgresql/test_postgresql_runner.py
@@ -0,0 +1,111 @@
+# 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 pytest
+from unittest.mock import MagicMock, patch
+
+from cosmotech.coal.postgresql.runner import send_runner_metadata_to_postgresql
+
+
+class TestRunnerFunctions:
+    """Tests for top-level functions in the runner module."""
+
+    @patch("cosmotech.coal.postgresql.runner.get_api_client")
+    @patch("cosmotech.coal.postgresql.runner.get_runner_metadata")
+    @patch("cosmotech.coal.postgresql.runner.generate_postgresql_full_uri")
+    @patch("cosmotech.coal.postgresql.runner.dbapi.connect")
+    def test_send_runner_metadata_to_postgresql(
+        self, mock_connect, mock_generate_uri, mock_get_runner_metadata, mock_get_api_client
+    ):
+        """Test the send_runner_metadata_to_postgresql function."""
+        # Arrange
+        # Mock API client with context manager behavior
+        mock_api_client = MagicMock()
+        mock_api_client_context = MagicMock()
+        mock_api_client.__enter__.return_value = mock_api_client_context
+        mock_get_api_client.return_value = (mock_api_client, "Test Connection")
+
+        # Mock runner metadata
+        mock_runner = {
+            "id": "test-runner-id",
+            "name": "Test Runner",
+            "lastRunId": "test-run-id",
+            "runTemplateId": "test-template-id",
+        }
+        mock_get_runner_metadata.return_value = mock_runner
+
+        # Mock PostgreSQL URI
+        mock_uri = "postgresql://user:password@host:5432/db"
+        mock_generate_uri.return_value = mock_uri
+
+        # Mock PostgreSQL connection and cursor
+        mock_conn = MagicMock()
+        mock_cursor = MagicMock()
+        mock_conn.cursor.return_value.__enter__.return_value = mock_cursor
+        mock_connect.return_value.__enter__.return_value = mock_conn
+
+        # Test parameters
+        organization_id = "test-org"
+        workspace_id = "test-workspace"
+        runner_id = "test-runner-id"
+        postgres_host = "localhost"
+        postgres_port = 5432
+        postgres_db = "testdb"
+        postgres_schema = "public"
+        postgres_user = "user"
+        postgres_password = "password"
+        table_prefix = "Test_"
+
+        # Act
+        send_runner_metadata_to_postgresql(
+            organization_id,
+            workspace_id,
+            runner_id,
+            postgres_host,
+            postgres_port,
+            postgres_db,
+            postgres_schema,
+            postgres_user,
+            postgres_password,
+            table_prefix,
+        )
+
+        # Assert
+        # Check that API client was used correctly
+        mock_get_api_client.assert_called_once()
+        mock_get_runner_metadata.assert_called_once_with(
+            mock_api_client_context, organization_id, workspace_id, runner_id
+        )
+
+        # Check that PostgreSQL URI was generated correctly
+        mock_generate_uri.assert_called_once_with(
+            postgres_host, postgres_port, postgres_db, postgres_user, postgres_password
+        )
+
+        # Check that PostgreSQL connection was established
+        mock_connect.assert_called_once_with(mock_uri, autocommit=True)
+
+        # Check that SQL statements were executed
+        assert mock_cursor.execute.call_count == 2
+
+        # Verify the SQL statements (partially, since the exact SQL is complex)
+        create_table_call = mock_cursor.execute.call_args_list[0]
+        assert "CREATE TABLE IF NOT EXISTS" in create_table_call[0][0]
+        assert f"{postgres_schema}.{table_prefix}RunnerMetadata" in create_table_call[0][0]
+
+        upsert_call = mock_cursor.execute.call_args_list[1]
+        assert "INSERT INTO" in upsert_call[0][0]
+        assert f"{postgres_schema}.{table_prefix}RunnerMetadata" in upsert_call[0][0]
+        assert upsert_call[0][1] == (
+            mock_runner["id"],
+            mock_runner["name"],
+            mock_runner["lastRunId"],
+            mock_runner["runTemplateId"],
+        )
+
+        # Check that commits were called
+        assert mock_conn.commit.call_count == 2
diff --git a/tests/unit/coal/test_postgresql/test_postgresql_store.py b/tests/unit/coal/test_postgresql/test_postgresql_store.py
