cognite-toolkit 0.6.82__py3-none-any.whl → 0.6.84__py3-none-any.whl

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,
@@ -125,7 +126,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 +143,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 +226,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)

cognite_toolkit/_cdf_tk/utils/modules.py

@@ -4,7 +4,6 @@ from pathlib import Path
 from typing import Literal, overload
 
 from cognite_toolkit._cdf_tk.constants import (
-    BUILTIN_MODULES,
     MODULE_PATH_SEP,
     ROOT_MODULES,
 )
@@ -23,9 +22,6 @@ def iterate_modules(root_dir: Path) -> Iterator[tuple[Path, list[Path]]]:
     if root_dir.name in ROOT_MODULES:
         yield from _iterate_modules(root_dir)
         return
-    elif root_dir.name == BUILTIN_MODULES:
-        yield from _iterate_modules(root_dir)
-        return
     for root_module in ROOT_MODULES:
         module_dir = root_dir / root_module
         if module_dir.exists():

cognite_toolkit/_repo_files/GitHub/.github/workflows/deploy.yaml

@@ -12,7 +12,7 @@ jobs:
     environment: dev
     name: Deploy
     container:
-      image: cognite/toolkit:0.6.82
+      image: cognite/toolkit:0.6.84
       env:
         CDF_CLUSTER: ${{ vars.CDF_CLUSTER }}
         CDF_PROJECT: ${{ vars.CDF_PROJECT }}

cognite_toolkit/_repo_files/GitHub/.github/workflows/dry-run.yaml

@@ -10,7 +10,7 @@ jobs:
     environment: dev
     name: Deploy Dry Run
     container:
-      image: cognite/toolkit:0.6.82
+      image: cognite/toolkit:0.6.84
       env:
         CDF_CLUSTER: ${{ vars.CDF_CLUSTER }}
         CDF_PROJECT: ${{ vars.CDF_PROJECT }}

cognite_toolkit/_resources/cdf.toml

@@ -4,10 +4,14 @@ default_env = "<DEFAULT_ENV_PLACEHOLDER>"
 [modules]
 # This is the version of the modules. It should not be changed manually.
 # It will be updated by the 'cdf modules upgrade' command.
-version = "0.6.82"
+version = "0.6.84"
 
+[alpha_flags]
+external-libraries = true
 
 [plugins]
 run = true
-pull = false
-dump = false
+
+[library.toolkit-data]
+url = "https://github.com/cognitedata/toolkit-data/releases/download/latest/packages.zip" 
+checksum = "sha256:b70c207a78d6c94e0977a733571420b250ca641ac8802d8310f1ffe05bf7b2eb"

cognite_toolkit/_version.py

@@ -1 +1 @@
-__version__ = "0.6.82"
+__version__ = "0.6.84"

{cognite_toolkit-0.6.82.dist-info → cognite_toolkit-0.6.84.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cognite_toolkit
-Version: 0.6.82
+Version: 0.6.84
 Summary: Official Cognite Data Fusion tool for project templates and configuration deployment
 Project-URL: Homepage, https://docs.cognite.com/cdf/deploy/cdf_toolkit/
 Project-URL: Changelog, https://github.com/cognitedata/toolkit/releases