cognite-toolkit 0.6.83__py3-none-any.whl → 0.6.85__py3-none-any.whl

cognite_toolkit/_cdf_tk/cruds/_resource_cruds/function.py

@@ -36,7 +36,6 @@ from cognite_toolkit._cdf_tk.client.data_classes.functions import FunctionSchedu
 from cognite_toolkit._cdf_tk.cruds._base_cruds import ResourceCRUD
 from cognite_toolkit._cdf_tk.exceptions import (
     ResourceCreationError,
-    ResourceRetrievalError,
     ToolkitRequiredValueError,
 )
 from cognite_toolkit._cdf_tk.feature_flags import Flags
@@ -565,22 +564,7 @@ class FunctionScheduleCRUD(
 
     def create(self, items: FunctionScheduleWriteList) -> FunctionSchedulesList:
         created_list = FunctionSchedulesList([], cognite_client=self.client)
-        functions_to_lookup = list({item.function_external_id for item in items if item.function_external_id})
-        function_id_by_external_id: dict[str, int] = {}
-        if functions_to_lookup:
-            try:
-                function_ids = self.client.lookup.functions.id(functions_to_lookup)
-            except ResourceRetrievalError as e:
-                failed_items = self.get_ids(items)
-                missing_functions = functions_to_lookup
-                if e.resources:
-                    missing_functions = list(e.resources)
-                    failed_items = [id_ for id_ in failed_items if id_.function_external_id in set(missing_functions)]
-                raise ResourceCreationError(
-                    f"Failed to create function schedule(s) {humanize_collection(failed_items)}. "
-                    f"Could not find function(s) {humanize_collection(missing_functions)!r}"
-                ) from e
-            function_id_by_external_id = dict(zip(functions_to_lookup, function_ids))
+        function_id_by_external_id = self._get_function_ids_by_external_id(items)
 
         for item in items:
             id_ = self.get_id(item)
@@ -607,6 +591,28 @@ class FunctionScheduleCRUD(
             created_list.append(created)
         return created_list
 
+    def _get_function_ids_by_external_id(self, items: FunctionScheduleWriteList) -> dict[str, int]:
+        functions_to_lookup = list({item.function_external_id for item in items if item.function_external_id})
+        if not functions_to_lookup:
+            return {}
+        function_ids = self.client.lookup.functions.id(functions_to_lookup)
+
+        if len(function_ids) == len(functions_to_lookup):
+            return dict(zip(functions_to_lookup, function_ids))
+        # The lookup API is cached, so it is cheap to do individual lookups to find the missing ones.
+        lookup_pair = (
+            (function_external_id, self.client.lookup.functions.id(function_external_id))
+            for function_external_id in functions_to_lookup
+        )
+        missing_functions = {func for func, func_id in lookup_pair if func_id is None}
+        failed_schedules = [self.get_id(item) for item in items if item.function_external_id in missing_functions]
+        plural_schedules = "s" if len(failed_schedules) > 1 else ""
+        plural_fun = "s" if len(missing_functions) > 1 else ""
+        raise ResourceCreationError(
+            f"Failed to create function schedule{plural_schedules} {humanize_collection(failed_schedules)}. "
+            f"Could not find function{plural_fun} {humanize_collection(missing_functions)!r}"
+        )
+
     def delete(self, ids: SequenceNotStr[FunctionScheduleID]) -> int:
         schedules = self.retrieve(ids)
         count = 0

cognite_toolkit/_cdf_tk/storageio/_asset_centric.py

@@ -36,10 +36,9 @@ from cognite_toolkit._cdf_tk.cruds import (
     EventCRUD,
     FileMetadataCRUD,
     LabelCRUD,
-    ResourceCRUD,
     TimeSeriesCRUD,
 )
-from cognite_toolkit._cdf_tk.exceptions import ToolkitNotImplementedError
+from cognite_toolkit._cdf_tk.exceptions import ToolkitMissingResourceError, ToolkitNotImplementedError
 from cognite_toolkit._cdf_tk.utils.aggregators import (
     AssetAggregator,
     AssetCentricAggregator,
@@ -85,19 +84,10 @@ class BaseAssetCentricIO(
 
     def __init__(self, client: ToolkitClient) -> None:
         super().__init__(client)
-        self._loader = self._get_loader()
         self._aggregator = self._get_aggregator()
         self._downloaded_data_sets_by_selector: dict[AssetCentricSelector, set[int]] = defaultdict(set)
         self._downloaded_labels_by_selector: dict[AssetCentricSelector, set[str]] = defaultdict(set)
 
-    @abstractmethod
-    def _get_loader(
-        self,
-    ) -> ResourceCRUD[
-        T_ID, T_WriteClass, T_WritableCogniteResource, T_CogniteResourceList, T_WritableCogniteResourceList
-    ]:
-        raise NotImplementedError()
-
     @abstractmethod
     def _get_aggregator(self) -> AssetCentricAggregator:
         raise NotImplementedError()
@@ -113,11 +103,6 @@ class BaseAssetCentricIO(
             return self._aggregator.count(hierarchy=selector.hierarchy)
         return None
 
-    def data_to_json_chunk(
-        self, data_chunk: Sequence[T_WritableCogniteResource], selector: AssetCentricSelector | None = None
-    ) -> list[dict[str, JsonVal]]:
-        return [self._loader.dump_resource(item) for item in data_chunk]
-
     def configurations(self, selector: AssetCentricSelector) -> Iterable[StorageIOConfig]:
         data_set_ids = self._downloaded_data_sets_by_selector[selector]
         if data_set_ids:
@@ -178,6 +163,22 @@ class BaseAssetCentricIO(
     def create_internal_identifier(cls, internal_id: int, project: str) -> str:
         return f"INTERNAL_ID_project_{project}_{internal_id!s}"
 
+    def _populate_data_set_cache(self, chunk: Sequence[Asset | FileMetadata | TimeSeries | Event]) -> None:
+        data_set_ids = {item.data_set_id for item in chunk if item.data_set_id is not None}
+        self.client.lookup.data_sets.external_id(list(data_set_ids))
+
+    def _populate_security_category_cache(self, chunk: Sequence[FileMetadata | TimeSeries]) -> None:
+        security_category_ids: set[int] = set()
+        for item in chunk:
+            security_category_ids.update(item.security_categories or [])
+        self.client.lookup.security_categories.external_id(list(security_category_ids))
+
+    def _populate_asset_cache(self, chunk: Sequence[FileMetadata | Event]) -> None:
+        asset_ids: set[int] = set()
+        for item in chunk:
+            asset_ids.update(item.asset_ids or [])
+        self.client.lookup.assets.external_id(list(asset_ids))
+
 
 class AssetIO(BaseAssetCentricIO[str, AssetWrite, Asset, AssetWriteList, AssetList]):
     KIND = "Assets"
@@ -187,22 +188,31 @@ class AssetIO(BaseAssetCentricIO[str, AssetWrite, Asset, AssetWriteList, AssetLi
     SUPPORTED_READ_FORMATS = frozenset({".parquet", ".csv", ".ndjson", ".yaml", ".yml"})
     UPLOAD_ENDPOINT = "/assets"
 
+    def __init__(self, client: ToolkitClient) -> None:
+        super().__init__(client)
+        self._crud = AssetCRUD.create_loader(self.client)
+
     def as_id(self, item: Asset) -> str:
         return item.external_id if item.external_id is not None else self._create_identifier(item.id)
 
-    def _get_loader(self) -> AssetCRUD:
-        return AssetCRUD.create_loader(self.client)
-
     def _get_aggregator(self) -> AssetCentricAggregator:
         return AssetAggregator(self.client)
 
     def get_schema(self, selector: AssetCentricSelector) -> list[SchemaColumn]:
         data_set_ids: list[int] = []
         if isinstance(selector, DataSetSelector):
-            data_set_ids.append(self.client.lookup.data_sets.id(selector.data_set_external_id))
+            data_set_id = self.client.lookup.data_sets.id(selector.data_set_external_id)
+            if data_set_id is None:
+                raise ToolkitMissingResourceError(
+                    f"Data set with external ID {selector.data_set_external_id} not found."
+                )
+            data_set_ids.append(data_set_id)
         hierarchy: list[int] = []
         if isinstance(selector, AssetSubtreeSelector):
-            hierarchy.append(self.client.lookup.assets.id(selector.hierarchy))
+            asset_id = self.client.lookup.assets.id(selector.hierarchy)
+            if asset_id is None:
+                raise ToolkitMissingResourceError(f"Asset with external ID {selector.hierarchy} not found.")
+            hierarchy.append(asset_id)
 
         if hierarchy or data_set_ids:
             metadata_keys = metadata_key_counts(
@@ -238,8 +248,17 @@ class AssetIO(BaseAssetCentricIO[str, AssetWrite, Asset, AssetWriteList, AssetLi
             self._collect_dependencies(asset_list, selector)
             yield Page(worker_id="main", items=asset_list)
 
+    def data_to_json_chunk(
+        self, data_chunk: Sequence[Asset], selector: AssetCentricSelector | None = None
+    ) -> list[dict[str, JsonVal]]:
+        # Ensure data sets are looked up to populate cache.
+        # This is to avoid looking up each data set id individually in the .dump_resource call.
+        self._populate_data_set_cache(data_chunk)
+
+        return [self._crud.dump_resource(item) for item in data_chunk]
+
     def json_to_resource(self, item_json: dict[str, JsonVal]) -> AssetWrite:
-        return self._loader.load_resource(item_json)
+        return self._crud.load_resource(item_json)
 
     def retrieve(self, ids: Sequence[int]) -> AssetList:
         return self.client.assets.retrieve_multiple(ids)
@@ -253,19 +272,25 @@ class FileMetadataIO(BaseAssetCentricIO[str, FileMetadataWrite, FileMetadata, Fi
     SUPPORTED_READ_FORMATS = frozenset({".parquet", ".csv", ".ndjson"})
     UPLOAD_ENDPOINT = "/files"
 
+    def __init__(self, client: ToolkitClient) -> None:
+        super().__init__(client)
+        self._crud = FileMetadataCRUD.create_loader(self.client)
+
     def as_id(self, item: FileMetadata) -> str:
         return item.external_id if item.external_id is not None else self._create_identifier(item.id)
 
-    def _get_loader(self) -> FileMetadataCRUD:
-        return FileMetadataCRUD.create_loader(self.client)
-
     def _get_aggregator(self) -> AssetCentricAggregator:
         return FileAggregator(self.client)
 
     def get_schema(self, selector: AssetCentricSelector) -> list[SchemaColumn]:
         data_set_ids: list[int] = []
         if isinstance(selector, DataSetSelector):
-            data_set_ids.append(self.client.lookup.data_sets.id(selector.data_set_external_id))
+            data_set_id = self.client.lookup.data_sets.id(selector.data_set_external_id)
+            if data_set_id is None:
+                raise ToolkitMissingResourceError(
+                    f"Data set with external ID {selector.data_set_external_id} not found."
+                )
+            data_set_ids.append(data_set_id)
         if isinstance(selector, AssetSubtreeSelector):
             raise ToolkitNotImplementedError(f"Selector type {type(selector)} not supported for FileIO.")
 
@@ -345,8 +370,19 @@ class FileMetadataIO(BaseAssetCentricIO[str, FileMetadataWrite, FileMetadata, Fi
     def retrieve(self, ids: Sequence[int]) -> FileMetadataList:
         return self.client.files.retrieve_multiple(ids)
 
+    def data_to_json_chunk(
+        self, data_chunk: Sequence[FileMetadata], selector: AssetCentricSelector | None = None
+    ) -> list[dict[str, JsonVal]]:
+        # Ensure data sets/assets/security-categories are looked up to populate cache.
+        # This is to avoid looking up each data set id individually in the .dump_resource call
+        self._populate_data_set_cache(data_chunk)
+        self._populate_asset_cache(data_chunk)
+        self._populate_security_category_cache(data_chunk)
+
+        return [self._crud.dump_resource(item) for item in data_chunk]
+
     def json_to_resource(self, item_json: dict[str, JsonVal]) -> FileMetadataWrite:
-        return self._loader.load_resource(item_json)
+        return self._crud.load_resource(item_json)
 
 
 class TimeSeriesIO(BaseAssetCentricIO[str, TimeSeriesWrite, TimeSeries, TimeSeriesWriteList, TimeSeriesList]):
@@ -357,12 +393,13 @@ class TimeSeriesIO(BaseAssetCentricIO[str, TimeSeriesWrite, TimeSeries, TimeSeri
     UPLOAD_ENDPOINT = "/timeseries"
     RESOURCE_TYPE = "timeseries"
 
+    def __init__(self, client: ToolkitClient) -> None:
+        super().__init__(client)
+        self._crud = TimeSeriesCRUD.create_loader(self.client)
+
     def as_id(self, item: TimeSeries) -> str:
         return item.external_id if item.external_id is not None else self._create_identifier(item.id)
 
-    def _get_loader(self) -> TimeSeriesCRUD:
-        return TimeSeriesCRUD.create_loader(self.client)
-
     def _get_aggregator(self) -> AssetCentricAggregator:
         return TimeSeriesAggregator(self.client)
 
@@ -380,13 +417,29 @@ class TimeSeriesIO(BaseAssetCentricIO[str, TimeSeriesWrite, TimeSeries, TimeSeri
             self._collect_dependencies(ts_list, selector)
             yield Page(worker_id="main", items=ts_list)
 
+    def data_to_json_chunk(
+        self, data_chunk: Sequence[TimeSeries], selector: AssetCentricSelector | None = None
+    ) -> list[dict[str, JsonVal]]:
+        # Ensure data sets/assets/security categories are looked up to populate cache.
+        self._populate_data_set_cache(data_chunk)
+        self._populate_security_category_cache(data_chunk)
+        asset_ids = {item.asset_id for item in data_chunk if item.asset_id is not None}
+        self.client.lookup.assets.external_id(list(asset_ids))
+
+        return [self._crud.dump_resource(item) for item in data_chunk]
+
     def json_to_resource(self, item_json: dict[str, JsonVal]) -> TimeSeriesWrite:
-        return self._loader.load_resource(item_json)
+        return self._crud.load_resource(item_json)
 
     def get_schema(self, selector: AssetCentricSelector) -> list[SchemaColumn]:
         data_set_ids: list[int] = []
         if isinstance(selector, DataSetSelector):
-            data_set_ids.append(self.client.lookup.data_sets.id(selector.data_set_external_id))
+            data_set_id = self.client.lookup.data_sets.id(selector.data_set_external_id)
+            if data_set_id is None:
+                raise ToolkitMissingResourceError(
+                    f"Data set with external ID {selector.data_set_external_id} not found."
+                )
+            data_set_ids.append(data_set_id)
         elif isinstance(selector, AssetSubtreeSelector):
             raise ToolkitNotImplementedError(f"Selector type {type(selector)} not supported for {type(self).__name__}.")
 
@@ -424,19 +477,25 @@ class EventIO(BaseAssetCentricIO[str, EventWrite, Event, EventWriteList, EventLi
     UPLOAD_ENDPOINT = "/events"
     RESOURCE_TYPE = "event"
 
+    def __init__(self, client: ToolkitClient) -> None:
+        super().__init__(client)
+        self._crud = EventCRUD.create_loader(self.client)
+
     def as_id(self, item: Event) -> str:
         return item.external_id if item.external_id is not None else self._create_identifier(item.id)
 
-    def _get_loader(self) -> EventCRUD:
-        return EventCRUD.create_loader(self.client)
-
     def _get_aggregator(self) -> AssetCentricAggregator:
         return EventAggregator(self.client)
 
     def get_schema(self, selector: AssetCentricSelector) -> list[SchemaColumn]:
         data_set_ids: list[int] = []
         if isinstance(selector, DataSetSelector):
-            data_set_ids.append(self.client.lookup.data_sets.id(selector.data_set_external_id))
+            data_set_id = self.client.lookup.data_sets.id(selector.data_set_external_id)
+            if data_set_id is None:
+                raise ToolkitMissingResourceError(
+                    f"Data set with external ID {selector.data_set_external_id} not found."
+                )
+            data_set_ids.append(data_set_id)
         hierarchy: list[int] = []
         if isinstance(selector, AssetSubtreeSelector):
             raise ToolkitNotImplementedError(f"Selector type {type(selector)} not supported for {type(self).__name__}.")
@@ -476,8 +535,17 @@ class EventIO(BaseAssetCentricIO[str, EventWrite, Event, EventWriteList, EventLi
             self._collect_dependencies(event_list, selector)
             yield Page(worker_id="main", items=event_list)
 
+    def data_to_json_chunk(
+        self, data_chunk: Sequence[Event], selector: AssetCentricSelector | None = None
+    ) -> list[dict[str, JsonVal]]:
+        # Ensure data sets/assets are looked up to populate cache.
+        self._populate_data_set_cache(data_chunk)
+        self._populate_asset_cache(data_chunk)
+
+        return [self._crud.dump_resource(item) for item in data_chunk]
+
     def json_to_resource(self, item_json: dict[str, JsonVal]) -> EventWrite:
-        return self._loader.load_resource(item_json)
+        return self._crud.load_resource(item_json)
 
     def retrieve(self, ids: Sequence[int]) -> EventList:
         return self.client.events.retrieve_multiple(ids)

cognite_toolkit/_cdf_tk/utils/aggregators.py

@@ -18,6 +18,7 @@ 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,
@@ -105,6 +106,15 @@ class AssetCentricAggregator(ABC):
                 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]
@@ -125,7 +135,9 @@ class MetadataAggregator(AssetCentricAggregator, ABC, Generic[T_CogniteFilter]):
         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)
+        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)
@@ -140,22 +152,44 @@ class MetadataAggregator(AssetCentricAggregator, ABC, Generic[T_CogniteFilter]):
         )
 
     def _lookup_hierarchy_data_set_pair(
-        self, hierarchy: str | list[str] | None = None, data_sets: str | list[str] | None = None
+        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):
-            hierarchy_ids = tuple(sorted(self.client.lookup.assets.id(external_id=hierarchy, allow_empty=False)))
+            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 = tuple(sorted(self.client.lookup.data_sets.id(external_id=data_sets, allow_empty=False)))
+            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
 
@@ -201,7 +235,9 @@ class LabelAggregator(MetadataAggregator, ABC, Generic[T_CogniteFilter]):
         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)
+        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)
@@ -359,9 +395,10 @@ class RelationshipAggregator(AssetCentricAggregator):
     def count(
         self, hierarchy: str | list[str] | None = None, data_set_external_id: str | list[str] | None = None
     ) -> int:
-        if hierarchy is not None or data_set_external_id is not None:
+        if hierarchy is not None:
             raise NotImplementedError()
-        results = relationship_aggregate_count(self.client)
+        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]:
@@ -378,9 +415,10 @@ class LabelCountAggregator(AssetCentricAggregator):
     def count(
         self, hierarchy: str | list[str] | None = None, data_set_external_id: str | list[str] | None = None
     ) -> int:
-        if hierarchy is not None or data_set_external_id is not None:
+        if hierarchy is not None:
             raise NotImplementedError()
-        return label_aggregate_count(self.client)
+        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()