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

cosmotech/coal/store/output/aws_channel.py

@@ -0,0 +1,73 @@
+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
+from cosmotech.coal.store.store import Store
+from cosmotech.coal.utils.configuration import Configuration
+from cosmotech.coal.utils.logger import LOGGER
+
+
+class AwsChannel(ChannelInterface):
+    required_keys = {
+        "cosmotech": [
+            "dataset_absolute_path",
+        ],
+        "s3": ["access_key_id", "endpoint_url", "secret_access_key"],
+    }
+    requirement_string = required_keys
+
+    def __init__(self, dct: dict = None):
+        self.configuration = Configuration(dct)
+        self._s3 = S3(self.configuration)
+
+    def send(self, tables_filter: Optional[list[str]] = None) -> bool:
+
+        _s = Store(store_location=self.configuration.cosmotech.parameters_absolute_path)
+
+        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 tables_filter:
+                tables = [t for t in tables if t in tables_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,42 @@
+from typing import Optional
+
+from cosmotech.coal.azure.blob import dump_store_to_azure
+from cosmotech.coal.store.output.channel_interface import ChannelInterface
+from cosmotech.coal.utils.configuration import Configuration
+
+
+class AzureStorageChannel(ChannelInterface):
+    required_keys = {
+        "cosmotech": [
+            "dataset_absolute_path",
+        ],
+        "azure": [
+            "account_name",
+            "container_name",
+            "tenant_id",
+            "client_id",
+            "client_secret",
+            "output_type",
+            "file_prefix",
+        ],
+    }
+    requirement_string = required_keys
+
+    def __init__(self, dct: dict = None):
+        self.configuration = Configuration(dct)
+
+    def send(self, tables_filter: Optional[list[str]] = None) -> bool:
+        dump_store_to_azure(
+            store_folder=self.configuration.cosmotech.dataset_absolute_path,
+            account_name=self.configuration.azure.account_name,
+            container_name=self.configuration.azure.container_name,
+            tenant_id=self.configuration.azure.tenant_id,
+            client_id=self.configuration.azure.client_id,
+            client_secret=self.configuration.azure.client_secret,
+            output_type=self.configuration.azure.output_type,
+            file_prefix=self.configuration.azure.file_prefix,
+            selected_tables=tables_filter,
+        )
+
+    def delete(self):
+        pass

cosmotech/coal/store/output/channel_interface.py

@@ -0,0 +1,23 @@
+from typing import Optional
+
+from cosmotech.orchestrator.utils.translate import T
+
+
+class ChannelInterface:
+    required_keys = {}
+    requirement_string: str = T("coal.store.output.data_interface.requirements")
+
+    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

cosmotech/coal/store/output/channel_spliter.py

@@ -0,0 +1,55 @@
+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 Configuration
+from cosmotech.coal.utils.logger import LOGGER
+
+
+class ChannelSpliter(ChannelInterface):
+    requirement_string: str = "(Requires any working interface)"
+    targets = []
+    available_interfaces: dict[str, ChannelInterface] = {
+        "s3": AwsChannel,
+        "az_storage": AzureStorageChannel,
+        "postgres": PostgresChannel,
+    }
+
+    def __init__(self, configuration: Configuration):
+        self.configuration = configuration
+        self.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__))
+        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__))
+        return any_ok

cosmotech/coal/store/output/postgres_channel.py

@@ -0,0 +1,40 @@
+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
+from cosmotech.coal.utils.configuration import Configuration
+
+
+class PostgresChannel(ChannelInterface):
+    required_keys = {
+        "cosmotech": ["dataset_absolute_path", "organization_id", "workspace_id", "runner_id"],
+        "postgres": [
+            "host",
+            "post",
+            "db_name",
+            "db_schema",
+            "user_name",
+            "user_password",
+        ],
+    }
+    requirement_string = required_keys
+
+    def __init__(self, dct: dict = None):
+        self.configuration = Configuration(dct)
+
+    def send(self, tables_filter: Optional[list[str]] = None) -> bool:
+        run_id = send_runner_metadata_to_postgresql(self.configuration)
+        dump_store_to_postgresql_from_conf(
+            self.configuration,
+            store_folder=self.configuration.cosmotech.dataset_absolute_path,
+            selected_tables=tables_filter,
+            fk_id=run_id,
+        )
+
+    def delete(self):
+        # removing metadata will trigger cascade delete on real data
+        remove_runner_metadata_from_postgresql(self.configuration)

