cognite-toolkit 0.5.61__py3-none-any.whl → 0.5.63__py3-none-any.whl

cognite_toolkit/_cdf_tk/client/data_classes/search_config.py

@@ -0,0 +1,198 @@
+from __future__ import annotations
+
+from abc import ABC
+from dataclasses import dataclass
+from typing import Any
+
+from cognite.client import CogniteClient
+from cognite.client.data_classes._base import (
+    CogniteObject,
+    CogniteResourceList,
+    WriteableCogniteResource,
+    WriteableCogniteResourceList,
+)
+from typing_extensions import Self
+
+
+@dataclass
+class SearchConfigView(CogniteObject):
+    external_id: str
+    space: str
+
+    @classmethod
+    def _load(cls, resource: dict[str, Any], cognite_client: CogniteClient | None = None) -> Self:
+        return cls(
+            external_id=resource["externalId"],
+            space=resource["space"],
+        )
+
+
+@dataclass
+class SearchConfigViewProperty(CogniteObject):
+    property: str
+    disabled: bool | None = None
+    selected: bool | None = None
+    hidden: bool | None = None
+
+    @classmethod
+    def _load(cls, resource: dict[str, Any], cognite_client: CogniteClient | None = None) -> Self:
+        return cls(
+            property=resource["property"],
+            disabled=resource.get("disabled"),
+            selected=resource.get("selected"),
+            hidden=resource.get("hidden"),
+        )
+
+
+class SearchConfigCore(WriteableCogniteResource["SearchConfigWrite"], ABC):
+    """
+    Core model for a single Configuration.
+
+    Args:
+        view: The configuration for one specific view.
+        id: A server-generated ID for the object.
+        use_as_name: The name of property to use for the name column in the UI.
+        use_as_description: The name of property to use for the description column in the UI.
+        column_layout: Array of column configurations per property.
+        filter_layout: Array of filter configurations per property.
+        properties_layout: Array of property configurations per property.
+    """
+
+    def __init__(
+        self,
+        view: SearchConfigView,
+        id: int | None = None,
+        use_as_name: str | None = None,
+        use_as_description: str | None = None,
+        column_layout: list[SearchConfigViewProperty] | None = None,
+        filter_layout: list[SearchConfigViewProperty] | None = None,
+        properties_layout: list[SearchConfigViewProperty] | None = None,
+    ) -> None:
+        self.view = view
+        self.id = id
+        self.use_as_name = use_as_name
+        self.use_as_description = use_as_description
+        self.column_layout = column_layout
+        self.filter_layout = filter_layout
+        self.properties_layout = properties_layout
+
+    def as_write(self) -> SearchConfigWrite:
+        return SearchConfigWrite(
+            view=self.view,
+            id=self.id,
+            use_as_name=self.use_as_name,
+            use_as_description=self.use_as_description,
+            column_layout=self.column_layout,
+            filter_layout=self.filter_layout,
+            properties_layout=self.properties_layout,
+        )
+
+    def dump(self, camel_case: bool = True) -> dict[str, Any]:
+        output = super().dump(camel_case)
+        if self.column_layout:
+            output["columLayout" if camel_case else "column_layout"] = [
+                _data.dump(camel_case) for _data in self.column_layout
+            ]
+        if self.filter_layout:
+            output["filterLayout" if camel_case else "filter_layout"] = [
+                _data.dump(camel_case) for _data in self.filter_layout
+            ]
+        if self.properties_layout:
+            output["propertiesLayout" if camel_case else "properties_layout"] = [
+                _data.dump(camel_case) for _data in self.properties_layout
+            ]
+        if self.view:
+            output["view"] = self.view.dump(camel_case)
+        return output
+
+
+class SearchConfigWrite(SearchConfigCore):
+    @classmethod
+    def _load(cls, resource: dict[str, Any], cognite_client: CogniteClient | None = None) -> Self:
+        return cls(
+            id=resource.get("id"),
+            view=SearchConfigView.load(resource["view"]),
+            use_as_name=resource.get("useAsName"),
+            use_as_description=resource.get("useAsDescription"),
+            column_layout=[SearchConfigViewProperty.load(item) for item in resource.get("columnLayout", [])]
+            if resource.get("columnLayout")
+            else None,
+            filter_layout=[SearchConfigViewProperty.load(item) for item in resource.get("filterLayout", [])]
+            if resource.get("filterLayout")
+            else None,
+            properties_layout=[SearchConfigViewProperty.load(item) for item in resource.get("propertiesLayout", [])]
+            if resource.get("propertiesLayout")
+            else None,
+        )
+
+
+class SearchConfig(SearchConfigCore):
+    """
+    Response model for a single Configuration.
+
+    Args:
+        view: The configuration for one specific view.
+        id: A server-generated ID for the object.
+        created_time: The time when the search config was created.
+        updated_time: The time when the search config was last updated.
+        use_as_name: The name of property to use for the name column in the UI.
+        use_as_description: The name of property to use for the description column in the UI.
+        column_layout: Array of column configurations per property.
+        filter_layout: Array of filter configurations per property.
+        properties_layout: Array of property configurations per property.
+    """
+
+    def __init__(
+        self,
+        view: SearchConfigView,
+        id: int,
+        created_time: int,
+        updated_time: int,
+        use_as_name: str | None = None,
+        use_as_description: str | None = None,
+        column_layout: list[SearchConfigViewProperty] | None = None,
+        filter_layout: list[SearchConfigViewProperty] | None = None,
+        properties_layout: list[SearchConfigViewProperty] | None = None,
+    ) -> None:
+        super().__init__(
+            view,
+            id,
+            use_as_name,
+            use_as_description,
+            column_layout,
+            filter_layout,
+            properties_layout,
+        )
+        self.created_time = created_time
+        self.updated_time = updated_time
+
+    @classmethod
+    def _load(cls, resource: dict[str, Any], cognite_client: CogniteClient | None = None) -> Self:
+        return cls(
+            view=SearchConfigView.load(resource["view"]),
+            id=resource["id"],
+            created_time=resource["createdTime"],
+            updated_time=resource["lastUpdatedTime"],
+            use_as_name=resource.get("useAsName"),
+            use_as_description=resource.get("useAsDescription"),
+            column_layout=[SearchConfigViewProperty.load(item) for item in resource.get("columnLayout", [])]
+            if resource.get("columnLayout")
+            else None,
+            filter_layout=[SearchConfigViewProperty.load(item) for item in resource.get("filterLayout", [])]
+            if resource.get("filterLayout")
+            else None,
+            properties_layout=[SearchConfigViewProperty.load(item) for item in resource.get("propertiesLayout", [])]
+            if resource.get("propertiesLayout")
+            else None,
+        )
+
+
+class SearchConfigWriteList(CogniteResourceList):
+    _RESOURCE = SearchConfigWrite
+
+
+class SearchConfigList(WriteableCogniteResourceList[SearchConfigWrite, SearchConfig]):
+    _RESOURCE = SearchConfig
+
+    def as_write(self) -> SearchConfigWriteList:
+        return SearchConfigWriteList([searchConfig.as_write() for searchConfig in self])

