cosmotech-acceleration-library 1.0.0__py3-none-any.whl

cosmotech/coal/azure/storage.py

@@ -0,0 +1,74 @@
+# 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.
+
+"""
+Azure Storage operations module.
+
+This module provides functions for interacting with Azure Storage, including
+uploading files to blob storage.
+"""
+
+import pathlib
+
+from azure.storage.blob import ContainerClient
+
+from cosmotech.coal.utils.logger import LOGGER
+from cosmotech.orchestrator.utils.translate import T
+
+
+def upload_file(
+    file_path: pathlib.Path,
+    blob_name: str,
+    az_storage_sas_url: str,
+    file_prefix: str = "",
+) -> None:
+    """
+    Upload a single file to Azure Blob Storage.
+
+    Args:
+        file_path: Path to the file to upload
+        blob_name: Name of the blob container
+        az_storage_sas_url: SAS URL for the Azure Storage account
+        file_prefix: Prefix to add to the file name in the blob
+    """
+    uploaded_file_name = blob_name + "/" + file_prefix + file_path.name
+    LOGGER.info(T("coal.common.data_transfer.file_sent").format(file_path=file_path, uploaded_name=uploaded_file_name))
+    ContainerClient.from_container_url(az_storage_sas_url).upload_blob(
+        uploaded_file_name, file_path.open("rb"), overwrite=True
+    )
+
+
+def upload_folder(
+    source_folder: str,
+    blob_name: str,
+    az_storage_sas_url: str,
+    file_prefix: str = "",
+    recursive: bool = False,
+) -> None:
+    """
+    Upload files from a folder to Azure Blob Storage.
+
+    Args:
+        source_folder: Path to the folder containing files to upload
+        blob_name: Name of the blob container
+        az_storage_sas_url: SAS URL for the Azure Storage account
+        file_prefix: Prefix to add to the file names in the blob
+        recursive: Whether to recursively upload files from subdirectories
+    """
+    source_path = pathlib.Path(source_folder)
+    if not source_path.exists():
+        LOGGER.error(T("coal.common.file_operations.not_found").format(source_folder=source_folder))
+        raise FileNotFoundError(T("coal.common.file_operations.not_found").format(source_folder=source_folder))
+
+    if source_path.is_dir():
+        _source_name = str(source_path)
+        for _file_path in source_path.glob("**/*" if recursive else "*"):
+            if _file_path.is_file():
+                _file_name = str(_file_path).removeprefix(_source_name).removeprefix("/")
+                upload_file(_file_path, blob_name, az_storage_sas_url, file_prefix)
+    else:
+        upload_file(source_path, blob_name, az_storage_sas_url, file_prefix)

cosmotech/coal/cosmotech_api/__init__.py

@@ -0,0 +1,36 @@
+# 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.
+
+"""
+Cosmotech API integration module.
+
+This module provides functions for interacting with the Cosmotech API.
+"""
+
+# Re-export functions from the parameters module
+from cosmotech.coal.cosmotech_api.parameters import (
+    write_parameters,
+)
+
+# Re-export functions from the twin_data_layer module
+from cosmotech.coal.cosmotech_api.twin_data_layer import (
+    get_dataset_id_from_runner,
+    send_files_to_tdl,
+    load_files_from_tdl,
+)
+
+# Re-export functions from the run_data module
+from cosmotech.coal.cosmotech_api.run_data import (
+    send_csv_to_run_data,
+    send_store_to_run_data,
+    load_csv_from_run_data,
+)
+
+# Re-export functions from the run_template module
+from cosmotech.coal.cosmotech_api.run_template import (
+    load_run_template_handlers,
+)

cosmotech/coal/cosmotech_api/connection.py