cosmotech/coal/utils/configuration.py

@@ -0,0 +1,169 @@
+import os
+import tomllib
+
+from cosmotech.orchestrator.utils.translate import T
+
+from cosmotech.coal.utils.logger import LOGGER
+
+
+class Dotdict(dict):
+    """dot.notation access to dictionary attributes"""
+
+    def __setattr__(self, key, val):
+        dd = Dotdict(val, root=self.root) if isinstance(val, dict) else val
+        self.__setitem__(key, dd)
+
+    def __getattr__(self, key):
+        _v = self.__getitem__(key)
+        if isinstance(_v, str) and _v.startswith("$"):
+            _r = self.root
+            for _p in _v[1:].split("."):
+                _r = _r[_p]
+            return _r
+        return _v
+
+    __delattr__ = dict.__delitem__
+
+    def __init__(self, dct=dict(), root=None):
+        object.__setattr__(self, "root", self if root is None else root)
+
+        def update(data):
+            if isinstance(data, dict):
+                return Dotdict(data, root=self.root)
+            elif isinstance(data, list):
+                return [update(d) for d in data]
+            return data
+
+        for k, v in dct.items():
+            self[k] = update(v)
+
+    def merge(self, d):
+        for k, v in d.items():
+            if isinstance(v, Dotdict) and k in self:
+                self[k].merge(v)
+            else:
+                self[k] = v
+
+
+class Configuration(Dotdict):
+
+    # HARD CODED ENVVAR CONVERSION
+    CONVERSION_DICT = {
+        "outputs": [
+            {
+                "type": "postgres",
+                "conf": {
+                    "cosmotech": {
+                        "dataset_absolute_path": "$cosmotech.dataset_absolute_path",
+                        "organization_id": "$cosmotech.organization_id",
+                        "workspace_id": "$cosmotech.workspace_id",
+                        "runner_id": "$cosmotech.runner_id",
+                    },
+                    "postgres": {
+                        "host": "$postgres.host",
+                        "post": "$postgres.post",
+                        "db_name": "$postgres.db_name",
+                        "db_schema": "$postgres.db_schema",
+                        "user_name": "$postgres.user_name",
+                        "user_password": "$postgres.user_password",
+                    },
+                },
+            }
+        ],
+        "secrets": {
+            "log_level": "LOG_LEVEL",
+            "s3": {
+                "access_key_id": "AWS_ACCESS_KEY_ID",
+                "endpoint_url": "AWS_ENDPOINT_URL",
+                "secret_access_key": "AWS_SECRET_ACCESS_KEY",
+                "bucket_prefix": "CSM_DATA_BUCKET_PREFIX",
+            },
+            "azure": {
+                "account_name": "AZURE_ACCOUNT_NAME",
+                "client_id": "AZURE_CLIENT_ID",
+                "client_secret": "AZURE_CLIENT_SECRET",
+                "container_name": "AZURE_CONTAINER_NAME",
+                "data_explorer_database_name": "AZURE_DATA_EXPLORER_DATABASE_NAME",
+                "data_explorer_resource_ingest_uri": "AZURE_DATA_EXPLORER_RESOURCE_INGEST_URI",
+                "data_explorer_resource_uri": "AZURE_DATA_EXPLORER_RESOURCE_URI",
+                "storage_blob_name": "AZURE_STORAGE_BLOB_NAME",
+                "storage_sas_url": "AZURE_STORAGE_SAS_URL",
+                "tenant_id": "AZURE_TENANT_ID",
+            },
+            "cosmotech": {
+                "api_url": "CSM_API_URL",
+                "container_mode": "CSM_CONTAINER_MODE",
+                "data_adx_tag": "CSM_DATA_ADX_TAG",
+                "data_adx_wait_ingestion": "CSM_DATA_ADX_WAIT_INGESTION",
+                "data_blob_prefix": "CSM_DATA_BLOB_PREFIX",
+                "data_bucket_name": "CSM_DATA_BUCKET_NAME",
+                "data_prefix": "CSM_DATA_PREFIX",
+                "dataset_absolute_path": "CSM_DATASET_ABSOLUTE_PATH",
+                "organization_id": "CSM_ORGANIZATION_ID",
+                "parameters_absolute_path": "CSM_PARAMETERS_ABSOLUTE_PATH",
+                "run_id": "CSM_RUN_ID",
+                "runner_id": "CSM_RUNNER_ID",
+                "run_template_id": "CSM_RUN_TEMPLATE_ID",
+                "s3_ca_bundle": "CSM_S3_CA_BUNDLE",
+                "scenario_id": "CSM_SCENARIO_ID",
+                "send_datawarehouse_datasets": "CSM_SEND_DATAWAREHOUSE_DATASETS",
+                "send_datawarehouse_parameters": "CSM_SEND_DATAWAREHOUSE_PARAMETERS",
+                "workspace_id": "CSM_WORKSPACE_ID",
+                "fetch_dataset": "FETCH_DATASET",
+                "fetch_datasets_in_parallel": "FETCH_DATASETS_IN_PARALLEL",
+            },
+            "postgres": {
+                "db_name": "POSTGRES_DB_NAME",
+                "db_schema": "POSTGRES_DB_SCHEMA",
+                "port": "POSTGRES_HOST_PORT",
+                "host": "POSTGRES_HOST_URI",
+                "user_name": "POSTGRES_USER_NAME",
+                "user_password": "POSTGRES_USER_PASSWORD",
+                "password_encoding": "CSM_PSQL_FORCE_PASSWORD_ENCODING",
+            },
+            "single_store": {
+                "db": "SINGLE_STORE_DB",
+                "host": "SINGLE_STORE_HOST",
+                "password": "SINGLE_STORE_PASSWORD",
+                "port": "SINGLE_STORE_PORT",
+                "tables": "SINGLE_STORE_TABLES",
+                "username": "SINGLE_STORE_USERNAME",
+            },
+        },
+    }
+
+    def __init__(self, dct: dict = None):
+        if dct:
+            super().__init__(dct)
+        elif config_path := os.environ.get("CONFIG_FILE_PATH", default=None):
+            with open(config_path, "rb") as f:
+                super().__init__(tomllib.load(f))
+        else:
+            LOGGER.info(T("coal.utils.configuration.no_config_file"))
+            super().__init__(self.CONVERSION_DICT)
+
+        if "secrets" in self:
+            self.secrets = self._env_swap_recusion(self.secrets)
+            # set secret section back to respective section
+            self.merge(self.secrets)
+            del self.secrets
+
+    # convert value to env
+    def _env_swap_recusion(self, dic):
+        for k, v in dic.items():
+            if isinstance(v, Dotdict):
+                dic[k] = self._env_swap_recusion(v)
+                # remove value not found
+                dic[k] = {k: v for k, v in dic[k].items() if v is not None}
+            elif isinstance(v, list):
+                dic[k] = list(self._env_swap_recusion(_v) for _v in v)
+            elif isinstance(v, str):
+                dic[k] = os.environ.get(v)
+        return dic
+
+    def merge_in(self, dic):
+        trans_dic = self._env_swap_recusion(dic)
+        self._merge(trans_dic)
+
+
+ENVIRONMENT_CONFIGURATION = Configuration()

