cognite-toolkit 0.7.39__py3-none-any.whl → 0.7.41__py3-none-any.whl

cognite_toolkit/_cdf_tk/commands/pull.py

@@ -1,5 +1,6 @@
 import dataclasses
 import itertools
+import json
 import re
 import sys
 import tempfile
@@ -28,6 +29,7 @@ from cognite_toolkit._cdf_tk.cruds import (
     HostedExtractorSourceCRUD,
     ResourceCRUD,
     StreamlitCRUD,
+    ViewCRUD,
 )
 from cognite_toolkit._cdf_tk.data_classes import (
     BuildEnvironment,
@@ -643,6 +645,44 @@ class PullCommand(ToolkitCommand):
         loader: ResourceCRUD[T_ID, T_ResourceRequest, T_ResourceResponse],
         source_file: Path,
     ) -> tuple[str, dict[Path, str]]:
+        """Convert resource data from CDF into YAML file content ready to be written to disk.
+
+        This method takes the raw CDF resource data and transforms it back into a properly
+        formatted YAML file that preserves:
+        - Template variables (e.g., {{ variable_name }}) instead of their resolved values
+        - YAML comments from the original source file
+        - The original key ordering in dictionaries
+
+        The transformation process:
+        1. Replace all template variables with unique placeholders
+        2. Load source YAML content while preserving comments
+        3. Update the resource data with placeholder values where variables were used
+        4. Dump the updated data back to YAML format
+        5. Replace placeholders with the original template variable syntax
+        6. Restore the YAML comments
+
+        Args:
+            source: The original YAML file content as a string.
+            to_write: A mapping from resource identifiers to their updated data dictionaries
+                pulled from CDF.
+            resources: The list of built resources containing build variables and metadata.
+            environment_variables: A mapping of environment variable names to their values,
+                used to resolve variables like ${VAR_NAME} in template values.
+            loader: The ResourceCRUD loader instance for this resource type.
+            source_file: The path to the source file being processed.
+
+        Returns:
+            A tuple containing:
+            - The final YAML content string ready to be written to disk.
+            - A dictionary mapping extra file paths to their content (for resources
+              that have additional files, like SQL queries for transformations).
+
+        Raises:
+            ValueError: If the loaded YAML structure doesn't match between the original
+                and placeholder versions.
+            ToolkitMissingResourceError: If a resource identifier is not found in the
+                to_write or resources mappings.
+        """
         # 1. Replace all variables with placeholders
         # 2. Load source and keep the comments
         # 3. Update the to_write dict with the placeholders
@@ -762,9 +802,22 @@ class PullCommand(ToolkitCommand):
 
 
 class ResourceReplacer:
-    """Replaces values in a local resource directory with the updated values from CDF.
+    """Replaces values in a local resource dictionary with the updated values from CDF.
 
     The local resource dict order is maintained. In addition, placeholders are used for variables.