cognite_toolkit/_cdf_tk/commands/__init__.py

@@ -1,6 +1,6 @@
 from ._migrate import MigrateTimeseriesCommand, MigrationPrepareCommand
 from ._populate import PopulateCommand
-from ._profile import ProfileCommand
+from ._profile import ProfileAssetCentricCommand, ProfileTransformationCommand
 from ._purge import PurgeCommand
 from .auth import AuthCommand
 from .build_cmd import BuildCommand
@@ -30,7 +30,8 @@ __all__ = [
     "MigrationPrepareCommand",
     "ModulesCommand",
     "PopulateCommand",
-    "ProfileCommand",
+    "ProfileAssetCentricCommand",
+    "ProfileTransformationCommand",
     "PullCommand",
     "PurgeCommand",
     "RepoCommand",

cognite_toolkit/_cdf_tk/commands/_profile.py

@@ -1,12 +1,21 @@
-from collections.abc import Callable
+import itertools
+from abc import ABC, abstractmethod
+from collections.abc import Callable, Iterable, Mapping
 from concurrent.futures import ThreadPoolExecutor, as_completed
+from functools import cached_property
+from typing import ClassVar, Literal, TypeAlias, overload
 
+from cognite.client.data_classes import Transformation
 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
 
 from cognite_toolkit._cdf_tk.client import ToolkitClient
+from cognite_toolkit._cdf_tk.exceptions import ToolkitValueError
+from cognite_toolkit._cdf_tk.utils import humanize_collection
 from cognite_toolkit._cdf_tk.utils.aggregators import (
     AssetAggregator,
     AssetCentricAggregator,
@@ -19,112 +28,285 @@ from cognite_toolkit._cdf_tk.utils.aggregators import (
     SequenceAggregator,
     TimeSeriesAggregator,
 )
+from cognite_toolkit._cdf_tk.utils.sql_parser import SQLParser, SQLTable
 
 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
 
-    @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
+    @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]]: ...
+
+    @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
+
+    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}")
+
+
+class ProfileTransformationCommand(ProfileCommand):
+    valid_destinations: frozenset[str] = frozenset({"assets", "files", "events", "timeseries", "sequences"})
+
+    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 = "Transformation Profile"
+        self.destination_type: Literal["assets", "files", "events", "timeseries", "sequences"] | None = None
+
+    class Columns:
+        Transformation = "Transformation"
+        Source = "Sources"
+        DestinationColumns = "Destination Columns"
+        Destination = "Destination"
+        ConflictMode = "Conflict Mode"
+        IsPaused = "Is Paused"
+
+    def transformation(
+        self, client: ToolkitClient, destination_type: str | None = None, verbose: bool = False
+    ) -> list[dict[str, CellValue]]:
+        self.destination_type = self._validate_destination_type(destination_type)
+        return self.create_profile_table(client)
+
+    @classmethod
+    def _validate_destination_type(
+        cls, destination_type: str | None
+    ) -> Literal["assets", "files", "events", "timeseries", "sequences"]:
+        if destination_type is None or destination_type not in cls.valid_destinations:
+            raise ToolkitValueError(
+                f"Invalid destination type: {destination_type}. Must be one of {humanize_collection(cls.valid_destinations)}."
+            )
+        # We validated the destination type above
+        return destination_type  # type: ignore[return-value]
+
+    def create_initial_table(self, client: ToolkitClient) -> dict[tuple[str, str], PendingCellValue]:
+        if self.valid_destinations is None:
+            raise ToolkitValueError("Destination type must be set before calling create_initial_table.")
+        iterable: Iterable[Transformation] = client.transformations.list(
+            destination_type=self.destination_type, limit=-1
+        )
+        if self.destination_type == "assets":
+            iterable = itertools.chain(iterable, client.transformations(destination_type="asset_hierarchy", limit=-1))
+        table: dict[tuple[str, str], PendingCellValue] = {}
+        for transformation in iterable:
+            sources: list[SQLTable] = []
+            destination_columns: list[str] = []
+            if transformation.query:
+                parser = SQLParser(transformation.query, operation="Profile transformations")
+                sources = parser.sources
+                destination_columns = parser.destination_columns
+            index = str(transformation.id)
+            table[(index, self.Columns.Transformation)] = transformation.name or transformation.external_id or "Unknown"
+            table[(index, self.Columns.Source)] = ", ".join(map(str, sources))
+            table[(index, self.Columns.DestinationColumns)] = (
+                ", ".join(destination_columns) or None if destination_columns else None
+            )
+            table[(index, self.Columns.Destination)] = (
+                transformation.destination.type or "Unknown" if transformation.destination else "Unknown"
+            )
+            table[(index, self.Columns.ConflictMode)] = transformation.conflict_mode or "Unknown"
+            table[(index, self.Columns.IsPaused)] = (
+                str(transformation.schedule.is_paused) if transformation.schedule else "No schedule"
+            )
+        return table
 
-        return styled_callable
+    def call_api(self, row: str, col: str, client: ToolkitClient) -> Callable:
+        raise NotImplementedError(f"{type(self).__name__} does not support API calls for {col} in row {row}.")