cosmotech-acceleration-library 2.0.0__tar.gz → 2.1.0rc1__tar.gz

{cosmotech_acceleration_library-2.0.0 → cosmotech_acceleration_library-2.1.0rc1}/MANIFEST.in

@@ -1,4 +1,4 @@
 include README.md
 include LICENSE
 graft cosmotech/orchestrator_plugins
-graft cosmotech/translation
+graft cosmotech/translation

{cosmotech_acceleration_library-2.0.0/cosmotech_acceleration_library.egg-info → cosmotech_acceleration_library-2.1.0rc1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cosmotech_acceleration_library
-Version: 2.0.0
+Version: 2.1.0rc1
 Summary: Acceleration libraries for CosmoTech cloud based solution development
 Author-email: Cosmo Tech <platform@cosmotech.com>
 Project-URL: Homepage, https://www.cosmotech.com
@@ -8,26 +8,27 @@ Project-URL: Source, https://github.com/Cosmo-Tech/CosmoTech-Acceleration-Librar
 Project-URL: Documentation, https://cosmo-tech.github.io/CosmoTech-Acceleration-Library
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: azure-functions~=1.20.0
-Requires-Dist: azure-digitaltwins-core~=1.2.0
-Requires-Dist: azure-identity~=1.16.1
-Requires-Dist: azure-kusto-data~=4.4.1
-Requires-Dist: azure-kusto-ingest~=4.4.1
-Requires-Dist: tenacity~=8.3.0
-Requires-Dist: python-keycloak~=4.7.3
+Requires-Dist: azure-core~=1.38
+Requires-Dist: azure-functions~=1.20
+Requires-Dist: azure-digitaltwins-core~=1.2
+Requires-Dist: azure-identity~=1.16
+Requires-Dist: azure-kusto-data~=4.4
+Requires-Dist: azure-kusto-ingest~=4.4
+Requires-Dist: tenacity~=9.1
+Requires-Dist: python-keycloak~=5.8
 Requires-Dist: cosmotech-api~=5.0.0rc1
-Requires-Dist: boto3~=1.35.19
-Requires-Dist: requests~=2.32.3
-Requires-Dist: singlestoredb~=1.10.0
+Requires-Dist: boto3~=1.41
+Requires-Dist: requests~=2.32
+Requires-Dist: singlestoredb~=1.16
 Requires-Dist: cosmotech-run-orchestrator~=2.0.0
-Requires-Dist: pyarrow~=20.0.0
-Requires-Dist: adbc-driver-manager~=1.7.0
-Requires-Dist: adbc-driver-sqlite~=1.7.0
-Requires-Dist: adbc-driver-postgresql~=1.7.0
-Requires-Dist: click~=8.1.7
-Requires-Dist: rich-click~=1.7.3
-Requires-Dist: click-log~=0.4.0
-Requires-Dist: tqdm~=4.67.1
+Requires-Dist: pyarrow~=22.0
+Requires-Dist: adbc-driver-manager~=1.7
+Requires-Dist: adbc-driver-sqlite~=1.7
+Requires-Dist: adbc-driver-postgresql~=1.7
+Requires-Dist: click~=8.1
+Requires-Dist: rich-click~=1.7
+Requires-Dist: click-log~=0.4
+Requires-Dist: tqdm~=4.67
 Requires-Dist: openpyxl~=3.1
 Requires-Dist: pandas~=2.3
 Requires-Dist: python-dateutil~=2.8
@@ -35,14 +36,12 @@ Requires-Dist: rich~=13.7
 Requires-Dist: setuptools
 Provides-Extra: doc
 Requires-Dist: mkdocs~=1.6.1; extra == "doc"
-Requires-Dist: mkdocs-click~=0.8.1; extra == "doc"
-Requires-Dist: mkdocs-gen-files~=0.5.0; extra == "doc"
-Requires-Dist: mkdocstrings[python]~=0.28.2; extra == "doc"
+Requires-Dist: mkdocs-click~=0.8; extra == "doc"
+Requires-Dist: mkdocs-gen-files~=0.5; extra == "doc"
 Requires-Dist: mkdocs-awesome-nav~=3.1.2; extra == "doc"
 Requires-Dist: pymdown-extensions~=10.7; extra == "doc"
-Requires-Dist: requirements-parser~=0.11.0; extra == "doc"
-Requires-Dist: mike~=2.0.0; extra == "doc"
-Requires-Dist: griffe~=1.5.7; extra == "doc"
+Requires-Dist: requirements-parser~=0.11; extra == "doc"
+Requires-Dist: mike~=2.1.0; extra == "doc"
 Requires-Dist: mkdocs-include-dir-to-nav~=1.2.0; extra == "doc"
 Requires-Dist: mkdocs-material[imaging]~=9.6.14; extra == "doc"
 Requires-Dist: mkdocs-table-reader-plugin~=2.0.3; extra == "doc"
@@ -52,9 +51,9 @@ Requires-Dist: pytest; extra == "test"
 Requires-Dist: pytest-docker; extra == "test"
 Requires-Dist: pytest-cov; extra == "test"
 Provides-Extra: dev
-Requires-Dist: black~=24.3.0; extra == "dev"
+Requires-Dist: black~=25.0; extra == "dev"
 Requires-Dist: pre-commit~=3.3.2; extra == "dev"
-Requires-Dist: isort~=7.0.0; extra == "dev"
+Requires-Dist: isort~=7.0; extra == "dev"
 Provides-Extra: all
 Requires-Dist: CosmoTech_Acceleration_Library[dev,doc,test]; extra == "all"
 Dynamic: license-file

{cosmotech_acceleration_library-2.0.0 → cosmotech_acceleration_library-2.1.0rc1}/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__ = "2.0.0"
+__version__ = "2.1.0-rc1"

{cosmotech_acceleration_library-2.0.0 → cosmotech_acceleration_library-2.1.0rc1}/cosmotech/coal/azure/adx/__init__.py

@@ -5,22 +5,36 @@
 # etc., to any person is prohibited unless it has been previously and
 # specifically authorized by written means by Cosmo Tech.
 
-from cosmotech.coal.azure.adx.auth import create_kusto_client, create_ingest_client, initialize_clients
-from cosmotech.coal.azure.adx.query import run_query, run_command_query
+from cosmotech.coal.azure.adx.auth import (
+    create_ingest_client,
+    create_kusto_client,
+    initialize_clients,
+)
 from cosmotech.coal.azure.adx.ingestion import (
-    ingest_dataframe,
-    send_to_adx,
+    IngestionStatus,
     check_ingestion_status,
-    monitor_ingestion,
     handle_failures,
-    IngestionStatus,
+    ingest_dataframe,
+    monitor_ingestion,
+    send_to_adx,
 )
-from cosmotech.coal.azure.adx.tables import table_exists, create_table, check_and_create_table, _drop_by_tag
-from cosmotech.coal.azure.adx.utils import type_mapping, create_column_mapping
-from cosmotech.coal.azure.adx.store import send_pyarrow_table_to_adx, send_table_data, process_tables, send_store_to_adx
+from cosmotech.coal.azure.adx.query import run_command_query, run_query
 from cosmotech.coal.azure.adx.runner import (
-    prepare_csv_content,
     construct_create_query,
     insert_csv_files,
+    prepare_csv_content,
     send_runner_data,
 )
+from cosmotech.coal.azure.adx.store import (
+    process_tables,
+    send_pyarrow_table_to_adx,
+    send_store_to_adx,
+    send_table_data,
+)
+from cosmotech.coal.azure.adx.tables import (
+    _drop_by_tag,
+    check_and_create_table,
+    create_table,
+    table_exists,
+)
+from cosmotech.coal.azure.adx.utils import create_column_mapping, type_mapping

{cosmotech_acceleration_library-2.0.0 → cosmotech_acceleration_library-2.1.0rc1}/cosmotech/coal/azure/adx/ingestion.py

@@ -5,28 +5,24 @@
 # etc., to any person is prohibited unless it has been previously and
 # specifically authorized by written means by Cosmo Tech.
 
+import os
+import time
 from enum import Enum
-from typing import Dict
-from typing import Iterator
-from typing import List
-from typing import Optional
-from typing import Tuple
+from typing import Dict, Iterator, List, Optional, Tuple
 
-import os
 import pandas as pd
-import time
 import tqdm
 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.status import FailureMessage
-from azure.kusto.ingest.status import KustoIngestStatusQueues
-from azure.kusto.ingest.status import SuccessMessage
+from azure.kusto.ingest import IngestionProperties, QueuedIngestClient, ReportLevel
+from azure.kusto.ingest.status import (
+    FailureMessage,
+    KustoIngestStatusQueues,
+    SuccessMessage,
+)
 from cosmotech.orchestrator.utils.translate import T
 
-from cosmotech.coal.azure.adx.tables import create_table, _drop_by_tag
+from cosmotech.coal.azure.adx.tables import _drop_by_tag, create_table
 from cosmotech.coal.azure.adx.utils import type_mapping
 from cosmotech.coal.utils.logger import LOGGER
 

{cosmotech_acceleration_library-2.0.0 → cosmotech_acceleration_library-2.1.0rc1}/cosmotech/coal/azure/blob.py

@@ -21,6 +21,7 @@ from azure.storage.blob import BlobServiceClient
 from cosmotech.orchestrator.utils.translate import T
 
 from cosmotech.coal.store.store import Store
+from cosmotech.coal.utils.configuration import Configuration
 from cosmotech.coal.utils.logger import LOGGER
 
 VALID_TYPES = (
@@ -31,42 +32,35 @@ VALID_TYPES = (
 
 
 def dump_store_to_azure(
-    store_folder: str,
-    account_name: str,
-    container_name: str,
-    tenant_id: str,
-    client_id: str,
-    client_secret: str,
-    output_type: str = "sqlite",
-    file_prefix: str = "",
+    configuration: Configuration = Configuration(),
     selected_tables: list[str] = [],
 ) -> None:
     """
     Dump Store data to Azure Blob Storage.
 
     Args:
-        store_folder: Folder containing the Store
-        account_name: Azure Storage account name
-        container_name: Azure Storage container name
-        tenant_id: Azure tenant ID
-        client_id: Azure client ID
-        client_secret: Azure client secret
-        output_type: Output file type (sqlite, csv, or parquet)
-        file_prefix: Prefix for uploaded files
+        configuration: Configuration utils class
+        selected_tables: List of tables name
 
     Raises:
         ValueError: If the output type is invalid
     """
-    _s = Store(store_location=store_folder)
+    _s = Store(configuration=configuration)
+    output_type = configuration.safe_get("azure.output_type", default="sqlite")
+    file_prefix = configuration.safe_get("azure.file_prefix", default="")
 
     if output_type not in VALID_TYPES:
         LOGGER.error(T("coal.common.validation.invalid_output_type").format(output_type=output_type))
         raise ValueError(T("coal.common.validation.invalid_output_type").format(output_type=output_type))
 
     container_client = BlobServiceClient(
-        account_url=f"https://{account_name}.blob.core.windows.net/",
-        credential=ClientSecretCredential(tenant_id=tenant_id, client_id=client_id, client_secret=client_secret),
-    ).get_container_client(container_name)
+        account_url=f"https://{configuration.azure.account_name}.blob.core.windows.net/",
+        credential=ClientSecretCredential(
+            tenant_id=configuration.azure.tenant_id,
+            client_id=configuration.azure.client_id,
+            client_secret=configuration.azure.client_secret,
+        ),
+    ).get_container_client(configuration.azure.container_name)
 
     def data_upload(data_stream: BytesIO, file_name: str):
         uploaded_file_name = file_prefix + file_name

cosmotech_acceleration_library-2.1.0rc1/cosmotech/coal/cosmotech_api/apis/dataset.py

@@ -0,0 +1,222 @@
+# 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 pathlib import Path
+from typing import Optional, Union
+
+from cosmotech.orchestrator.utils.translate import T
+from cosmotech_api import (
+    Dataset,
+)
+from cosmotech_api import DatasetApi as BaseDatasetApi
+from cosmotech_api import (
+    DatasetCreateRequest,
+    DatasetPartCreateRequest,
+    DatasetPartTypeEnum,
+)
+
+from cosmotech.coal.cosmotech_api.objects.connection import Connection
+from cosmotech.coal.utils.configuration import ENVIRONMENT_CONFIGURATION, Configuration
+from cosmotech.coal.utils.logger import LOGGER
+
+
+class DatasetApi(BaseDatasetApi, Connection):
+
+    def __init__(
+        self,
+        configuration: Configuration = ENVIRONMENT_CONFIGURATION,
+    ):
+        Connection.__init__(self, configuration)
+        BaseDatasetApi.__init__(self, self.api_client)
+
+        LOGGER.debug(T("coal.cosmotech_api.initialization.dataset_api_initialized"))
+
+    def download_dataset(self, dataset_id) -> Dataset:
+        LOGGER.debug(f"Downloading dataset {dataset_id}")
+        dataset = self.get_dataset(
+            organization_id=self.configuration.cosmotech.organization_id,
+            workspace_id=self.configuration.cosmotech.workspace_id,
+            dataset_id=dataset_id,
+        )
+        # send dataset files under dataset id folder
+        destination = Path(self.configuration.cosmotech.dataset_absolute_path) / dataset_id
+        for part in dataset.parts:
+            self._download_part(dataset_id, part, destination)
+        return dataset
+
+    def download_parameter(self, dataset_id) -> Dataset:
+        LOGGER.debug(f"Downloading dataset {dataset_id}")
+        dataset = self.get_dataset(
+            organization_id=self.configuration.cosmotech.organization_id,
+            workspace_id=self.configuration.cosmotech.workspace_id,
+            dataset_id=dataset_id,
+        )
+        # send parameters file under parameters_name folder
+        destination = Path(self.configuration.cosmotech.parameters_absolute_path) / dataset_id
+        for part in dataset.parts:
+            part_dst = destination / part.name
+            self._download_part(dataset_id, part, part_dst)
+        return dataset
+
+    def _download_part(self, dataset_id, dataset_part, destination):
+        part_file_path = destination / dataset_part.source_name
+        part_file_path.parent.mkdir(parents=True, exist_ok=True)
+        data_part = self.download_dataset_part(
+            organization_id=self.configuration.cosmotech.organization_id,
+            workspace_id=self.configuration.cosmotech.workspace_id,
+            dataset_id=dataset_id,
+            dataset_part_id=dataset_part.id,
+        )
+        with open(part_file_path, "wb") as binary_file:
+            binary_file.write(data_part)
+        LOGGER.debug(
+            T("coal.services.dataset.part_downloaded").format(
+                part_name=dataset_part.source_name, file_path=part_file_path
+            )
+        )
+
+    @staticmethod
+    def path_to_parts(_path, part_type) -> list[tuple[str, Path, DatasetPartTypeEnum]]:
+        if (_path := Path(_path)).is_dir():
+            return list((str(_p.relative_to(_path)), _p, part_type) for _p in _path.rglob("*") if _p.is_file())
+        return list(((_path.name, _path, part_type),))
+
+    def upload_dataset(
+        self,
+        dataset_name: str,
+        as_files: Optional[list[Union[Path, str]]] = (),
+        as_db: Optional[list[Union[Path, str]]] = (),
+        tags: Optional[list[str]] = None,
+        additional_data: Optional[dict] = None,
+    ) -> Dataset:
+        """Upload a new dataset with optional tags and additional data.
+
+        Args:
+            dataset_name: The name of the dataset to create
+            as_files: List of file paths to upload as FILE type parts
+            as_db: List of file paths to upload as DB type parts
+            tags: Optional list of tags to associate with the dataset
+            additional_data: Optional dictionary of additional metadata
+
+        Returns:
+            The created Dataset object
+        """
+        _parts = list()
+
+        for _f in as_files:
+            _parts.extend(self.path_to_parts(_f, DatasetPartTypeEnum.FILE))
+
+        for _db in as_db:
+            _parts.extend(self.path_to_parts(_db, DatasetPartTypeEnum.DB))
+
+        d_request = DatasetCreateRequest(
+            name=dataset_name,
+            tags=tags,
+            additional_data=additional_data,
+            parts=list(
+                DatasetPartCreateRequest(
+                    name=_p_name,
+                    description=_p_name,
+                    sourceName=_p_name,
+                    type=_type,
+                )
+                for _p_name, _, _type in _parts
+            ),
+        )
+
+        _files = []
+        for _p in _parts:
+            with _p[1].open("rb") as _p_file:
+                _files.append((_p[0], _p_file.read()))
+
+        d_ret = self.create_dataset(
+            self.configuration.cosmotech.organization_id,
+            self.configuration.cosmotech.workspace_id,
+            d_request,
+            files=_files,
+        )
+
+        LOGGER.info(T("coal.services.dataset.dataset_created").format(dataset_id=d_ret.id))
+        return d_ret
+
+    def upload_dataset_parts(
+        self,
+        dataset_id: str,
+        as_files: Optional[list[Union[Path, str]]] = (),
+        as_db: Optional[list[Union[Path, str]]] = (),
+        replace_existing: bool = False,
+    ) -> Dataset:
+        """Upload parts to an existing dataset.
+
+        Args:
+            dataset_id: The ID of the existing dataset
+            as_files: List of file paths to upload as FILE type parts
+            as_db: List of file paths to upload as DB type parts
+            replace_existing: If True, replace existing parts with same name
+
+        Returns:
+            The updated Dataset object
+        """
+        # Get current dataset to check existing parts
+        current_dataset = self.get_dataset(
+            organization_id=self.configuration.cosmotech.organization_id,
+            workspace_id=self.configuration.cosmotech.workspace_id,
+            dataset_id=dataset_id,
+        )
+
+        # Build set of existing part names and their IDs for quick lookup
+        existing_parts = {part.source_name: part.id for part in (current_dataset.parts or [])}
+
+        # Collect parts to upload
+        _parts = list()
+        for _f in as_files:
+            _parts.extend(self.path_to_parts(_f, DatasetPartTypeEnum.FILE))
+        for _db in as_db:
+            _parts.extend(self.path_to_parts(_db, DatasetPartTypeEnum.DB))
+
+        # Process each part
+        for _p_name, _p_path, _type in _parts:
+            if _p_name in existing_parts:
+                if replace_existing:
+                    # Delete existing part before creating new one
+                    self.delete_dataset_part(
+                        organization_id=self.configuration.cosmotech.organization_id,
+                        workspace_id=self.configuration.cosmotech.workspace_id,
+                        dataset_id=dataset_id,
+                        dataset_part_id=existing_parts[_p_name],
+                    )
+                    LOGGER.info(T("coal.services.dataset.part_replaced").format(part_name=_p_name))
+                else:
+                    LOGGER.warning(T("coal.services.dataset.part_skipped").format(part_name=_p_name))
+                    continue
+
+            # Create new part
+            part_request = DatasetPartCreateRequest(
+                name=_p_name,
+                description=_p_name,
+                sourceName=_p_name,
+                type=_type,
+            )
+
+            with _p_path.open("rb") as _p_file:
+                self.create_dataset_part(
+                    organization_id=self.configuration.cosmotech.organization_id,
+                    workspace_id=self.configuration.cosmotech.workspace_id,
+                    dataset_id=dataset_id,
+                    dataset_part_create_request=part_request,
+                    file=(_p_name, _p_file.read()),
+                )
+                LOGGER.debug(T("coal.services.dataset.part_uploaded").format(part_name=_p_name))
+
+        # Return updated dataset
+        updated_dataset = self.get_dataset(
+            organization_id=self.configuration.cosmotech.organization_id,
+            workspace_id=self.configuration.cosmotech.workspace_id,
+            dataset_id=dataset_id,
+        )
+
+        LOGGER.info(T("coal.services.dataset.parts_uploaded").format(dataset_id=dataset_id))
+        return updated_dataset

{cosmotech_acceleration_library-2.0.0 → cosmotech_acceleration_library-2.1.0rc1}/cosmotech/coal/cosmotech_api/apis/runner.py

@@ -29,43 +29,47 @@ class RunnerApi(BaseRunnerApi, Connection):
 
     def get_runner_metadata(
         self,
-        organization_id: str,
-        workspace_id: str,
-        runner_id: str,
+        runner_id: Optional[str] = None,
         include: Optional[list[str]] = None,
         exclude: Optional[list[str]] = None,
     ) -> dict[str, Any]:
-        runner = self.get_runner(organization_id, workspace_id, runner_id)
+        runner = self.get_runner(
+            self.configuration.cosmotech.organization_id,
+            self.configuration.cosmotech.workspace_id,
+            runner_id or self.configuration.cosmotech.runner_id,
+        )
 
         return runner.model_dump(by_alias=True, exclude_none=True, include=include, exclude=exclude, mode="json")
 
     def download_runner_data(
         self,
-        organization_id: str,
-        workspace_id: str,
-        runner_id: str,
-        parameter_folder: str,
-        dataset_folder: Optional[str] = None,
+        download_datasets: Optional[str] = None,
     ):
         LOGGER.info(T("coal.cosmotech_api.runner.starting_download"))
 
         # Get runner data
-        runner_data = self.get_runner(organization_id, workspace_id, runner_id)
+        runner = self.get_runner(
+            self.configuration.cosmotech.organization_id,
+            self.configuration.cosmotech.workspace_id,
+            self.configuration.cosmotech.runner_id,
+        )
 
         # Skip if no parameters found
-        if not runner_data.parameters_values:
+        if not runner.parameters_values:
             LOGGER.warning(T("coal.cosmotech_api.runner.no_parameters"))
         else:
             LOGGER.info(T("coal.cosmotech_api.runner.loaded_data"))
-            parameters = Parameters(runner_data)
-            parameters.write_parameters_to_json(parameter_folder)
+            parameters = Parameters(runner)
+            parameters.write_parameters_to_json(self.configuration.cosmotech.parameters_absolute_path)
 
-        # Download datasets if requested
-        if dataset_folder:
-            datasets_ids = runner_data.datasets.bases
+        if runner.datasets.parameter:
+            ds_api = DatasetApi(self.configuration)
+            ds_api.download_parameter(runner.datasets.parameter)
 
-            if datasets_ids:
-                LOGGER.info(T("coal.cosmotech_api.runner.downloading_datasets").format(count=len(datasets_ids)))
+        # Download datasets if requested
+        if download_datasets:
+            LOGGER.info(T("coal.cosmotech_api.runner.downloading_datasets").format(count=len(runner.datasets.bases)))
+            if runner.datasets.bases:
                 ds_api = DatasetApi(self.configuration)
-                for dataset_id in datasets_ids:
+                for dataset_id in runner.datasets.bases:
                     ds_api.download_dataset(dataset_id)

{cosmotech_acceleration_library-2.0.0 → cosmotech_acceleration_library-2.1.0rc1}/cosmotech/coal/postgresql/runner.py

@@ -38,8 +38,6 @@ def send_runner_metadata_to_postgresql(
     # 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,
     )
 
@@ -51,7 +49,7 @@ def send_runner_metadata_to_postgresql(
                 CREATE TABLE IF NOT EXISTS {schema_table}  (
                   id varchar(32) PRIMARY KEY,
                   name varchar(256),
-                  last_run_id varchar(32),
+                  last_csm_run_id varchar(32),
                   run_template_id varchar(32)
                 );
             """
@@ -60,24 +58,25 @@ def send_runner_metadata_to_postgresql(
             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)
+                INSERT INTO {schema_table} (id, name, last_csm_run_id, run_template_id)
+                  VALUES ($1, $2, $3, $4)
                   ON CONFLICT (id)
                   DO
-                    UPDATE SET name = EXCLUDED.name, last_run_id = EXCLUDED.last_run_id;
+                    UPDATE SET name = EXCLUDED.name, last_csm_run_id = EXCLUDED.last_csm_run_id;
             """
+            LOGGER.debug(runner)
             curs.execute(
                 sql_upsert,
                 (
                     runner.get("id"),
                     runner.get("name"),
-                    runner.get("lastRunId"),
+                    runner.get("lastRunInfo").get("lastRunId"),
                     runner.get("runTemplateId"),
                 ),
             )
             conn.commit()
             LOGGER.info(T("coal.services.postgresql.metadata_updated"))
-    return runner.get("lastRunId")
+    return runner.get("lastRunInfo").get("lastRunId")
 
 
 def remove_runner_metadata_from_postgresql(
@@ -97,8 +96,6 @@ def remove_runner_metadata_from_postgresql(
     # 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,
     )
 
@@ -108,7 +105,7 @@ def remove_runner_metadata_from_postgresql(
             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")};
+                WHERE last_csm_run_id={runner.get("lastRunId")};
             """
             curs.execute(sql_delete_from_metatable)
             conn.commit()