@@ -0,0 +1,96 @@
+# 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 os
+import pathlib
+import ssl
+
+import cosmotech_api
+
+from cosmotech.coal.utils.logger import LOGGER
+from cosmotech.orchestrator.utils.translate import T
+
+api_env_keys = {"CSM_API_KEY", "CSM_API_URL"}
+azure_env_keys = {
+    "AZURE_CLIENT_ID",
+    "AZURE_CLIENT_SECRET",
+    "AZURE_TENANT_ID",
+    "CSM_API_URL",
+    "CSM_API_SCOPE",
+}
+keycloak_env_keys = {
+    "IDP_TENANT_ID",
+    "IDP_CLIENT_ID",
+    "IDP_CLIENT_SECRET",
+    "IDP_BASE_URL",
+    "CSM_API_URL",
+}
+
+
+def get_api_client() -> (cosmotech_api.ApiClient, str):
+    existing_keys = set(os.environ.keys())
+    missing_azure_keys = azure_env_keys - existing_keys
+    missing_api_keys = api_env_keys - existing_keys
+    missing_keycloak_keys = keycloak_env_keys - existing_keys
+    if all((missing_api_keys, missing_azure_keys, missing_keycloak_keys)):
+        LOGGER.error(T("coal.common.errors.no_env_vars"))
+        LOGGER.error(T("coal.cosmotech_api.connection.existing_sets"))
+        LOGGER.error(T("coal.cosmotech_api.connection.azure_connection").format(keys=", ".join(azure_env_keys)))
+        LOGGER.error(T("coal.cosmotech_api.connection.api_key_connection").format(keys=", ".join(api_env_keys)))
+        LOGGER.error(T("coal.cosmotech_api.connection.keycloak_connection").format(keys=", ".join(keycloak_env_keys)))
+        raise EnvironmentError(T("coal.common.errors.no_env_vars"))
+
+    if not missing_keycloak_keys:
+        LOGGER.info(T("coal.cosmotech_api.connection.found_keycloak"))
+        from keycloak import KeycloakOpenID
+
+        server_url = os.environ.get("IDP_BASE_URL")
+        if server_url[-1] != "/":
+            server_url = server_url + "/"
+        keycloack_parameters = dict(
+            server_url=server_url,
+            client_id=os.environ.get("IDP_CLIENT_ID"),
+            realm_name=os.environ.get("IDP_TENANT_ID"),
+            client_secret_key=os.environ.get("IDP_CLIENT_SECRET"),
+        )
+        if (ca_cert_path := os.environ.get("IDP_CA_CERT")) and pathlib.Path(ca_cert_path).exists():
+            LOGGER.info(T("coal.cosmotech_api.connection.found_cert_authority"))
+            keycloack_parameters["verify"] = ca_cert_path
+        keycloak_openid = KeycloakOpenID(**keycloack_parameters)
+
+        access_token = keycloak_openid.token(grant_type="client_credentials")
+
+        configuration = cosmotech_api.Configuration(
+            host=os.environ.get("CSM_API_URL"),
+            access_token=access_token["access_token"],
+        )
+        return cosmotech_api.ApiClient(configuration), "Keycloak Connection"
+
+    if not missing_api_keys:
+        LOGGER.info(T("coal.cosmotech_api.connection.found_api_key"))
+        configuration = cosmotech_api.Configuration(
+            host=os.environ.get("CSM_API_URL"),
+        )
+        return (
+            cosmotech_api.ApiClient(
+                configuration,
+                os.environ.get("CSM_API_KEY_HEADER", "X-CSM-API-KEY"),
+                os.environ.get("CSM_API_KEY"),
+            ),
+            "Cosmo Tech API Key",
+        )
+
+    if not missing_azure_keys:
+        LOGGER.info(T("coal.cosmotech_api.connection.found_azure"))
+        from azure.identity import EnvironmentCredential
+
+        credentials = EnvironmentCredential()
+        token = credentials.get_token(os.environ.get("CSM_API_SCOPE"))
+
+        configuration = cosmotech_api.Configuration(host=os.environ.get("CSM_API_URL"), access_token=token.token)
+        return cosmotech_api.ApiClient(configuration), "Azure Entra Connection"
+
+    raise EnvironmentError(T("coal.common.errors.no_valid_connection"))

cosmotech/coal/cosmotech_api/dataset/__init__.py

@@ -0,0 +1,26 @@
+# 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.
+
+# Re-export all download functions from download submodule
+from cosmotech.coal.cosmotech_api.dataset.download import (
+    download_adt_dataset,
+    download_twingraph_dataset,
+    download_legacy_twingraph_dataset,
+    download_file_dataset,
+    download_dataset_by_id,
+)
+
+from cosmotech.coal.cosmotech_api.dataset.converters import (
+    convert_dataset_to_files,
+    convert_graph_dataset_to_files,
+    convert_file_dataset_to_files,
+)
+
+from cosmotech.coal.cosmotech_api.dataset.utils import (
+    get_content_from_twin_graph_data,
+    sheet_to_header,
+)

