cognite-toolkit 0.5.44__py3-none-any.whl → 0.7.34__py3-none-any.whl

cognite_toolkit/_cdf_tk/utils/aggregators.py

@@ -0,0 +1,424 @@
+from abc import ABC, abstractmethod
+from functools import lru_cache
+from typing import ClassVar, Generic, Literal, TypeVar
+
+from cognite.client.data_classes import (
+    AssetFilter,
+    EventFilter,
+    FileMetadataFilter,
+    SequenceFilter,
+    TimeSeriesFilter,
+    Transformation,
+    filters,
+)
+from cognite.client.data_classes.assets import AssetProperty
+from cognite.client.data_classes.documents import SourceFileProperty
+from cognite.client.data_classes.events import EventProperty
+from cognite.client.data_classes.sequences import SequenceProperty
+from cognite.client.data_classes.time_series import TimeSeriesProperty
+
+from cognite_toolkit._cdf_tk.client import ToolkitClient
+from cognite_toolkit._cdf_tk.exceptions import ToolkitMissingResourceError
+from cognite_toolkit._cdf_tk.utils.cdf import (
+    label_aggregate_count,
+    label_count,
+    metadata_key_counts,
+    relationship_aggregate_count,
+)
+from cognite_toolkit._cdf_tk.utils.sql_parser import SQLParser
+
+T_CogniteFilter = TypeVar(
+    "T_CogniteFilter",
+    bound=AssetFilter | EventFilter | FileMetadataFilter | TimeSeriesFilter | SequenceFilter,
+    contravariant=True,
+)
+
+
+class AssetCentricAggregator(ABC):
+    _transformation_destination: ClassVar[tuple[str, ...]]
+
+    def __init__(self, client: ToolkitClient) -> None:
+        self.client = client
+
+    @property
+    @abstractmethod
+    def display_name(self) -> str:
+        raise NotImplementedError()
+
+    @abstractmethod
+    def count(
+        self, hierarchy: str | list[str] | None = None, data_set_external_id: str | list[str] | None = None
+    ) -> int:
+        raise NotImplementedError
+
+    def transformation_count(self) -> int:
+        """Returns the number of transformations associated with the resource."""
+        transformation_count = 0
+        for destination in self._transformation_destination:
+            for chunk in self.client.transformations(chunk_size=1000, destination_type=destination, limit=None):
+                transformation_count += len(chunk)
+        return transformation_count
+
+    @abstractmethod
+    def used_data_sets(self, hierarchy: str | None = None) -> list[str]:
+        """Returns a list of data sets used by the resource."""
+        raise NotImplementedError
+
+    def used_transformations(self, data_set_external_ids: list[str]) -> list[Transformation]:
+        """Returns a list of transformations used by the resource."""
+        data_set_ids = self.client.lookup.data_sets.id(data_set_external_ids, allow_empty=True)
+        found_transformations: list[Transformation] = []
+        for destination in self._transformation_destination:
+            for chunk in self.client.transformations(chunk_size=1000, destination_type=destination, limit=None):
+                for transformation in chunk:
+                    if SQLParser(transformation.query or "", operation="profiling").is_using_data_set(
+                        data_set_ids, data_set_external_ids
+                    ):
+                        found_transformations.append(transformation)
+        return found_transformations
+
+    @staticmethod
+    def _to_unique_int_list(results: list) -> list[int]:
+        """Converts a list of results to a unique list of integers.
+
+        This method does the following:
+        * Converts each item in the results to an integer, if possible.
+        * Filters out None.
+        * Removes duplicates
+
+        This is used as the aggregation results are inconsistently implemented for the different resources,
+        when aggregating dataSetIds, Sequences, TimeSeries, and Files return a list of strings, while
+        Assets and Events return a list of integers. In addition, the files aggregation can return
+        duplicated values.
+        """
+        seen: set[int] = set()
+        ids: list[int] = []
+        for id_ in results:
+            if isinstance(id_, int) and id_ not in seen:
+                ids.append(id_)
+                seen.add(id_)
+            try:
+                int_id = int(id_)
+            except (ValueError, TypeError):
+                continue
+            if int_id not in seen:
+                ids.append(int_id)
+                seen.add(int_id)
+        return ids
+
+    def _to_dataset_id(self, data_set_external_id: str | list[str] | None) -> list[int] | None:
+        """Converts data set external IDs to data set IDs."""
+        dataset_id: list[int] | None = None
+        if data_set_external_id is not None:
+            if isinstance(data_set_external_id, str):
+                data_set_external_id = [data_set_external_id]
+            dataset_id = self.client.lookup.data_sets.id(data_set_external_id, allow_empty=False)
+        return dataset_id
+
+
+class MetadataAggregator(AssetCentricAggregator, ABC, Generic[T_CogniteFilter]):
+    filter_cls: type[T_CogniteFilter]
+
+    def __init__(
+        self, client: ToolkitClient, resource_name: Literal["assets", "events", "files", "timeseries", "sequences"]
+    ) -> None:
+        super().__init__(client)
+        self.resource_name = resource_name
+
+    def metadata_key_count(
+        self, hierarchy: str | list[str] | None = None, data_sets: str | list[str] | None = None
+    ) -> int:
+        """Returns the number of metadata keys used by the resource."""
+        return len(self.used_metadata_keys(hierarchy=hierarchy, data_sets=data_sets))
+
+    def used_metadata_keys(
+        self, hierarchy: str | list[str] | None = None, data_sets: str | list[str] | None = None
+    ) -> list[tuple[str, int]]:
+        """Returns a list of metadata keys and their counts."""
+        hierarchy_ids, data_set_ids = self._lookup_hierarchy_data_set_pair(
+            hierarchy, data_sets, operation="find metadata keys"
+        )
+        return self._used_metadata_keys(hierarchy=hierarchy_ids, data_sets=data_set_ids)
+
+    @lru_cache(maxsize=1)
+    def _used_metadata_keys(
+        self, hierarchy: tuple[int, ...] | None = None, data_sets: tuple[int, ...] | None = None
+    ) -> list[tuple[str, int]]:
+        return metadata_key_counts(
+            self.client,
+            self.resource_name,
+            hierarchies=list(hierarchy) if hierarchy else None,
+            data_sets=list(data_sets) if data_sets else None,
+        )
+
+    def _lookup_hierarchy_data_set_pair(
+        self, hierarchy: str | list[str] | None, data_sets: str | list[str] | None, operation: str
+    ) -> tuple[tuple[int, ...] | None, tuple[int, ...] | None]:
+        """Returns a tuple of hierarchy and data sets."""
+        hierarchy_ids: tuple[int, ...] | None = None
+        if isinstance(hierarchy, str):
+            asset_id = self.client.lookup.assets.id(external_id=hierarchy, allow_empty=False)
+            if asset_id is None:
+                raise ToolkitMissingResourceError(f"Cannot {operation}. Asset with external ID {hierarchy!r} not found")
+            hierarchy_ids = (asset_id,)
+        elif isinstance(hierarchy, list) and all(isinstance(item, str) for item in hierarchy):
+            asset_ids = self.client.lookup.assets.id(external_id=hierarchy, allow_empty=False)
+            if len(asset_ids) != len(hierarchy):
+                missing = set(hierarchy) - set(
+                    self.client.lookup.assets.external_id([id_ for id_ in asset_ids if id_ is not None])
+                )
+                raise ToolkitMissingResourceError(
+                    f"Cannot {operation}. Assets with external IDs {sorted(missing)!r} not found"
+                )
+            hierarchy_ids = tuple(sorted(asset_ids))
+
+        data_set_ids: tuple[int, ...] | None = None
+        if isinstance(data_sets, str):
+            data_set_id = self.client.lookup.data_sets.id(external_id=data_sets, allow_empty=False)
+            if data_set_id is None:
+                raise ToolkitMissingResourceError(
+                    f"Cannot {operation}. Data set with external ID {data_sets!r} not found"
+                )
+            data_set_ids = (data_set_id,)
+        elif isinstance(data_sets, list) and all(isinstance(item, str) for item in data_sets):
+            data_set_ids_list = self.client.lookup.data_sets.id(external_id=data_sets, allow_empty=False)
+            if len(data_set_ids_list) != len(data_sets):
+                missing = set(data_sets) - set(
+                    self.client.lookup.data_sets.external_id([id_ for id_ in data_set_ids_list if id_ is not None])
+                )
+                raise ToolkitMissingResourceError(
+                    f"Cannot {operation}. Data sets with external IDs {sorted(missing)!r} not found"
+                )
+            data_set_ids = tuple(sorted(data_set_ids_list))
+
+        return hierarchy_ids, data_set_ids
+
+    @classmethod
+    def create_filter(
+        cls,
+        hierarchy: str | list[str] | tuple[str, ...] | None = None,
+        data_set_external_id: str | list[str] | tuple[str, ...] | None = None,
+    ) -> T_CogniteFilter | None:
+        """Creates a filter for the resource based on hierarchy and data set external ID."""
+        if cls._is_empty(hierarchy) and cls._is_empty(data_set_external_id):
+            return None
+        asset_subtree_ids: list[dict[str, str]] | None = None
+        if isinstance(hierarchy, str):
+            asset_subtree_ids = [{"externalId": hierarchy}]
+        elif isinstance(hierarchy, list | tuple) and hierarchy:
+            asset_subtree_ids = [{"externalId": item} for item in hierarchy]
+        data_set_ids: list[dict[str, str]] | None = None
+        if isinstance(data_set_external_id, str):
+            data_set_ids = [{"externalId": data_set_external_id}]
+        elif isinstance(data_set_external_id, list | tuple) and data_set_external_id:
+            data_set_ids = [{"externalId": item} for item in data_set_external_id]
+
+        # MyPy fails to understand that filter_cls() produce a T_CogniteFilter
+        return cls.filter_cls(asset_subtree_ids=asset_subtree_ids, data_set_ids=data_set_ids)  # type: ignore[return-value]
+
+    @classmethod
+    def _is_empty(cls, items: str | list[str] | tuple[str, ...] | None) -> bool:
+        """Checks if the provided items are empty."""
+        if items is None:
+            return True
+        if isinstance(items, list | tuple):
+            return not items
+        return False
+
+
+class LabelAggregator(MetadataAggregator, ABC, Generic[T_CogniteFilter]):
+    def label_count(self, hierarchy: str | list[str] | None = None, data_sets: str | list[str] | None = None) -> int:
+        """Returns the number of labels used by the resource."""
+        return len(self.used_labels(hierarchy=hierarchy, data_sets=data_sets))
+
+    def used_labels(
+        self, hierarchy: str | list[str] | None = None, data_sets: str | list[str] | None = None
+    ) -> list[tuple[str, int]]:
+        """Returns a list of labels and their counts."""
+        hierarchy_ids, data_set_ids = self._lookup_hierarchy_data_set_pair(
+            hierarchy, data_sets, operation="find labels"
+        )
+        return self._used_labels(hierarchy=hierarchy_ids, data_sets=data_set_ids)
+
+    @lru_cache(maxsize=1)
+    def _used_labels(
+        self, hierarchy: tuple[int, ...] | None = None, data_sets: tuple[int, ...] | None = None
+    ) -> list[tuple[str, int]]:
+        """Returns a list of labels and their counts."""
+        return label_count(
+            self.client,
+            self.resource_name,
+            hierarchies=list(hierarchy) if hierarchy else None,
+            data_sets=list(data_sets) if data_sets else None,
+        )
+
+
+class AssetAggregator(LabelAggregator[AssetFilter]):
+    _transformation_destination = ("assets", "asset_hierarchy")
+    filter_cls = AssetFilter
+
+    def __init__(self, client: ToolkitClient) -> None:
+        super().__init__(client, "assets")
+
+    @property
+    def display_name(self) -> str:
+        return "Assets"
+
+    def count(
+        self,
+        hierarchy: str | list[str] | tuple[str, ...] | None = None,
+        data_set_external_id: str | list[str] | tuple[str, ...] | None = None,
+    ) -> int:
+        return self.client.assets.aggregate_count(filter=self.create_filter(hierarchy, data_set_external_id))
+
+    def used_data_sets(self, hierarchy: str | None = None) -> list[str]:
+        """Returns a list of data sets used by the resource."""
+        results = self.client.assets.aggregate_unique_values(
+            AssetProperty.data_set_id, filter=self.create_filter(hierarchy)
+        )
+        return self.client.lookup.data_sets.external_id([id_ for id_ in results.unique if isinstance(id_, int)])
+
+
+class EventAggregator(MetadataAggregator[EventFilter]):
+    _transformation_destination = ("events",)
+    filter_cls = EventFilter
+
+    def __init__(self, client: ToolkitClient) -> None:
+        super().__init__(client, "events")
+
+    @property
+    def display_name(self) -> str:
+        return "Events"
+
+    def count(
+        self, hierarchy: str | list[str] | None = None, data_set_external_id: str | list[str] | None = None
+    ) -> int:
+        return self.client.events.aggregate_count(filter=self.create_filter(hierarchy, data_set_external_id))
+
+    def used_data_sets(self, hierarchy: str | None = None) -> list[str]:
+        """Returns a list of data sets used by the resource."""
+        results = self.client.events.aggregate_unique_values(
+            property=EventProperty.data_set_id, filter=self.create_filter(hierarchy)
+        )
+        return self.client.lookup.data_sets.external_id([id_ for id_ in results.unique if isinstance(id_, int)])
+
+
+class FileAggregator(LabelAggregator[FileMetadataFilter]):
+    _transformation_destination = ("files",)
+    filter_cls = FileMetadataFilter
+
+    def __init__(self, client: ToolkitClient) -> None:
+        super().__init__(client, "files")
+
+    @property
+    def display_name(self) -> str:
+        return "Files"
+
+    def count(
+        self, hierarchy: str | list[str] | None = None, data_set_external_id: str | list[str] | None = None
+    ) -> int:
+        response = self.client.files.aggregate(filter=self.create_filter(hierarchy, data_set_external_id))
+        if response:
+            return response[0].count
+        else:
+            return 0
+
+    def used_data_sets(self, hierarchy: str | None = None) -> list[str]:
+        """Returns a list of data sets used by the resource."""
+        filter_: filters.Filter | None = None
+        if hierarchy is not None:
+            filter_ = filters.InAssetSubtree("assetExternalIds", [hierarchy])
+        results = self.client.documents.aggregate_unique_values(
+            property=SourceFileProperty.data_set_id, filter=filter_, limit=1000
+        )
+        ids = self._to_unique_int_list(results.unique)
+        return self.client.lookup.data_sets.external_id(list(ids))
+
+
+class TimeSeriesAggregator(MetadataAggregator[TimeSeriesFilter]):
+    _transformation_destination = ("timeseries",)
+    filter_cls = TimeSeriesFilter
+
+    def __init__(self, client: ToolkitClient) -> None:
+        super().__init__(client, "timeseries")
+
+    @property
+    def display_name(self) -> str:
+        return "TimeSeries"
+
+    def count(
+        self, hierarchy: str | list[str] | None = None, data_set_external_id: str | list[str] | None = None
+    ) -> int:
+        return self.client.time_series.aggregate_count(filter=self.create_filter(hierarchy, data_set_external_id))
+
+    def used_data_sets(self, hierarchy: str | None = None) -> list[str]:
+        """Returns a list of data sets used by the resource."""
+        results = self.client.time_series.aggregate_unique_values(
+            property=TimeSeriesProperty.data_set_id, filter=self.create_filter(hierarchy)
+        )
+        ids = self._to_unique_int_list(results.unique)
+        return self.client.lookup.data_sets.external_id(ids)
+
+
+class SequenceAggregator(MetadataAggregator):
+    _transformation_destination = ("sequences",)
+    filter_cls = SequenceFilter
+
+    def __init__(self, client: ToolkitClient) -> None:
+        super().__init__(client, "sequences")
+
+    @property
+    def display_name(self) -> str:
+        return "Sequences"
+
+    def count(
+        self, hierarchy: str | list[str] | None = None, data_set_external_id: str | list[str] | None = None
+    ) -> int:
+        return self.client.sequences.aggregate_count(filter=self.create_filter(hierarchy, data_set_external_id))
+
+    def used_data_sets(self, hierarchy: str | None = None) -> list[str]:
+        """Returns a list of data sets used by the resource."""
+        results = self.client.sequences.aggregate_unique_values(
+            property=SequenceProperty.data_set_id, filter=self.create_filter(hierarchy)
+        )
+        ids = self._to_unique_int_list(results.unique)
+        return self.client.lookup.data_sets.external_id(ids)
+
+
+class RelationshipAggregator(AssetCentricAggregator):
+    _transformation_destination = ("relationships",)
+
+    @property
+    def display_name(self) -> str:
+        return "Relationships"
+
+    def count(
+        self, hierarchy: str | list[str] | None = None, data_set_external_id: str | list[str] | None = None
+    ) -> int:
+        if hierarchy is not None:
+            raise NotImplementedError()
+        dataset_id = self._to_dataset_id(data_set_external_id)
+        results = relationship_aggregate_count(self.client, dataset_id)
+        return sum(result.count for result in results)
+
+    def used_data_sets(self, hierarchy: str | None = None) -> list[str]:
+        raise NotImplementedError()
+
+
+class LabelCountAggregator(AssetCentricAggregator):
+    _transformation_destination = ("labels",)
+
+    @property
+    def display_name(self) -> str:
+        return "Labels"
+
+    def count(
+        self, hierarchy: str | list[str] | None = None, data_set_external_id: str | list[str] | None = None
+    ) -> int:
+        if hierarchy is not None:
+            raise NotImplementedError()
+        data_set_id = self._to_dataset_id(data_set_external_id)
+        return label_aggregate_count(self.client, data_set_id)
+
+    def used_data_sets(self, hierarchy: str | None = None) -> list[str]:
+        raise NotImplementedError()

