cognite-toolkit 0.5.60__py3-none-any.whl → 0.5.62__py3-none-any.whl

cognite_toolkit/_cdf_tk/commands/_profile.py

@@ -1,7 +1,12 @@
-from collections.abc import Callable
+from abc import ABC, abstractmethod
+from collections.abc import Callable, Mapping
 from concurrent.futures import ThreadPoolExecutor, as_completed
+from functools import cached_property
+from typing import ClassVar, Literal, TypeAlias, overload
 
 from cognite.client.exceptions import CogniteException
+from rich import box
+from rich.console import Console
 from rich.live import Live
 from rich.spinner import Spinner
 from rich.table import Table
@@ -23,108 +28,212 @@ from cognite_toolkit._cdf_tk.utils.aggregators import (
 from ._base import ToolkitCommand
 
 
-class ProfileCommand(ToolkitCommand):
-    class Columns:
-        Resource = "Resource"
-        Count = "Count"
-        MetadataKeyCount = "Metadata Key Count"
-        LabelCount = "Label Count"
-        Transformation = "Transformations"
+class WaitingAPICallClass:
+    def __bool__(self) -> bool:
+        return False
 
-    columns = (
-        Columns.Resource,
-        Columns.Count,
-        Columns.MetadataKeyCount,
-        Columns.LabelCount,
-        Columns.Transformation,
-    )
-    spinner_speed = 1.0
 
-    @classmethod
-    def asset_centric(
-        cls,
-        client: ToolkitClient,
-        verbose: bool = False,
-    ) -> list[dict[str, str]]:
-        aggregators: list[AssetCentricAggregator] = [
-            AssetAggregator(client),
-            EventAggregator(client),
-            FileAggregator(client),
-            TimeSeriesAggregator(client),
-            SequenceAggregator(client),
-            RelationshipAggregator(client),
-            LabelCountAggregator(client),
-        ]
-        results, api_calls = cls._create_initial_table(aggregators)
-        with Live(cls.create_profile_table(results), refresh_per_second=4) as live:
-            with ThreadPoolExecutor(max_workers=8) as executor:
-                future_to_cell = {
-                    executor.submit(api_calls[(index, col)]): (index, col)
-                    for index in range(len(aggregators))
-                    for col in cls.columns
-                    if (index, col) in api_calls
+WaitingAPICall = WaitingAPICallClass()
+
+PendingCellValue: TypeAlias = int | float | str | bool | None | WaitingAPICallClass
+CellValue: TypeAlias = int | float | str | bool | None
+PendingTable: TypeAlias = dict[tuple[str, str], PendingCellValue]
+
+
+class ProfileCommand(ToolkitCommand, ABC):
+    def __init__(self, print_warning: bool = True, skip_tracking: bool = False, silent: bool = False) -> None:
+        super().__init__(print_warning, skip_tracking, silent)
+        self.table_title = self.__class__.__name__.removesuffix("Command")
+
+    class Columns:  # Placeholder for columns, subclasses should define their own Columns class
+        ...
+
+    spinner_args: ClassVar[Mapping] = dict(name="arc", text="loading...", style="bold green", speed=1.0)
+
+    max_workers = 8
+    is_dynamic_table = False
+
+    @cached_property
+    def columns(self) -> tuple[str, ...]:
+        return (
+            tuple([attr for attr in self.Columns.__dict__.keys() if not attr.startswith("_")])
+            if hasattr(self, "Columns")
+            else tuple()
+        )
+
+    def create_profile_table(self, client: ToolkitClient) -> list[dict[str, CellValue]]:
+        console = Console()
+        with console.status("Setting up", spinner="aesthetic", speed=0.4) as _:
+            table = self.create_initial_table(client)
+        with (
+            Live(self.draw_table(table), refresh_per_second=4, console=console) as live,
+            ThreadPoolExecutor(max_workers=self.max_workers) as executor,
+        ):
+            while True:
+                current_calls = {
+                    executor.submit(self.call_api(row, col, client)): (row, col)
+                    for (row, col), cell in table.items()
+                    if cell is WaitingAPICall
                 }
-                for future in as_completed(future_to_cell):
-                    index, col = future_to_cell[future]
-                    results[index][col] = future.result()
-                    live.update(cls.create_profile_table(results))
-        return [{col: str(value) for col, value in row.items()} for row in results]
+                if not current_calls:
+                    break
+                for future in as_completed(current_calls):
+                    row, col = current_calls[future]
+                    try:
+                        result = future.result()
+                    except CogniteException as e:
+                        result = type(e).__name__
+                    table[(row, col)] = self.format_result(result, row, col)
+                    if self.is_dynamic_table:
+                        table = self.update_table(table, result, row, col)
+                    live.update(self.draw_table(table))
+        return self.as_record_format(table, allow_waiting_api_call=False)
 
-    @classmethod
-    def _create_initial_table(
-        cls, aggregators: list[AssetCentricAggregator]
-    ) -> tuple[list[dict[str, str | Spinner]], dict[tuple[int, str], Callable[[], str]]]:
-        rows: list[dict[str, str | Spinner]] = []
-        api_calls: dict[tuple[int, str], Callable[[], str]] = {}
-        for index, aggregator in enumerate(aggregators):
-            row: dict[str, str | Spinner] = {
-                cls.Columns.Resource: aggregator.display_name,
-                cls.Columns.Count: Spinner("arc", text="loading...", style="bold green", speed=cls.spinner_speed),
-            }
-            api_calls[(index, cls.Columns.Count)] = cls._call_api(aggregator.count)
-            count: str | Spinner = "-"
-            if isinstance(aggregator, MetadataAggregator):
-                count = Spinner("arc", text="loading...", style="bold green", speed=cls.spinner_speed)
-                api_calls[(index, cls.Columns.MetadataKeyCount)] = cls._call_api(aggregator.metadata_key_count)
-            row[cls.Columns.MetadataKeyCount] = count
+    @abstractmethod
+    def create_initial_table(self, client: ToolkitClient) -> PendingTable:
+        """
+        Create the initial table with placeholders for API calls.
+        Each cell that requires an API call should be initialized with WaitingAPICall.
+        """
+        raise NotImplementedError("Subclasses must implement create_initial_table.")
 
-            count = "-"
-            if isinstance(aggregator, LabelAggregator):
-                count = Spinner("arc", text="loading...", style="bold green", speed=cls.spinner_speed)
-                api_calls[(index, cls.Columns.LabelCount)] = cls._call_api(aggregator.label_count)
-            row[cls.Columns.LabelCount] = count
+    @abstractmethod
+    def call_api(self, row: str, col: str, client: ToolkitClient) -> Callable:
+        raise NotImplementedError("Subclasses must implement call_api.")
 
-            row[cls.Columns.Transformation] = Spinner(
-                "arc", text="loading...", style="bold green", speed=cls.spinner_speed
-            )
-            api_calls[(index, cls.Columns.Transformation)] = cls._call_api(aggregator.transformation_count)
+    def format_result(self, result: object, row: str, col: str) -> CellValue:
+        """
+        Format the result of an API call for display in the table.
+        This can be overridden by subclasses to customize formatting.
+        """
+        if isinstance(result, int | float | bool | str):
+            return result
+        raise NotImplementedError("Subclasses must implement format_result.")
 
-            rows.append(row)
-        return rows, api_calls
+    def update_table(
+        self,
+        current_table: PendingTable,
+        result: object,
+        row: str,
+        col: str,
+    ) -> PendingTable:
+        raise NotImplementedError("Subclasses must implement update_table.")
 
-    @classmethod
-    def create_profile_table(cls, rows: list[dict[str, str | Spinner]]) -> Table:
-        table = Table(
-            title="Asset Centric Profile",
+    def draw_table(self, table: PendingTable) -> Table:
+        rich_table = Table(
+            title=self.table_title,
             title_justify="left",
             show_header=True,
             header_style="bold magenta",
+            box=box.MINIMAL,
         )
-        for col in cls.columns:
-            table.add_column(col)
+        for col in self.columns:
+            rich_table.add_column(col)
+
+        rows = self.as_record_format(table)
 
         for row in rows:
-            table.add_row(*row.values())
-        return table
+            rich_table.add_row(*[self._as_cell(value) for value in row.values()])
+        return rich_table
+
+    @classmethod
+    @overload
+    def as_record_format(
+        cls, table: PendingTable, allow_waiting_api_call: Literal[True] = True
+    ) -> list[dict[str, PendingCellValue]]: ...
+
+    @classmethod
+    @overload
+    def as_record_format(
+        cls,
+        table: PendingTable,
+        allow_waiting_api_call: Literal[False],
+    ) -> list[dict[str, CellValue]]: ...
 
-    @staticmethod
-    def _call_api(call_fun: Callable[[], int]) -> Callable[[], str]:
-        def styled_callable() -> str:
-            try:
-                value = call_fun()
-            except CogniteException as e:
-                return type(e).__name__
+    @classmethod
+    def as_record_format(
+        cls,
+        table: PendingTable,
+        allow_waiting_api_call: bool = True,
+    ) -> list[dict[str, PendingCellValue]] | list[dict[str, CellValue]]:
+        rows: list[dict[str, PendingCellValue]] = []
+        row_indices: dict[str, int] = {}
+        for (row, col), value in table.items():
+            if value is WaitingAPICall and not allow_waiting_api_call:
+                value = None
+            if row not in row_indices:
+                row_indices[row] = len(rows)
+                rows.append({col: value})
             else:
-                return f"{value:,}"
+                rows[row_indices[row]][col] = value
+        return rows
+
+    def _as_cell(self, value: PendingCellValue) -> str | Spinner:
+        if isinstance(value, WaitingAPICallClass):
+            return Spinner(**self.spinner_args)
+        elif isinstance(value, int):
+            return f"{value:,}"
+        elif isinstance(value, float):
+            return f"{value:.2f}"
+        elif value is None:
+            return "-"
+        return str(value)
+
+
+class ProfileAssetCentricCommand(ProfileCommand):
+    def __init__(self, print_warning: bool = True, skip_tracking: bool = False, silent: bool = False) -> None:
+        super().__init__(print_warning, skip_tracking, silent)
+        self.table_title = "Asset Centric Profile"
+        self.aggregators: dict[str, AssetCentricAggregator] = {}
+
+    class Columns:
+        Resource = "Resource"
+        Count = "Count"
+        MetadataKeyCount = "Metadata Key Count"
+        LabelCount = "Label Count"
+        Transformation = "Transformations"
+
+    def asset_centric(self, client: ToolkitClient, verbose: bool = False) -> list[dict[str, CellValue]]:
+        self.aggregators.update(
+            {
+                agg.display_name: agg
+                for agg in [
+                    AssetAggregator(client),
+                    EventAggregator(client),
+                    FileAggregator(client),
+                    TimeSeriesAggregator(client),
+                    SequenceAggregator(client),
+                    RelationshipAggregator(client),
+                    LabelCountAggregator(client),
+                ]
+            }
+        )
+        return self.create_profile_table(client)
+
+    def create_initial_table(self, client: ToolkitClient) -> PendingTable:
+        table: dict[tuple[str, str], str | int | float | bool | None | WaitingAPICallClass] = {}
+        for index, aggregator in self.aggregators.items():
+            table[(index, self.Columns.Resource)] = aggregator.display_name
+            table[(index, self.Columns.Count)] = WaitingAPICall
+            if isinstance(aggregator, MetadataAggregator):
+                table[(index, self.Columns.MetadataKeyCount)] = WaitingAPICall
+            else:
+                table[(index, self.Columns.MetadataKeyCount)] = None
+            if isinstance(aggregator, LabelAggregator):
+                table[(index, self.Columns.LabelCount)] = WaitingAPICall
+            else:
+                table[(index, self.Columns.LabelCount)] = None
+            table[(index, self.Columns.Transformation)] = WaitingAPICall
+        return table
 
-        return styled_callable
+    def call_api(self, row: str, col: str, client: ToolkitClient) -> Callable:
+        aggregator = self.aggregators[row]
+        if col == self.Columns.Count:
+            return aggregator.count
+        elif col == self.Columns.MetadataKeyCount and isinstance(aggregator, MetadataAggregator):
+            return aggregator.metadata_key_count
+        elif col == self.Columns.LabelCount and isinstance(aggregator, LabelAggregator):
+            return aggregator.label_count
+        elif col == self.Columns.Transformation:
+            return aggregator.transformation_count
+        raise ValueError(f"Unknown column: {col} for row: {row}")

cognite_toolkit/_cdf_tk/commands/modules.py

@@ -1,12 +1,17 @@
 from __future__ import annotations
 
 import shutil
+import tempfile
+import zipfile
 from collections import Counter
+from hashlib import sha256
 from importlib import resources
 from pathlib import Path
+from types import TracebackType
 from typing import Any, Literal, Optional
 
 import questionary
+import requests
 import typer
 from packaging.version import Version
 from packaging.version import parse as parse_version
@@ -14,7 +19,7 @@ from rich import print
 from rich.markdown import Markdown
 from rich.padding import Padding
 from rich.panel import Panel
-from rich.progress import track
+from rich.progress import Progress, track
 from rich.rule import Rule
 from rich.table import Table
 from rich.tree import Tree
@@ -47,7 +52,8 @@ from cognite_toolkit._cdf_tk.data_classes import (
     Package,
     Packages,
 )
-from cognite_toolkit._cdf_tk.exceptions import ToolkitRequiredValueError, ToolkitValueError
+from cognite_toolkit._cdf_tk.exceptions import ToolkitError, ToolkitRequiredValueError, ToolkitValueError
+from cognite_toolkit._cdf_tk.feature_flags import Flags
 from cognite_toolkit._cdf_tk.hints import verify_module_directory
 from cognite_toolkit._cdf_tk.tk_warnings import MediumSeverityWarning
 from cognite_toolkit._cdf_tk.utils import humanize_collection, read_yaml_file
@@ -85,6 +91,28 @@ class ModulesCommand(ToolkitCommand):
     def __init__(self, print_warning: bool = True, skip_tracking: bool = False, silent: bool = False):
         super().__init__(print_warning, skip_tracking, silent)
         self._builtin_modules_path = Path(resources.files(cognite_toolkit.__name__)) / BUILTIN_MODULES  # type: ignore [arg-type]
+        self._temp_download_dir = Path(tempfile.gettempdir()) / "library_downloads"
+        if not self._temp_download_dir.exists():
+            self._temp_download_dir.mkdir(parents=True, exist_ok=True)
+
+    def __enter__(self) -> ModulesCommand:
+        """
+        Context manager to ensure the temporary download directory is cleaned up after use. It requires the command to be used in a `with` block.
+        """
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,  # Type of the exception
+        exc_value: BaseException | None,  # Exception instance
+        traceback: TracebackType | None,  # Traceback object
+    ) -> None:
+        """
+        Clean up the temporary download directory.
+        """
+
+        if self._temp_download_dir.exists():
+            safe_rmtree(self._temp_download_dir)
 
     @classmethod
     def _create_tree(cls, item: Packages) -> Tree:
@@ -128,6 +156,7 @@ class ModulesCommand(ToolkitCommand):
         downloader_by_repo: dict[str, FileDownloader] = {}
 
         extra_resources: set[Path] = set()
+
         for package_name, package in selected_packages.items():
             print(f"{INDENT}[{'yellow' if mode == 'clean' else 'green'}]Creating {package_name}[/]")
 
@@ -280,7 +309,8 @@ default_organization_dir = "{organization_dir.name}"''',
             organization_dir = Path(organization_dir_raw.strip())
 
         modules_root_dir = organization_dir / MODULES
-        packages = Packages().load(self._builtin_modules_path)
+
+        packages = self._get_available_packages()
 
         if select_all:
             print(Panel("Instantiating all available modules"))
@@ -680,9 +710,102 @@ default_organization_dir = "{organization_dir.name}"''',
             build_env = default.environment.validation_type
 
         existing_module_names = [module.name for module in ModuleResources(organization_dir, build_env).list()]
-        available_packages = Packages().load(self._builtin_modules_path)
-
+        available_packages = self._get_available_packages()
         added_packages = self._select_packages(available_packages, existing_module_names)
 
         download_data = self._get_download_data(added_packages)
         self._create(organization_dir, added_packages, environments, "update", download_data)
+
+    def _get_available_packages(self) -> Packages:
+        """
+        Returns a list of available packages, either from the CDF TOML file or from external libraries if the feature flag is enabled.
+        If the feature flag is not enabled and no libraries are specified, it returns the built-in modules.
+        """
+
+        cdf_toml = CDFToml.load()
+        if Flags.EXTERNAL_LIBRARIES.is_enabled() and cdf_toml.libraries:
+            for library_name, library in cdf_toml.libraries.items():
+                try:
+                    print(f"[green]Adding library {library_name}[/]")
+                    file_path = self._temp_download_dir / f"{library_name}.zip"
+                    self._download(library.url, file_path)
+                    self._validate_checksum(library.checksum, file_path)
+                    self._unpack(file_path)
+                    return Packages().load(file_path.parent)
+                except Exception as e:
+                    if isinstance(e, ToolkitError):
+                        raise e
+                    else:
+                        raise ToolkitError(
+                            f"An unexpected error occurred during downloading {library.url} to {file_path}: {e}"
+                        ) from e
+
+                    raise ToolkitError(f"Failed to add library {library_name}, {e}")
+            # If no libraries are specified or the flag is not enabled, load the built-in modules
+            raise ValueError("No valid libraries found.")
+        else:
+            return Packages.load(self._builtin_modules_path)
+
+    def _download(self, url: str, file_path: Path) -> None:
+        """
+        Downloads a file from a URL to the specified output path.
+        If the file already exists, it skips the download.
+        """
+        try:
+            response = requests.get(url, stream=True)
+            response.raise_for_status()  # Raise an exception for HTTP errors
+
+            total_size = int(response.headers.get("content-length", 0))
+
+            with Progress() as progress:
+                task = progress.add_task("Download", total=total_size)
+                with open(file_path, "wb") as f:
+                    for chunk in response.iter_content(chunk_size=8192):
+                        f.write(chunk)
+                        progress.update(task, advance=len(chunk))
+
+        except requests.exceptions.RequestException as e:
+            raise ToolkitError(f"Error downloading file from {url}: {e}") from e
+
+    def _validate_checksum(self, checksum: str, file_path: Path) -> None:
+        """
+        Compares the checksum of the downloaded file with the expected checksum.
+        """
+
+        if checksum.lower().startswith("sha256:"):
+            checksum = checksum[7:]
+        else:
+            raise ToolkitValueError(f"Unsupported checksum format: {checksum}. Expected 'sha256:' prefix")
+
+        chunk_size: int = 8192
+        sha256_hash = sha256()
+        try:
+            with open(file_path, "rb") as f:
+                # Read the file in chunks to handle large files efficiently
+                for chunk in iter(lambda: f.read(chunk_size), b""):
+                    sha256_hash.update(chunk)
+            calculated = sha256_hash.hexdigest()
+            if calculated != checksum:
+                raise ToolkitError(f"Checksum mismatch. Expected {checksum}, got {calculated}.")
+            else:
+                print("Checksum verified")
+        except Exception as e:
+            raise ToolkitError(f"Failed to calculate checksum for {file_path}: {e}") from e
+
+    def _unpack(self, file_path: Path) -> None:
+        """
+        Unzips the downloaded file to the specified output path.
+        If the file is not a zip file, it raises an error.
+        """
+        total_size = file_path.stat().st_size if file_path.exists() else 0
+
+        try:
+            with Progress() as progress:
+                unzip = progress.add_task("Unzipping", total=total_size)
+                with zipfile.ZipFile(file_path, "r") as zip_ref:
+                    zip_ref.extractall(file_path.parent)
+                    progress.update(unzip, advance=total_size)
+        except zipfile.BadZipFile as e:
+            raise ToolkitError(f"Error unpacking zip file {file_path}: {e}") from e
+        except Exception as e:
+            raise ToolkitError(f"An unexpected error occurred while unpacking {file_path}: {e}") from e

cognite_toolkit/_cdf_tk/data_classes/_packages.py

@@ -66,8 +66,8 @@ class Packages(dict, MutableMapping[str, Package]):
             root_module_dir: The module directories to load the packages from.
         """
 
-        package_definition_path = root_module_dir / "package.toml"
-        if not package_definition_path.exists():
+        package_definition_path = next(root_module_dir.rglob("packages.toml"), None)
+        if not package_definition_path or not package_definition_path.exists():
             raise ToolkitFileNotFoundError(f"Package manifest toml not found at {package_definition_path}")
         package_definitions = toml.loads(package_definition_path.read_text(encoding="utf-8"))["packages"]
 

cognite_toolkit/_cdf_tk/feature_flags.py

@@ -52,6 +52,10 @@ class Flags(Enum):
         "visible": True,
         "description": "Enables the migrate command",
     }
+    EXTERNAL_LIBRARIES: ClassVar[dict[str, Any]] = {  # type: ignore[misc]
+        "visible": True,
+        "description": "Enables the support for external libraries in the config file",
+    }
 
     def is_enabled(self) -> bool:
         return FeatureFlag.is_enabled(self)

cognite_toolkit/_cdf_tk/resource_classes/capabilities.py

@@ -48,6 +48,11 @@ class AllScope(Scope):
     _scope_name = "all"
 
 
+class AppConfigScope(Scope):
+    _scope_name = "appScope"
+    apps: list[Literal["SEARCH"]]
+
+
 class CurrentUserScope(Scope):
     _scope_name = "currentuserscope"
 
@@ -175,6 +180,12 @@ class AnnotationsAcl(Capability):
     scope: AllScope
 
 
+class AppConfigAcl(Capability):
+    _capability_name = "appConfigAcl"
+    actions: list[Literal["READ", "WRITE"]]
+    scope: AllScope | AppConfigScope
+
+
 class AssetsAcl(Capability):
     _capability_name = "assetsAcl"
     actions: list[Literal["READ", "WRITE"]]