cosmotech/coal/cosmotech_api/dataset/converters.py

@@ -0,0 +1,164 @@
+# 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 csv
+import json
+import os
+import tempfile
+from pathlib import Path
+from typing import Dict, List, Any, Optional, Union
+
+from cosmotech.coal.utils.logger import LOGGER
+from cosmotech.orchestrator.utils.translate import T
+from cosmotech.coal.cosmotech_api.dataset.utils import sheet_to_header
+
+
+def convert_dataset_to_files(dataset_info: Dict[str, Any], target_folder: Optional[Union[str, Path]] = None) -> Path:
+    """
+    Convert dataset info to files.
+
+    Args:
+        dataset_info: Dataset info dict with type, content, name
+        target_folder: Optional folder to save files (if None, uses temp dir)
+
+    Returns:
+        Path to folder containing files
+    """
+    dataset_type = dataset_info["type"]
+    content = dataset_info["content"]
+    name = dataset_info["name"]
+
+    LOGGER.info(T("coal.services.dataset.converting_to_files").format(dataset_type=dataset_type, dataset_name=name))
+
+    if target_folder is None:
+        target_folder = Path(tempfile.mkdtemp())
+        LOGGER.debug(T("coal.services.dataset.created_temp_folder").format(folder=target_folder))
+    else:
+        target_folder = Path(target_folder)
+        target_folder.mkdir(parents=True, exist_ok=True)
+        LOGGER.debug(T("coal.services.dataset.using_folder").format(folder=target_folder))
+
+    if dataset_type in ["adt", "twincache"]:
+        return convert_graph_dataset_to_files(content, target_folder)
+    else:
+        return convert_file_dataset_to_files(content, target_folder, dataset_type)
+
+
+def convert_graph_dataset_to_files(
+    content: Dict[str, List[Dict]], target_folder: Optional[Union[str, Path]] = None
+) -> Path:
+    """
+    Convert graph dataset content to CSV files.
+
+    Args:
+        content: Dictionary mapping entity types to lists of entities
+        target_folder: Folder to save files (if None, uses temp dir)
+
+    Returns:
+        Path to folder containing files
+    """
+    if target_folder is None:
+        target_folder = Path(tempfile.mkdtemp())
+        LOGGER.debug(T("coal.services.dataset.created_temp_folder").format(folder=target_folder))
+    else:
+        target_folder = Path(target_folder)
+        target_folder.mkdir(parents=True, exist_ok=True)
+        LOGGER.debug(T("coal.services.dataset.using_folder").format(folder=target_folder))
+    file_count = 0
+
+    LOGGER.info(
+        T("coal.services.dataset.converting_graph_data").format(entity_types=len(content), folder=target_folder)
+    )
+
+    for entity_type, entities in content.items():
+        if not entities:
+            LOGGER.debug(T("coal.services.dataset.skipping_empty_entity").format(entity_type=entity_type))
+            continue
+
+        file_path = target_folder / f"{entity_type}.csv"
+        LOGGER.debug(T("coal.services.dataset.writing_csv").format(file_name=file_path.name, count=len(entities)))
+
+        fieldnames = sheet_to_header(entities)
+
+        with open(file_path, "w", newline="") as file:
+            writer = csv.DictWriter(file, fieldnames=fieldnames, dialect="unix", quoting=csv.QUOTE_MINIMAL)
+            writer.writeheader()
+
+            for entity in entities:
+                # Convert values to strings and handle boolean values
+                row = {
+                    k: str(v).replace("'", '"').replace("True", "true").replace("False", "false")
+                    for k, v in entity.items()
+                }
+                writer.writerow(row)
+
+        file_count += 1
+        LOGGER.debug(T("coal.services.dataset.file_written").format(file_path=file_path))
+
+    LOGGER.info(T("coal.services.dataset.files_created").format(count=file_count, folder=target_folder))
+
+    return target_folder
+
+
+def convert_file_dataset_to_files(
+    content: Dict[str, Any],
+    target_folder: Optional[Union[str, Path]] = None,
+    file_type: str = "",
+) -> Path:
+    """
+    Convert file dataset content to files.
+
+    Args:
+        content: Dictionary mapping file names to content
+        target_folder: Folder to save files (if None, uses temp dir)
+        file_type: Type of file (csv, json, etc.)
+
+    Returns:
+        Path to folder containing files
+    """
+    if target_folder is None:
+        target_folder = Path(tempfile.mkdtemp())
+        LOGGER.debug(T("coal.services.dataset.created_temp_folder").format(folder=target_folder))
+    else:
+        target_folder = Path(target_folder)
+        target_folder.mkdir(parents=True, exist_ok=True)
+        LOGGER.debug(T("coal.services.dataset.using_folder").format(folder=target_folder))
+    file_count = 0
+
+    LOGGER.info(
+        T("coal.services.dataset.converting_file_data").format(
+            file_count=len(content), file_type=file_type, folder=target_folder
+        )
+    )
+
+    for file_name, file_content in content.items():
+        file_path = target_folder / file_name
+
+        # Ensure parent directories exist
+        file_path.parent.mkdir(parents=True, exist_ok=True)
+
+        LOGGER.debug(T("coal.services.dataset.writing_file").format(file_name=file_path.name, file_type=file_type))
+
+        if isinstance(file_content, str):
+            # Text content
+            with open(file_path, "w") as file:
+                file.write(file_content)
+        elif isinstance(file_content, dict) or isinstance(file_content, list):
+            # JSON content
+            with open(file_path, "w") as file:
+                json.dump(file_content, file, indent=2)
+        else:
+            # Other content types
+            with open(file_path, "w") as file:
+                file.write(str(file_content))
+
+        file_count += 1
+        LOGGER.debug(T("coal.services.dataset.file_written").format(file_path=file_path))
+
+    LOGGER.info(T("coal.services.dataset.files_created").format(count=file_count, folder=target_folder))
+
+    return target_folder