cosmotech/coal/utils/decorator.py

@@ -0,0 +1,22 @@
+import time
+from functools import wraps
+
+from cosmotech.orchestrator.utils.translate import T
+
+from cosmotech.coal.utils.logger import LOGGER
+
+
+def timed(operation, debug=False):
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            process_start = time.time()
+            r = func(*args, **kwargs)
+            process_time = time.time() - process_start
+            log_function = LOGGER.debug if debug else LOGGER.info
+            log_function(T("coal.common.timing.operation_completed").format(operation=operation, time=process_time))
+            return r
+
+        return wrapper
+
+    return decorator

cosmotech/csm_data/commands/api/api.py

@@ -4,25 +4,18 @@
 # 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.orchestrator.utils.translate import T
+
+from cosmotech.coal.cosmotech_api.objects.connection import Connection
+from cosmotech.coal.utils.logger import LOGGER
 from cosmotech.csm_data.commands.api.postgres_send_runner_metadata import (
     postgres_send_runner_metadata,
 )
-from cosmotech.csm_data.commands.api.rds_load_csv import rds_load_csv
-from cosmotech.csm_data.commands.api.rds_send_csv import rds_send_csv
-from cosmotech.csm_data.commands.api.rds_send_store import rds_send_store
 from cosmotech.csm_data.commands.api.run_load_data import run_load_data