+
+    This class is responsible for merging CDF resource values back into local configuration files
+    while preserving:
+    - The original key ordering in dictionaries
+    - Template variable placeholders (e.g., {{ variable_name }})
+    - Comments and formatting where possible
+
+    Args:
+        value_by_placeholder: A mapping from placeholder strings to their corresponding
+            BuildVariable objects. Placeholders are temporary substitutes for template
+            variables during processing.
+        loader: The ResourceCRUD loader instance used to determine how to diff lists
+            and handle resource-specific logic.
     """
 
     def __init__(self, value_by_placeholder: dict[str, BuildVariable], loader: ResourceCRUD) -> None:
@@ -777,7 +830,49 @@ class ResourceReplacer:
         placeholder: dict[str, Any],
         to_write: dict[str, Any],
     ) -> dict[str, Any]:
-        return self._replace_dict(current, placeholder, to_write, tuple())
+        """Replace values in a local resource dict with updated values from CDF.
+
+        Merges the CDF resource values into the local configuration while maintaining
+        the original dictionary key ordering and preserving template variable placeholders.
+
+        Args:
+            current: The current local resource dictionary with resolved variable values.
+                This represents the source file content after template variables have
+                been substituted with their actual values.
+            placeholder: The local resource dictionary with placeholder strings instead
+                of resolved values. Used to identify which values contain template
+                variables that should be preserved.
+            to_write: The resource dictionary from CDF containing the updated values
+                to merge into the local configuration.
+
+        Returns:
+            A new dictionary with CDF values merged in, maintaining the original key
+            order from `current`. Template variables are preserved as placeholders
+            (to be converted back to {{ variable }} syntax by the caller). New keys
+            from CDF are appended at the end, and removed keys are omitted.
+
+        Raises:
+            ToolkitValueError: If a list variable has changed and cannot be updated,
+                or if there's a type mismatch between local and CDF values.
+        """
+        has_stringified_view_filter = False
+        if isinstance(self._loader, ViewCRUD):
+            # view.filter are recursive nested dicts that are complex. To avoid issues with comparing
+            # lists inside the filters, we stringify them before processing such that they are compared
+            # as strings.
+            processed = []
+            for d in (current, placeholder, to_write):
+                if isinstance(d.get("filter"), dict):
+                    d = d.copy()
+                    d["filter"] = json.dumps(d["filter"])
+                    has_stringified_view_filter = True
+                processed.append(d)
+            current, placeholder, to_write = processed
+        output = self._replace_dict(current, placeholder, to_write, tuple())
+        if has_stringified_view_filter and "filter" in output:
+            # Special case for ViewCRUD where the filter is stringified in CDF
+            output["filter"] = json.loads(output["filter"])
+        return output
 
     def _replace_dict(
         self,

cognite_toolkit/_cdf_tk/cruds/_resource_cruds/classic.py

@@ -6,10 +6,6 @@ from typing import Any, final
 
 import pandas as pd
 from cognite.client.data_classes import (
-    AggregateResultItem,
-    Asset,
-    AssetList,
-    AssetWrite,
     Event,
     EventList,
     EventWrite,
@@ -24,6 +20,12 @@ from cognite.client.utils.useful_types import SequenceNotStr
 from rich.console import Console
 
 from cognite_toolkit._cdf_tk.client import ToolkitClient
+from cognite_toolkit._cdf_tk.client.data_classes.asset import (
+    AssetAggregateItem,
+    AssetRequest,
+    AssetResponse,
+)
+from cognite_toolkit._cdf_tk.client.data_classes.identifiers import ExternalId, InternalOrExternalId
 from cognite_toolkit._cdf_tk.client.data_classes.legacy.sequences import (
     ToolkitSequenceRows,
     ToolkitSequenceRowsList,
@@ -44,10 +46,10 @@ _DEPRECATION_WARNING_ISSUED = False
 
 
 @final
-class AssetCRUD(ResourceCRUD[str, AssetWrite, Asset]):
+class AssetCRUD(ResourceCRUD[ExternalId, AssetRequest, AssetResponse]):
     folder_name = "classic"
-    resource_cls = Asset
-    resource_write_cls = AssetWrite
+    resource_cls = AssetResponse
+    resource_write_cls = AssetRequest
     yaml_cls = AssetYAML
     kind = "Asset"
     dependencies = frozenset({DataSetsCRUD, LabelCRUD})
@@ -78,15 +80,15 @@ class AssetCRUD(ResourceCRUD[str, AssetWrite, Asset]):
         return "assets"
 
     @classmethod
-    def get_id(cls, item: Asset | AssetWrite | dict) -> str:
+    def get_id(cls, item: AssetRequest | AssetResponse | dict) -> ExternalId:
         if isinstance(item, dict):
-            return item["externalId"]
+            return ExternalId(external_id=item["externalId"])
         if not item.external_id:
             raise KeyError("Asset must have external_id")
-        return item.external_id
+        return ExternalId(external_id=item.external_id)
 
     @classmethod
-    def get_internal_id(cls, item: Asset | dict) -> int:
+    def get_internal_id(cls, item: AssetResponse | dict) -> int:
         if isinstance(item, dict):
             return item["id"]
         if not item.id:
@@ -94,12 +96,12 @@ class AssetCRUD(ResourceCRUD[str, AssetWrite, Asset]):
         return item.id
 
     @classmethod
-    def dump_id(cls, id: str) -> dict[str, Any]:
-        return {"externalId": id}
+    def dump_id(cls, id: ExternalId) -> dict[str, Any]:
+        return id.dump()
 
     @classmethod
     def get_required_capability(
-        cls, items: collections.abc.Sequence[AssetWrite] | None, read_only: bool
+        cls, items: collections.abc.Sequence[AssetRequest] | None, read_only: bool
     ) -> Capability | list[Capability]:
         if not items and items is not None:
             return []
@@ -119,45 +121,39 @@ class AssetCRUD(ResourceCRUD[str, AssetWrite, Asset]):
 
         return capabilities.AssetsAcl(actions, scope)
 
-    def create(self, items: Sequence[AssetWrite]) -> AssetList:
-        return self.client.assets.create(items)
+    def create(self, items: collections.abc.Sequence[AssetRequest]) -> list[AssetResponse]:
+        return self.client.tool.assets.create(items)
 
-    def retrieve(self, ids: SequenceNotStr[str]) -> AssetList:
-        return self.client.assets.retrieve_multiple(external_ids=ids, ignore_unknown_ids=True)
+    def retrieve(self, ids: SequenceNotStr[ExternalId]) -> list[AssetResponse]:
+        return self.client.tool.assets.retrieve(list(ids), ignore_unknown_ids=True)
 
-    def update(self, items: Sequence[AssetWrite]) -> AssetList:
-        return self.client.assets.update(items, mode="replace")
+    def update(self, items: collections.abc.Sequence[AssetRequest]) -> list[AssetResponse]:
+        return self.client.tool.assets.update(items, mode="replace")
 
-    def delete(self, ids: SequenceNotStr[str | int]) -> int:
+    def delete(self, ids: SequenceNotStr[InternalOrExternalId]) -> int:
         if not ids:
             return 0
-        internal_ids, external_ids = self._split_ids(ids)
-        try:
-            self.client.assets.delete(id=internal_ids, external_id=external_ids)
-        except CogniteNotFoundError as e:
-            # Do a CogniteNotFoundError instead of passing 'ignore_unknown_ids=True' to the delete method
-            # to obtain an accurate list of deleted assets.
-            non_existing = set(e.failed or [])
-            if existing := [id_ for id_ in ids if id_ not in non_existing]:
-                internal_ids, external_ids = self._split_ids(existing)
-                self.client.assets.delete(id=internal_ids, external_id=external_ids)
-            return len(existing)
-        else:
-            return len(ids)
+        self.client.tool.assets.delete(list(ids), ignore_unknown_ids=True)
+        return len(ids)
 
     def _iterate(
         self,
         data_set_external_id: str | None = None,
         space: str | None = None,
         parent_ids: list[Hashable] | None = None,
-    ) -> Iterable[Asset]:
-        return iter(
-            self.client.assets(
+    ) -> Iterable[AssetResponse]:
+        cursor: str | None = None
+        while True:
+            page = self.client.tool.assets.iterate(
+                limit=1000,
+                cursor=cursor,
                 data_set_external_ids=[data_set_external_id] if data_set_external_id else None,
-                # This is used in the purge command to delete the children before the parent.
-                aggregated_properties=["depth", "child_count", "path"],
+                aggregated_properties=True,
             )
-        )
+            yield from page.items
+            if not page.next_cursor or not page.items:
+                break
+            cursor = page.next_cursor
 
     @classmethod
     def get_dependent_items(cls, item: dict) -> Iterable[tuple[type[ResourceCRUD], Hashable]]:
@@ -202,7 +198,7 @@ class AssetCRUD(ResourceCRUD[str, AssetWrite, Asset]):
 
         return resources
 
-    def load_resource(self, resource: dict[str, Any], is_dry_run: bool = False) -> AssetWrite:
+    def load_resource(self, resource: dict[str, Any], is_dry_run: bool = False) -> AssetRequest:
         # Unpack metadata keys from table formats (e.g. csv, parquet)
         metadata: dict = resource.get("metadata", {})
         for key, value in list(resource.items()):
@@ -219,10 +215,10 @@ class AssetCRUD(ResourceCRUD[str, AssetWrite, Asset]):
 
         if ds_external_id := resource.pop("dataSetExternalId", None):
             resource["dataSetId"] = self.client.lookup.data_sets.id(ds_external_id, is_dry_run)
-        return AssetWrite._load(resource)
+        return AssetRequest.model_validate(resource)
 
-    def dump_resource(self, resource: Asset, local: dict[str, Any] | None = None) -> dict[str, Any]:
-        dumped = resource.as_write().dump()
+    def dump_resource(self, resource: AssetResponse, local: dict[str, Any] | None = None) -> dict[str, Any]:
+        dumped = resource.as_request_resource().dump()
         local = local or {}
         if data_set_id := dumped.pop("dataSetId", None):
             dumped["dataSetExternalId"] = self.client.lookup.data_sets.external_id(data_set_id)
@@ -234,7 +230,7 @@ class AssetCRUD(ResourceCRUD[str, AssetWrite, Asset]):
             # This is only included when the asset is downloaded/migrated with aggregated properties
             aggregates = (
                 resource.aggregates.dump()
-                if isinstance(resource.aggregates, AggregateResultItem)
+                if isinstance(resource.aggregates, AssetAggregateItem)
                 else resource.aggregates
             )
             if "path" in aggregates:
@@ -372,7 +368,7 @@ class SequenceCRUD(ResourceCRUD[str, SequenceWrite, CDFSequence]):
         if "dataSetExternalId" in item:
             yield DataSetsCRUD, item["dataSetExternalId"]
         if "assetExternalId" in item:
-            yield AssetCRUD, item["assetExternalId"]
+            yield AssetCRUD, ExternalId(external_id=item["assetExternalId"])
 
 
 @final
@@ -596,7 +592,7 @@ class EventCRUD(ResourceCRUD[str, EventWrite, Event]):
             yield DataSetsCRUD, item["dataSetExternalId"]
         for asset_id in item.get("assetExternalIds", []):
             if isinstance(asset_id, str):
-                yield AssetCRUD, asset_id
+                yield AssetCRUD, ExternalId(external_id=asset_id)
 
     def load_resource(self, resource: dict[str, Any], is_dry_run: bool = False) -> EventWrite:
         if ds_external_id := resource.get("dataSetExternalId", None):

cognite_toolkit/_cdf_tk/cruds/_resource_cruds/fieldops.py

@@ -8,6 +8,7 @@ from cognite.client.data_classes.data_modeling import NodeApplyResultList, NodeI
 from cognite.client.exceptions import CogniteAPIError
 from cognite.client.utils.useful_types import SequenceNotStr
 
+from cognite_toolkit._cdf_tk.client.data_classes.identifiers import ExternalId
 from cognite_toolkit._cdf_tk.client.data_classes.infield import (
     InFieldCDMLocationConfig,
     InfieldLocationConfig,
@@ -139,7 +140,7 @@ class InfieldV1CRUD(ResourceCRUD[str, APMConfigWrite, APMConfig]):
             yield SpaceCRUD, customer_data_space_id
         for config in cls._get_root_location_configurations(item) or []:
             if isinstance(asset_external_id := config.get("assetExternalId"), str):
-                yield AssetCRUD, asset_external_id
+                yield AssetCRUD, ExternalId(external_id=asset_external_id)
             if isinstance(data_set_external_id := config.get("dataSetExternalId"), str):
                 yield DataSetsCRUD, data_set_external_id
             if isinstance(app_data_instance_space := config.get("appDataInstanceSpace"), str):
@@ -162,7 +163,7 @@ class InfieldV1CRUD(ResourceCRUD[str, APMConfigWrite, APMConfig]):
                         yield DataSetsCRUD, data_set_external_id
                 for asset_external_id in filter_.get("assetSubtreeExternalIds", []):
                     if isinstance(asset_external_id, str):
-                        yield AssetCRUD, asset_external_id
+                        yield AssetCRUD, ExternalId(external_id=asset_external_id)
                 if app_data_instance_space := filter_.get("appDataInstanceSpace"):
                     if isinstance(app_data_instance_space, str):
                         yield SpaceCRUD, app_data_instance_space
@@ -267,7 +268,10 @@ class InFieldLocationConfigCRUD(ResourceCRUD[TypedNodeIdentifier, InfieldLocatio
 
     @classmethod
     def dump_id(cls, id: TypedNodeIdentifier) -> dict[str, Any]:
-        return id.dump(include_type=False)
+        return {
+            "space": id.space,
+            "externalId": id.external_id,
+        }
 
     @classmethod
     def get_required_capability(
@@ -368,7 +372,10 @@ class InFieldCDMLocationConfigCRUD(
 
     @classmethod
     def dump_id(cls, id: TypedNodeIdentifier) -> dict[str, Any]:
-        return id.dump(include_type=False)
+        return {
+            "space": id.space,
+            "externalId": id.external_id,
+        }
 
     @classmethod
     def get_required_capability(

cognite_toolkit/_cdf_tk/cruds/_resource_cruds/file.py

@@ -33,6 +33,7 @@ from cognite.client.utils._time import convert_data_modelling_timestamp
 from cognite.client.utils.useful_types import SequenceNotStr
 from rich import print
 
+from cognite_toolkit._cdf_tk.client.data_classes.identifiers import ExternalId
 from cognite_toolkit._cdf_tk.client.data_classes.legacy.extendable_cognite_file import (
     ExtendableCogniteFile,
     ExtendableCogniteFileApply,
@@ -118,7 +119,7 @@ class FileMetadataCRUD(ResourceContainerCRUD[str, FileMetadataWrite, FileMetadat
                 elif isinstance(label, str):
                     yield LabelCRUD, label
         for asset_external_id in item.get("assetExternalIds", []):
-            yield AssetCRUD, asset_external_id
+            yield AssetCRUD, ExternalId(external_id=asset_external_id)
 
     def load_resource(self, resource: dict[str, Any], is_dry_run: bool = False) -> FileMetadataWrite:
         if resource.get("dataSetExternalId") is not None:

cognite_toolkit/_cdf_tk/cruds/_resource_cruds/location.py

@@ -7,6 +7,7 @@ from cognite.client.data_classes.capabilities import Capability, LocationFilters
 from cognite.client.data_classes.data_modeling import DataModelId, ViewId
 from cognite.client.utils.useful_types import SequenceNotStr
 
+from cognite_toolkit._cdf_tk.client.data_classes.identifiers import ExternalId
 from cognite_toolkit._cdf_tk.client.data_classes.legacy.location_filters import (
     LocationFilter,
     LocationFilterList,
@@ -239,14 +240,14 @@ class LocationFilterCRUD(ResourceCRUD[str, LocationFilterWrite, LocationFilter])
                 yield DataSetsCRUD, data_set_external_id
             for asset in asset_centric.get("assetSubtreeIds", []):
                 if "externalId" in asset:
-                    yield AssetCRUD, asset["externalId"]
+                    yield AssetCRUD, ExternalId(external_id=asset["externalId"])
             for subfilter_name in cls.subfilter_names:
                 subfilter = asset_centric.get(subfilter_name, {})
                 for data_set_external_id in subfilter.get("dataSetExternalIds", []):
                     yield DataSetsCRUD, data_set_external_id
                 for asset in subfilter.get("assetSubtreeIds", []):
                     if "externalId" in asset:
-                        yield AssetCRUD, asset["externalId"]
+                        yield AssetCRUD, ExternalId(external_id=asset["externalId"])
         for view in item.get("views", []):
             if in_dict(["space", "externalId", "version"], view):
                 yield ViewCRUD, ViewId(view["space"], view["externalId"], view["version"])

cognite_toolkit/_cdf_tk/cruds/_resource_cruds/relationship.py

@@ -11,6 +11,7 @@ from cognite.client.data_classes.capabilities import Capability
 from cognite.client.exceptions import CogniteAPIError, CogniteNotFoundError
 from cognite.client.utils.useful_types import SequenceNotStr
 
+from cognite_toolkit._cdf_tk.client.data_classes.identifiers import ExternalId
 from cognite_toolkit._cdf_tk.cruds._base_cruds import ResourceCRUD
 from cognite_toolkit._cdf_tk.resource_classes import RelationshipYAML
 
@@ -123,7 +124,7 @@ class RelationshipCRUD(ResourceCRUD[str, RelationshipWrite, Relationship]):
                 if isinstance(id_value, str) and isinstance(type_value, str):
                     type_value = type_value.strip().casefold()
                     if type_value == "asset":
-                        yield AssetCRUD, id_value
+                        yield AssetCRUD, ExternalId(external_id=id_value)
                     elif type_value == "sequence":
                         yield SequenceCRUD, id_value
                     elif type_value == "timeseries":

cognite_toolkit/_cdf_tk/cruds/_resource_cruds/streams.py

@@ -9,9 +9,9 @@ from cognite_toolkit._cdf_tk.client.data_classes.streams import (
     StreamResponse,
     StreamResponseList,
 )
+from cognite_toolkit._cdf_tk.client.http_client import ToolkitAPIError
 from cognite_toolkit._cdf_tk.cruds._base_cruds import ResourceCRUD
 from cognite_toolkit._cdf_tk.resource_classes import StreamYAML
-from cognite_toolkit._cdf_tk.utils.http_client import ToolkitAPIError
 
 from .datamodel import ContainerCRUD
 

cognite_toolkit/_cdf_tk/cruds/_resource_cruds/timeseries.py

@@ -23,6 +23,7 @@ from cognite.client.data_classes.datapoints_subscriptions import TimeSeriesIDLis
 from cognite.client.exceptions import CogniteAPIError, CogniteNotFoundError
 from cognite.client.utils.useful_types import SequenceNotStr
 
+from cognite_toolkit._cdf_tk.client.data_classes.identifiers import ExternalId
 from cognite_toolkit._cdf_tk.constants import MAX_TIMESTAMP_MS, MIN_TIMESTAMP_MS
 from cognite_toolkit._cdf_tk.cruds._base_cruds import ResourceContainerCRUD, ResourceCRUD
 from cognite_toolkit._cdf_tk.exceptions import (
@@ -97,7 +98,7 @@ class TimeSeriesCRUD(ResourceContainerCRUD[str, TimeSeriesWrite, TimeSeries]):
             for security_category in item["securityCategoryNames"]:
                 yield SecurityCategoryCRUD, security_category
         if "assetExternalId" in item:
-            yield AssetCRUD, item["assetExternalId"]
+            yield AssetCRUD, ExternalId(external_id=item["assetExternalId"])
 
     def load_resource(self, resource: dict[str, Any], is_dry_run: bool = False) -> TimeSeriesWrite:
         if ds_external_id := resource.pop("dataSetExternalId", None):

cognite_toolkit/_cdf_tk/storageio/_applications.py

@@ -7,10 +7,10 @@ from cognite_toolkit._cdf_tk.client.data_classes.legacy.canvas import (
     IndustrialCanvasApply,
 )
 from cognite_toolkit._cdf_tk.client.data_classes.legacy.charts import Chart, ChartList, ChartWrite
+from cognite_toolkit._cdf_tk.client.http_client import HTTPClient, HTTPMessage, SimpleBodyRequest
 from cognite_toolkit._cdf_tk.exceptions import ToolkitNotImplementedError
 from cognite_toolkit._cdf_tk.tk_warnings import HighSeverityWarning, MediumSeverityWarning
 from cognite_toolkit._cdf_tk.utils.collection import chunker_sequence
-from cognite_toolkit._cdf_tk.utils.http_client import HTTPClient, HTTPMessage, SimpleBodyRequest
 from cognite_toolkit._cdf_tk.utils.useful_types import JsonVal
 
 from ._base import Page, UploadableStorageIO, UploadItem

cognite_toolkit/_cdf_tk/storageio/_asset_centric.py

@@ -4,10 +4,6 @@ from collections.abc import Iterable, Sequence
 from typing import Any, ClassVar, Generic
 
 from cognite.client.data_classes import (
-    AggregateResultItem,
-    Asset,
-    AssetList,
-    AssetWrite,
     Event,
     EventList,
     EventWrite,
@@ -21,6 +17,8 @@ from cognite.client.data_classes import (
 )
 
 from cognite_toolkit._cdf_tk.client import ToolkitClient
+from cognite_toolkit._cdf_tk.client.data_classes.asset import AssetAggregateItem, AssetRequest, AssetResponse
+from cognite_toolkit._cdf_tk.client.data_classes.identifiers import InternalId
 from cognite_toolkit._cdf_tk.cruds import (
     AssetCRUD,
     DataSetsCRUD,
@@ -42,10 +40,10 @@ from cognite_toolkit._cdf_tk.utils.cdf import metadata_key_counts
 from cognite_toolkit._cdf_tk.utils.fileio import FileReader, SchemaColumn
 from cognite_toolkit._cdf_tk.utils.fileio._readers import TableReader
 from cognite_toolkit._cdf_tk.utils.useful_types import (
-    AssetCentricResource,
     AssetCentricType,
     JsonVal,
 )
+from cognite_toolkit._cdf_tk.utils.useful_types2 import AssetCentricResource
 
 from ._base import (
     ConfigurableStorageIO,
@@ -112,12 +110,14 @@ class AssetCentricIO(
         return asset_subtree_external_ids, data_set_external_ids
 
     def _collect_dependencies(
-        self, resources: AssetList | FileMetadataList | TimeSeriesList | EventList, selector: AssetCentricSelector
+        self,
+        resources: Sequence[AssetResponse] | FileMetadataList | TimeSeriesList | EventList,
+        selector: AssetCentricSelector,
     ) -> None:
         for resource in resources:
             if resource.data_set_id:
                 self._downloaded_data_sets_by_selector[selector].add(resource.data_set_id)
-            if isinstance(resource, Asset | FileMetadata):
+            if isinstance(resource, AssetResponse | FileMetadata):
                 for label in resource.labels or []:
                     if isinstance(label, str):
                         self._downloaded_labels_by_selector[selector].add(label)
@@ -149,7 +149,7 @@ class AssetCentricIO(
     def create_internal_identifier(cls, internal_id: int, project: str) -> str:
         return f"INTERNAL_ID_project_{project}_{internal_id!s}"
 
-    def _populate_data_set_id_cache(self, chunk: Sequence[Asset | FileMetadata | TimeSeries | Event]) -> None:
+    def _populate_data_set_id_cache(self, chunk: Sequence[AssetResponse | 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))
 
@@ -249,7 +249,7 @@ class UploadableAssetCentricIO(
         )
 
 
-class AssetIO(UploadableAssetCentricIO[Asset, AssetWrite]):
+class AssetIO(UploadableAssetCentricIO[AssetResponse, AssetRequest]):
     KIND = "Assets"
     RESOURCE_TYPE = "asset"
     SUPPORTED_DOWNLOAD_FORMATS = frozenset({".parquet", ".csv", ".ndjson"})
@@ -261,7 +261,7 @@ class AssetIO(UploadableAssetCentricIO[Asset, AssetWrite]):
         super().__init__(client)
         self._crud = AssetCRUD.create_loader(self.client)
 
-    def as_id(self, item: Asset) -> str:
+    def as_id(self, item: AssetResponse) -> str:
         return item.external_id if item.external_id is not None else self._create_identifier(item.id)
 
     def _get_aggregator(self) -> AssetCentricAggregator:
@@ -311,21 +311,25 @@ class AssetIO(UploadableAssetCentricIO[Asset, AssetWrite]):
 
     def stream_data(self, selector: AssetCentricSelector, limit: int | None = None) -> Iterable[Page]:
         asset_subtree_external_ids, data_set_external_ids = self._get_hierarchy_dataset_pair(selector)
-        for asset_list in self.client.assets(
-            chunk_size=self.CHUNK_SIZE,
-            limit=limit,
-            asset_subtree_external_ids=asset_subtree_external_ids,
-            data_set_external_ids=data_set_external_ids,
-            aggregated_properties=["child_count", "path", "depth"],
-            # We cannot use partitions here as it is not thread safe. This spawn multiple threads
-            # that are not shut down until all data is downloaded. We need to be able to abort.
-            partitions=None,
-        ):
-            self._collect_dependencies(asset_list, selector)
-            yield Page(worker_id="main", items=asset_list)
+        cursor: str | None = None
+        total_count = 0
+        while True:
+            page = self.client.tool.assets.iterate(
+                aggregated_properties=True,
+                data_set_external_ids=data_set_external_ids,
+                asset_subtree_external_ids=asset_subtree_external_ids,
+                limit=self.CHUNK_SIZE,
+                cursor=cursor,
+            )
+            self._collect_dependencies(page.items, selector)
+            yield Page(worker_id="main", items=page.items)
+            total_count += len(page.items)
+            if page.next_cursor is None or (limit is not None and total_count >= limit):
+                break
+            cursor = page.next_cursor
 
     def data_to_json_chunk(
-        self, data_chunk: Sequence[Asset], selector: AssetCentricSelector | None = None
+        self, data_chunk: Sequence[AssetResponse], 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.
@@ -333,19 +337,18 @@ class AssetIO(UploadableAssetCentricIO[Asset, AssetWrite]):
         asset_ids = {
             segment["id"]
             for item in data_chunk
-            if isinstance(item.aggregates, AggregateResultItem)
-            for segment in item.aggregates.path or []
-            if "id" in segment
+            if isinstance(item.aggregates, AssetAggregateItem)
+            for segment in item.aggregates.path
         }
         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]) -> AssetWrite:
+    def json_to_resource(self, item_json: dict[str, JsonVal]) -> AssetRequest:
         return self._crud.load_resource(item_json)
 
-    def retrieve(self, ids: Sequence[int]) -> AssetList:
-        return self.client.assets.retrieve_multiple(ids)
+    def retrieve(self, ids: Sequence[int]) -> list[AssetResponse]:
+        return self.client.tool.assets.retrieve(InternalId.from_ids(ids))
 
     @classmethod
     def read_chunks(

cognite_toolkit/_cdf_tk/storageio/_base.py

@@ -4,11 +4,11 @@ from dataclasses import dataclass
 from typing import ClassVar, Generic, Literal, TypeVar
 
 from cognite_toolkit._cdf_tk.client import ToolkitClient
+from cognite_toolkit._cdf_tk.client.http_client import HTTPClient, HTTPMessage, ItemsRequest
 from cognite_toolkit._cdf_tk.exceptions import ToolkitNotImplementedError
 from cognite_toolkit._cdf_tk.protocols import T_ResourceRequest, T_ResourceResponse
 from cognite_toolkit._cdf_tk.utils.collection import chunker
 from cognite_toolkit._cdf_tk.utils.fileio import MultiFileReader, SchemaColumn
-from cognite_toolkit._cdf_tk.utils.http_client import HTTPClient, HTTPMessage, ItemsRequest
 from cognite_toolkit._cdf_tk.utils.useful_types import JsonVal
 
 from .selectors import DataSelector

cognite_toolkit/_cdf_tk/storageio/_datapoints.py

@@ -15,6 +15,13 @@ from cognite.client.data_classes.filters import Exists
 from cognite.client.data_classes.time_series import TimeSeriesProperty
 
 from cognite_toolkit._cdf_tk.client import ToolkitClient
+from cognite_toolkit._cdf_tk.client.http_client import (
+    DataBodyRequest,
+    HTTPClient,
+    HTTPMessage,
+    SimpleBodyRequest,
+    SuccessResponse,
+)
 from cognite_toolkit._cdf_tk.exceptions import ToolkitNotImplementedError
 from cognite_toolkit._cdf_tk.tk_warnings import HighSeverityWarning
 from cognite_toolkit._cdf_tk.utils import humanize_collection
@@ -26,13 +33,6 @@ from cognite_toolkit._cdf_tk.utils.dtype_conversion import (
 )
 from cognite_toolkit._cdf_tk.utils.fileio import SchemaColumn
 from cognite_toolkit._cdf_tk.utils.fileio._readers import MultiFileReader
-from cognite_toolkit._cdf_tk.utils.http_client import (
-    DataBodyRequest,
-    HTTPClient,
-    HTTPMessage,
-    SimpleBodyRequest,
-    SuccessResponse,
-)
 from cognite_toolkit._cdf_tk.utils.useful_types import JsonVal
 
 from ._base import Page, TableStorageIO, TableUploadableStorageIO, UploadItem