cognite_toolkit/_cdf_tk/utils/auth.py

@@ -337,9 +337,11 @@ class EnvironmentVariables:
             max_workers=self.CDF_CLIENT_MAX_WORKERS,
         )
 
-    def get_client(self, is_strict_validation: bool = True) -> ToolkitClient:
+    def get_client(self, is_strict_validation: bool = True, enable_set_pending_ids: bool = False) -> ToolkitClient:
         if self._client is None:
-            self._client = ToolkitClient(config=self.get_config(is_strict_validation))
+            self._client = ToolkitClient(
+                config=self.get_config(is_strict_validation), enable_set_pending_ids=enable_set_pending_ids
+            )
         return self._client
 
     def dump(self, include_os: bool = True) -> dict[str, str | None]:

cognite_toolkit/_cdf_tk/utils/auxiliary.py

@@ -0,0 +1,23 @@
+import functools
+import platform
+
+from cognite_toolkit import _version
+
+
+def get_current_toolkit_version() -> str:
+    return _version.__version__
+
+
+@functools.lru_cache(maxsize=1)
+def get_user_agent() -> str:
+    toolkit_version = f"CogniteToolkit/{get_current_toolkit_version()}"
+    python_version = (
+        f"{platform.python_implementation()}/{platform.python_version()} "
+        f"({platform.python_build()};{platform.python_compiler()})"
+    )
+    os_version_info = [platform.release(), platform.machine(), platform.architecture()[0]]
+    os_version_info = [s for s in os_version_info if s]  # Ignore empty strings
+    os_version_info_str = "-".join(os_version_info)
+    operating_system = f"{platform.system()}/{os_version_info_str}"
+
+    return f"{toolkit_version} {python_version} {operating_system}"