-from cosmotech.csm_data.commands.api.runtemplate_load_handler import (
-    runtemplate_load_handler,
-)
-from cosmotech.csm_data.commands.api.tdl_load_files import tdl_load_files
-from cosmotech.csm_data.commands.api.tdl_send_files import tdl_send_files
 from cosmotech.csm_data.commands.api.wsf_load_file import wsf_load_file
 from cosmotech.csm_data.commands.api.wsf_send_file import wsf_send_file
 from cosmotech.csm_data.utils.click import click
 from cosmotech.csm_data.utils.decorators import translate_help, web_help
-from cosmotech.coal.cosmotech_api.connection import get_api_client
-from cosmotech.coal.utils.logger import LOGGER
-from cosmotech.orchestrator.utils.translate import T
 
 
 @click.group(invoke_without_command=True)
@@ -32,19 +25,13 @@ from cosmotech.orchestrator.utils.translate import T
 def api(ctx: click.Context):
     if ctx.invoked_subcommand is None:
         try:
-            api_client, description = get_api_client()
-            LOGGER.info(T("coal.cosmotech_api.connection.found_valid").format(type=description))
+            connection = Connection()
+            LOGGER.info(T("coal.cosmotech_api.connection.found_valid").format(type=connection.api_type))
         except EnvironmentError:
             raise click.Abort()
 
 
-api.add_command(rds_send_csv, "rds-send-csv")
-api.add_command(rds_send_store, "rds-send-store")
-api.add_command(rds_load_csv, "rds-load-csv")
 api.add_command(wsf_send_file, "wsf-send-file")
 api.add_command(wsf_load_file, "wsf-load-file")
-api.add_command(tdl_send_files, "tdl-send-files")
-api.add_command(tdl_load_files, "tdl-load-files")
-api.add_command(runtemplate_load_handler, "runtemplate-load-handler")
 api.add_command(run_load_data, "run-load-data")
 api.add_command(postgres_send_runner_metadata, "postgres-send-runner-metadata")

cosmotech/csm_data/commands/api/postgres_send_runner_metadata.py

@@ -5,10 +5,11 @@
 # etc., to any person is prohibited unless it has been previously and
 # specifically authorized by written means by Cosmo Tech.
 
-from cosmotech.csm_data.utils.click import click
-from cosmotech.csm_data.utils.decorators import web_help, translate_help
 from cosmotech.orchestrator.utils.translate import T
 
+from cosmotech.csm_data.utils.click import click
+from cosmotech.csm_data.utils.decorators import translate_help, web_help
+
 
 @click.command()
 @web_help("csm-data/api/postgres-send-runner-metadata")