cosmotech/coal/cosmotech_api/dataset/download/__init__.py

@@ -0,0 +1,19 @@
+# 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.
+
+"""
+Dataset download submodules.
+"""
+
+# Re-export all download functions
+from cosmotech.coal.cosmotech_api.dataset.download.adt import download_adt_dataset
+from cosmotech.coal.cosmotech_api.dataset.download.twingraph import (
+    download_twingraph_dataset,
+    download_legacy_twingraph_dataset,
+)
+from cosmotech.coal.cosmotech_api.dataset.download.file import download_file_dataset
+from cosmotech.coal.cosmotech_api.dataset.download.common import download_dataset_by_id

cosmotech/coal/cosmotech_api/dataset/download/adt.py

@@ -0,0 +1,119 @@
+# 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 time
+import tempfile
+from pathlib import Path
+from typing import Dict, Any, Optional, Union, Tuple
+
+from azure.digitaltwins.core import DigitalTwinsClient
+from azure.identity import DefaultAzureCredential
+
+from cosmotech.coal.utils.logger import LOGGER
+from cosmotech.orchestrator.utils.translate import T
+from cosmotech.coal.cosmotech_api.connection import get_api_client
+from cosmotech.coal.cosmotech_api.dataset.converters import convert_dataset_to_files
+
+
+def download_adt_dataset(
+    adt_address: str,
+    target_folder: Optional[Union[str, Path]] = None,
+    credentials: Optional[DefaultAzureCredential] = None,
+) -> Tuple[Dict[str, Any], Path]:
+    """
+    Download dataset from Azure Digital Twins.
+
+    Args:
+        adt_address: The ADT instance address
+        target_folder: Optional folder to save files (if None, uses temp dir)
+        credentials: Optional Azure credentials (if None, uses DefaultAzureCredential)
+
+    Returns:
+        Tuple of (content dict, folder path)
+    """
+    start_time = time.time()
+    LOGGER.info(T("coal.services.dataset.download_started").format(dataset_type="ADT"))
+    LOGGER.debug(T("coal.services.dataset.adt_connecting").format(url=adt_address))
+
+    # Create credentials if not provided
+    if credentials is None:
+        if get_api_client()[1] == "Azure Entra Connection":
+            credentials = DefaultAzureCredential()
+        else:
+            LOGGER.error(T("coal.services.dataset.adt_no_credentials"))
+            raise ValueError("No credentials available for ADT connection")
+
+    # Create client and download data
+    client = DigitalTwinsClient(adt_address, credentials)
+
+    # Query twins
+    query_start = time.time()
+    LOGGER.debug(T("coal.services.dataset.adt_querying_twins"))
+    query_expression = "SELECT * FROM digitaltwins"
+    query_result = client.query_twins(query_expression)
+
+    json_content = dict()
+    twin_count = 0
+
+    for twin in query_result:
+        twin_count += 1
+        entity_type = twin.get("$metadata").get("$model").split(":")[-1].split(";")[0]
+        t_content = {k: v for k, v in twin.items()}
+        t_content["id"] = t_content["$dtId"]
+
+        # Remove system properties
+        for k in list(twin.keys()):
+            if k[0] == "$":
+                del t_content[k]
+
+        json_content.setdefault(entity_type, [])
+        json_content[entity_type].append(t_content)
+
+    query_time = time.time() - query_start
+    LOGGER.debug(T("coal.services.dataset.adt_twins_found").format(count=twin_count))
+    LOGGER.debug(T("coal.common.timing.operation_completed").format(operation="twins query", time=query_time))
+
+    # Query relationships
+    rel_start = time.time()
+    LOGGER.debug(T("coal.services.dataset.adt_querying_relations"))
+    relations_query = "SELECT * FROM relationships"
+    query_result = client.query_twins(relations_query)
+
+    relation_count = 0
+    for relation in query_result:
+        relation_count += 1
+        tr = {"$relationshipId": "id", "$sourceId": "source", "$targetId": "target"}
+        r_content = {k: v for k, v in relation.items()}
+
+        # Map system properties to standard names
+        for k, v in tr.items():
+            r_content[v] = r_content[k]
+
+        # Remove system properties
+        for k in list(relation.keys()):
+            if k[0] == "$":
+                del r_content[k]
+
+        json_content.setdefault(relation["$relationshipName"], [])
+        json_content[relation["$relationshipName"]].append(r_content)
+
+    rel_time = time.time() - rel_start
+    LOGGER.debug(T("coal.services.dataset.adt_relations_found").format(count=relation_count))
+    LOGGER.debug(T("coal.common.timing.operation_completed").format(operation="relations query", time=rel_time))
+
+    # Convert to files if target_folder is provided
+    if target_folder:
+        dataset_info = {"type": "adt", "content": json_content, "name": "ADT Dataset"}
+        target_folder = convert_dataset_to_files(dataset_info, target_folder)
+    else:
+        target_folder = tempfile.mkdtemp()
+
+    elapsed_time = time.time() - start_time
+    LOGGER.info(T("coal.common.timing.operation_completed").format(operation="ADT download", time=elapsed_time))
+    LOGGER.info(T("coal.services.dataset.download_completed").format(dataset_type="ADT"))
+
+    return json_content, Path(target_folder)