@@ -93,7 +94,7 @@ from cosmotech.orchestrator.utils.translate import T
 @click.option(
     "--encode-password/--no-encode-password",
     "force_encode",
-    help=T("csm_data.commands.store.postgres_send_runner_metadata.parameters.encode_password"),
+    help=T("csm_data.commands.api.postgres_send_runner_metadata.parameters.encode_password"),
     envvar="CSM_PSQL_FORCE_PASSWORD_ENCODING",
     show_envvar=True,
     default=True,
@@ -115,17 +116,20 @@ def postgres_send_runner_metadata(
 ):
     # Import the function at the start of the command
     from cosmotech.coal.postgresql import send_runner_metadata_to_postgresql
+    from cosmotech.coal.utils.configuration import Configuration
+
+    _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
+
+    _c.cosmotech.organization_id = organization_id
+    _c.cosmotech.workspace_id = workspace_id
+    _c.cosmotech.runner_id = runner_id
 
-    send_runner_metadata_to_postgresql(
-        organization_id=organization_id,
-        workspace_id=workspace_id,
-        runner_id=runner_id,
-        table_prefix=table_prefix,
-        postgres_host=postgres_host,
-        postgres_port=postgres_port,
-        postgres_db=postgres_db,
-        postgres_schema=postgres_schema,
-        postgres_user=postgres_user,
-        postgres_password=postgres_password,
-        force_encode=force_encode,
-    )
+    send_runner_metadata_to_postgresql(_c)

cosmotech/csm_data/commands/api/run_load_data.py

@@ -5,11 +5,11 @@
 # etc., to any person is prohibited unless it has been previously and
 # specifically authorized by written means by Cosmo Tech.
 
-from cosmotech.csm_data.utils.click import click
-from cosmotech.csm_data.utils.decorators import require_env
-from cosmotech.csm_data.utils.decorators import web_help, translate_help
 from cosmotech.orchestrator.utils.translate import T
 
+from cosmotech.csm_data.utils.click import click
+from cosmotech.csm_data.utils.decorators import require_env, translate_help, web_help
+
 
 @click.command()
 @click.option(
@@ -52,38 +52,6 @@ from cosmotech.orchestrator.utils.translate import T
     help=T("csm_data.commands.api.run_load_data.parameters.parameters_absolute_path"),
     required=True,
 )
-@click.option(
-    "--write-json/--no-write-json",
-    envvar="WRITE_JSON",
-    show_envvar=True,
-    default=True,
-    show_default=True,
-    help=T("csm_data.commands.api.run_load_data.parameters.write_json"),
-)
-@click.option(
-    "--write-csv/--no-write-csv",
-    envvar="WRITE_CSV",
-    show_envvar=True,
-    default=False,
-    show_default=True,
-    help=T("csm_data.commands.api.run_load_data.parameters.write_csv"),
-)
-@click.option(
-    "--fetch-dataset/--no-fetch-dataset",
-    envvar="FETCH_DATASET",
-    show_envvar=True,
-    default=True,
-    show_default=True,
-    help=T("csm_data.commands.api.run_load_data.parameters.fetch_dataset"),
-)
-@click.option(
-    "--parallel/--no-parallel",
-    envvar="FETCH_DATASETS_IN_PARALLEL",
-    show_envvar=True,
-    default=True,
-    show_default=True,
-    help=T("csm_data.commands.api.run_load_data.parameters.parallel"),
-)
 @require_env("CSM_API_SCOPE", "The identification scope of a Cosmotech API")
 @require_env("CSM_API_URL", "The URL to a Cosmotech API")
 @web_help("csm-data/api/run-load-data")
@@ -94,25 +62,18 @@ def run_load_data(
     organization_id: str,
     dataset_absolute_path: str,
     parameters_absolute_path: str,
-    write_json: bool,
-    write_csv: bool,
-    fetch_dataset: bool,
-    parallel: bool,
 ):
     # Import the function at the start of the command
-    from cosmotech.coal.cosmotech_api.runner.download import download_runner_data
+    from cosmotech.coal.cosmotech_api.apis.runner import RunnerApi
+
+    _r = RunnerApi()
 
-    return download_runner_data(
+    return _r.download_runner_data(
         organization_id=organization_id,
         workspace_id=workspace_id,
         runner_id=runner_id,
         parameter_folder=parameters_absolute_path,
         dataset_folder=dataset_absolute_path,
-        read_files=False,
-        parallel=parallel,
-        write_json=write_json,
-        write_csv=write_csv,
-        fetch_dataset=fetch_dataset,
     )
 
 

cosmotech/csm_data/commands/api/wsf_load_file.py

@@ -6,13 +6,11 @@
 # specifically authorized by written means by Cosmo Tech.
 import pathlib
 
-from cosmotech.csm_data.utils.click import click
-from cosmotech.csm_data.utils.decorators import web_help, translate_help
-from cosmotech.coal.cosmotech_api.connection import get_api_client
-from cosmotech.coal.cosmotech_api.workspace import download_workspace_file
-from cosmotech.coal.cosmotech_api.workspace import list_workspace_files
 from cosmotech.orchestrator.utils.translate import T
 
+from cosmotech.csm_data.utils.click import click
+from cosmotech.csm_data.utils.decorators import translate_help, web_help
+
 
 @click.command()
 @click.option(
@@ -53,14 +51,15 @@ from cosmotech.orchestrator.utils.translate import T
 @web_help("csm-data/api/wsf-load-file")
 @translate_help("csm_data.commands.api.wsf_load_file.description")
 def wsf_load_file(organization_id, workspace_id, workspace_path: str, target_folder: str):
-    with get_api_client()[0] as api_client:
-        target_list = list_workspace_files(api_client, organization_id, workspace_id, workspace_path)
+    from cosmotech.coal.cosmotech_api.apis.workspace import WorkspaceApi
+
+    ws_api = WorkspaceApi()
+    target_list = ws_api.list_filtered_workspace_files(organization_id, workspace_id, workspace_path)
 
-        for target in target_list:
-            download_workspace_file(
-                api_client,
-                organization_id,
-                workspace_id,
-                target,
-                pathlib.Path(target_folder),
-            )
+    for target in target_list:
+        ws_api.download_workspace_file(
+            organization_id,
+            workspace_id,
+            target,
+            pathlib.Path(target_folder),
+        )