cosmotech/coal/cosmotech_api/dataset/download/common.py

@@ -0,0 +1,140 @@
+# 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 time
+from pathlib import Path
+from typing import Dict, Any, Optional, Union, Tuple
+
+from cosmotech_api import DatasetApi
+
+from cosmotech.coal.utils.logger import LOGGER
+from cosmotech.orchestrator.utils.translate import T
+from cosmotech.coal.cosmotech_api.connection import get_api_client
+
+# Import specific download functions
+# These imports are defined here to avoid circular imports
+# The functions are imported directly from their modules
+from cosmotech.coal.cosmotech_api.dataset.download.adt import download_adt_dataset
+from cosmotech.coal.cosmotech_api.dataset.download.twingraph import (
+    download_twingraph_dataset,
+    download_legacy_twingraph_dataset,
+)
+from cosmotech.coal.cosmotech_api.dataset.download.file import download_file_dataset
+
+
+def download_dataset_by_id(
+    organization_id: str,
+    workspace_id: str,
+    dataset_id: str,
+    target_folder: Optional[Union[str, Path]] = None,
+) -> Tuple[Dict[str, Any], Path]:
+    """
+    Download dataset by ID.
+
+    Args:
+        organization_id: Organization ID
+        workspace_id: Workspace ID
+        dataset_id: Dataset ID
+        target_folder: Optional folder to save files (if None, uses temp dir)
+
+    Returns:
+        Tuple of (dataset info dict, folder path)
+    """
+    start_time = time.time()
+    LOGGER.info(T("coal.services.dataset.download_started").format(dataset_type="Dataset"))
+    LOGGER.debug(
+        T("coal.services.dataset.dataset_downloading").format(organization_id=organization_id, dataset_id=dataset_id)
+    )
+
+    with get_api_client()[0] as api_client:
+        api_instance = DatasetApi(api_client)
+
+        # Get dataset info
+        info_start = time.time()
+        dataset = api_instance.find_dataset_by_id(organization_id=organization_id, dataset_id=dataset_id)
+        info_time = time.time() - info_start
+
+        LOGGER.debug(
+            T("coal.services.dataset.dataset_info_retrieved").format(dataset_name=dataset.name, dataset_id=dataset_id)
+        )
+        LOGGER.debug(
+            T("coal.common.timing.operation_completed").format(operation="dataset info retrieval", time=info_time)
+        )
+
+        # Determine dataset type and download
+        if dataset.connector is None:
+            parameters = []
+        else:
+            parameters = dataset.connector.parameters_values
+
+        is_adt = "AZURE_DIGITAL_TWINS_URL" in parameters
+        is_storage = "AZURE_STORAGE_CONTAINER_BLOB_PREFIX" in parameters
+        is_legacy_twin_cache = "TWIN_CACHE_NAME" in parameters and dataset.twingraph_id is None
+        is_in_workspace_file = (
+            False if dataset.tags is None else "workspaceFile" in dataset.tags or "dataset_part" in dataset.tags
+        )
+
+        download_start = time.time()
+
+        if is_adt:
+            LOGGER.debug(T("coal.services.dataset.dataset_type_detected").format(type="ADT"))
+            content, folder = download_adt_dataset(
+                adt_address=parameters["AZURE_DIGITAL_TWINS_URL"],
+                target_folder=target_folder,
+            )
+            dataset_type = "adt"
+
+        elif is_legacy_twin_cache:
+            LOGGER.debug(T("coal.services.dataset.dataset_type_detected").format(type="Legacy TwinGraph"))
+            twin_cache_name = parameters["TWIN_CACHE_NAME"]
+            content, folder = download_legacy_twingraph_dataset(
+                organization_id=organization_id,
+                cache_name=twin_cache_name,
+                target_folder=target_folder,
+            )
+            dataset_type = "twincache"
+
+        elif is_storage or is_in_workspace_file:
+            if is_storage:
+                LOGGER.debug(T("coal.services.dataset.dataset_type_detected").format(type="Storage"))
+                _file_name = parameters["AZURE_STORAGE_CONTAINER_BLOB_PREFIX"].replace("%WORKSPACE_FILE%/", "")
+            else:
+                LOGGER.debug(T("coal.services.dataset.dataset_type_detected").format(type="Workspace File"))
+                _file_name = dataset.source.location
+
+            content, folder = download_file_dataset(
+                organization_id=organization_id,
+                workspace_id=workspace_id,
+                file_name=_file_name,
+                target_folder=target_folder,
+            )
+            dataset_type = _file_name.split(".")[-1]
+
+        else:
+            LOGGER.debug(T("coal.services.dataset.dataset_type_detected").format(type="TwinGraph"))
+            content, folder = download_twingraph_dataset(
+                organization_id=organization_id,
+                dataset_id=dataset_id,
+                target_folder=target_folder,
+            )
+            dataset_type = "twincache"
+
+        download_time = time.time() - download_start
+        LOGGER.debug(
+            T("coal.common.timing.operation_completed").format(operation="content download", time=download_time)
+        )
+
+    # Prepare result
+    dataset_info = {"type": dataset_type, "content": content, "name": dataset.name}
+
+    elapsed_time = time.time() - start_time
+    LOGGER.info(
+        T("coal.common.timing.operation_completed").format(operation="total dataset download", time=elapsed_time)
+    )
+    LOGGER.info(T("coal.services.dataset.download_completed").format(dataset_type="Dataset"))
+
+    return dataset_info, folder