cognite-toolkit 0.6.88__py3-none-any.whl → 0.6.89__py3-none-any.whl

cognite_toolkit/_cdf_tk/commands/_migrate/canvas.py

@@ -1,5 +1,11 @@
 from uuid import uuid4
 
+from cognite.client.data_classes.capabilities import (
+    Capability,
+    DataModelInstancesAcl,
+    DataModelsAcl,
+    SpaceIDScope,
+)
 from cognite.client.exceptions import CogniteException
 
 from cognite_toolkit._cdf_tk.client import ToolkitClient
@@ -10,16 +16,19 @@ from cognite_toolkit._cdf_tk.client.data_classes.canvas import (
     FdmInstanceContainerReferenceApply,
 )
 from cognite_toolkit._cdf_tk.client.data_classes.migration import InstanceSource
-from cognite_toolkit._cdf_tk.exceptions import ToolkitMigrationError
+from cognite_toolkit._cdf_tk.commands._base import ToolkitCommand
+from cognite_toolkit._cdf_tk.commands._migrate.data_model import (
+    INSTANCE_SOURCE_VIEW_ID,
+    MODEL_ID,
+    RESOURCE_VIEW_MAPPING_VIEW_ID,
+)
+from cognite_toolkit._cdf_tk.exceptions import AuthenticationError, ToolkitMigrationError
 from cognite_toolkit._cdf_tk.tk_warnings import HighSeverityWarning, LowSeverityWarning, MediumSeverityWarning
 from cognite_toolkit._cdf_tk.utils import humanize_collection
 from cognite_toolkit._cdf_tk.utils.interactive_select import InteractiveCanvasSelect
 
-from .base import BaseMigrateCommand
-from .data_model import INSTANCE_SOURCE_VIEW_ID
 
-
-class MigrationCanvasCommand(BaseMigrateCommand):
+class MigrationCanvasCommand(ToolkitCommand):
     canvas_schema_space = Canvas.get_source().space
     # Note sequences are not supported in Canvas, so we do not include them here.
     asset_centric_resource_types = frozenset({"asset", "event", "file", "timeseries"})
@@ -144,3 +153,49 @@ class MigrationCanvasCommand(BaseMigrateCommand):
             max_width=reference.max_width,
             max_height=reference.max_height,
         )
+
+    @staticmethod
+    def validate_access(
+        client: ToolkitClient,
+        instance_spaces: list[str] | None = None,
+        schema_spaces: list[str] | None = None,
+    ) -> None:
+        required_capabilities: list[Capability] = []
+        if instance_spaces is not None:
+            required_capabilities.append(
+                DataModelInstancesAcl(
+                    actions=[
+                        DataModelInstancesAcl.Action.Read,
+                        DataModelInstancesAcl.Action.Write,
+                        DataModelInstancesAcl.Action.Write_Properties,
+                    ],
+                    scope=SpaceIDScope(instance_spaces),
+                )
+            )
+        if schema_spaces is not None:
+            required_capabilities.append(
+                DataModelsAcl(actions=[DataModelsAcl.Action.Read], scope=SpaceIDScope(schema_spaces)),
+            )
+        if missing := client.iam.verify_capabilities(required_capabilities):
+            raise AuthenticationError(f"Missing required capabilities: {humanize_collection(missing)}.", missing)
+
+    @staticmethod
+    def validate_migration_model_available(client: ToolkitClient) -> None:
+        models = client.data_modeling.data_models.retrieve([MODEL_ID], inline_views=False)
+        if not models:
+            raise ToolkitMigrationError(
+                f"The migration data model {MODEL_ID!r} does not exist. "
+                "Please run the `cdf migrate prepare` command to deploy the migration data model."
+            )
+        elif len(models) > 1:
+            raise ToolkitMigrationError(
+                f"Multiple migration models {MODEL_ID!r}. "
+                "Please delete the duplicate models before proceeding with the migration."
+            )
+        model = models[0]
+        missing_views = {INSTANCE_SOURCE_VIEW_ID, RESOURCE_VIEW_MAPPING_VIEW_ID} - set(model.views or [])
+        if missing_views:
+            raise ToolkitMigrationError(
+                f"Invalid migration model. Missing views {humanize_collection(missing_views)}. "
+                f"Please run the `cdf migrate prepare` command to deploy the migration data model."
+            )

cognite_toolkit/_cdf_tk/commands/_migrate/command.py

@@ -162,12 +162,14 @@ class MigrationCommand(ToolkitCommand):
             for item in source:
                 target, issue = mapper.map(item)
                 id_ = data.as_id(item)
-                tracker.set_progress(id_, step=self.Steps.CONVERT, status="success")
+                result: Status = "failed" if target is None else "success"
+                tracker.set_progress(id_, step=self.Steps.CONVERT, status=result)
 
                 if issue.has_issues:
                     # MyPy fails to understand that dict[str, JsonVal] is a Chunk
                     issues.append(issue.dump())  # type: ignore[arg-type]
-                targets.append(UploadItem(source_id=id_, item=target))
+                if target is not None:
+                    targets.append(UploadItem(source_id=id_, item=target))
             if issues:
                 log_file.write_chunks(issues)
             return targets

cognite_toolkit/_cdf_tk/commands/_migrate/conversion.py

@@ -1,21 +1,29 @@
 from collections.abc import Mapping, Set
 from dataclasses import dataclass
-from typing import Any, ClassVar
+from typing import Any, ClassVar, overload
 
-from cognite.client.data_classes import Asset, Event, FileMetadata, Sequence, TimeSeries
-from cognite.client.data_classes._base import CogniteResource
-from cognite.client.data_classes.data_modeling import DirectRelationReference, MappedProperty, NodeApply, NodeId
-from cognite.client.data_classes.data_modeling.instances import NodeOrEdgeData, PropertyValueWrite
+from cognite.client.data_classes import Annotation, Asset, Event, FileMetadata, TimeSeries
+from cognite.client.data_classes.data_modeling import (
+    DirectRelation,
+    DirectRelationReference,
+    EdgeId,
+    MappedProperty,
+    NodeApply,
+    NodeId,
+)
+from cognite.client.data_classes.data_modeling.instances import EdgeApply, NodeOrEdgeData, PropertyValueWrite
 from cognite.client.data_classes.data_modeling.views import ViewProperty
 
-from cognite_toolkit._cdf_tk.client.data_classes.extended_filemetadata import ExtendedFileMetadata
-from cognite_toolkit._cdf_tk.client.data_classes.extended_timeseries import ExtendedTimeSeries
 from cognite_toolkit._cdf_tk.client.data_classes.migration import AssetCentricId, ResourceViewMapping
 from cognite_toolkit._cdf_tk.utils.collection import flatten_dict_json_path
 from cognite_toolkit._cdf_tk.utils.dtype_conversion import (
     asset_centric_convert_to_primary_property,
+    convert_to_primary_property,
+)
+from cognite_toolkit._cdf_tk.utils.useful_types import (
+    AssetCentricResourceExtended,
+    AssetCentricType,
 )
-from cognite_toolkit._cdf_tk.utils.useful_types import AssetCentricType
 
 from .data_model import INSTANCE_SOURCE_VIEW_ID
 from .issues import ConversionIssue, FailedConversion, InvalidPropertyDataType
@@ -49,37 +57,71 @@ class DirectRelationCache:
         ("event", "assetIds"),
         ("sequence", "assetId"),
         ("asset", "parentId"),
+        ("fileAnnotation", "data.assetRef.id"),
     }
     SOURCE_REFERENCE_PROPERTIES: ClassVar[Set[tuple[AssetCentricType, str]]] = {
         ("asset", "source"),
         ("event", "source"),
         ("file", "source"),
     }
+    FILE_REFERENCE_PROPERTIES: ClassVar[Set[tuple[AssetCentricType, str]]] = {
+        ("fileAnnotation", "data.fileRef.id"),
+        ("fileAnnotation", "annotatedResourceId"),
+    }
 
     asset: Mapping[int, DirectRelationReference]
     source: Mapping[str, DirectRelationReference]
+    file: Mapping[int, DirectRelationReference]
 
     def get(self, resource_type: AssetCentricType, property_id: str) -> Mapping[str | int, DirectRelationReference]:
-        if (resource_type, property_id) in self.ASSET_REFERENCE_PROPERTIES:
+        key = resource_type, property_id
+        if key in self.ASSET_REFERENCE_PROPERTIES:
             return self.asset  # type: ignore[return-value]
-        if (resource_type, property_id) in self.SOURCE_REFERENCE_PROPERTIES:
+        if key in self.SOURCE_REFERENCE_PROPERTIES:
             return self.source  # type: ignore[return-value]
+        if key in self.FILE_REFERENCE_PROPERTIES:
+            return self.file  # type: ignore[return-value]
         return {}
 
 
+@overload
 def asset_centric_to_dm(
-    resource: CogniteResource,
+    resource: AssetCentricResourceExtended,
     instance_id: NodeId,
     view_source: ResourceViewMapping,
     view_properties: dict[str, ViewProperty],
     asset_instance_id_by_id: Mapping[int, DirectRelationReference],
     source_instance_id_by_external_id: Mapping[str, DirectRelationReference],
-) -> tuple[NodeApply, ConversionIssue]:
+    file_instance_id_by_id: Mapping[int, DirectRelationReference],
+) -> tuple[NodeApply | None, ConversionIssue]: ...
+
+
+@overload
+def asset_centric_to_dm(
+    resource: AssetCentricResourceExtended,
+    instance_id: EdgeId,
+    view_source: ResourceViewMapping,
+    view_properties: dict[str, ViewProperty],
+    asset_instance_id_by_id: Mapping[int, DirectRelationReference],
+    source_instance_id_by_external_id: Mapping[str, DirectRelationReference],
+    file_instance_id_by_id: Mapping[int, DirectRelationReference],
+) -> tuple[EdgeApply | None, ConversionIssue]: ...
+
+
+def asset_centric_to_dm(
+    resource: AssetCentricResourceExtended,
+    instance_id: NodeId | EdgeId,
+    view_source: ResourceViewMapping,
+    view_properties: dict[str, ViewProperty],
+    asset_instance_id_by_id: Mapping[int, DirectRelationReference],
+    source_instance_id_by_external_id: Mapping[str, DirectRelationReference],
+    file_instance_id_by_id: Mapping[int, DirectRelationReference],
+) -> tuple[NodeApply | EdgeApply | None, ConversionIssue]:
     """Convert an asset-centric resource to a data model instance.
 
     Args:
         resource (CogniteResource): The asset-centric resource to convert.
-        instance_id (NodeId): The ID of the instance to create or update.
+        instance_id (NodeId | EdgeApply): The ID of the instance to create or update.
         view_source (ResourceViewMapping): The view source defining how to map the resource to the data model.
         view_properties (dict[str, ViewProperty]): The defined properties referenced in the view source mapping.
         asset_instance_id_by_id (dict[int, DirectRelationReference]): A mapping from asset IDs to their corresponding
@@ -88,12 +130,17 @@ def asset_centric_to_dm(
         source_instance_id_by_external_id (dict[str, DirectRelationReference]): A mapping from source strings to their
             corresponding DirectRelationReference in the data model. This is used to create direct relations for resources
             that reference sources.
+        file_instance_id_by_id (dict[int, DirectRelationReference]): A mapping from file IDs to their corresponding
+            DirectRelationReference in the data model. This is used to create direct relations for resources that
+            reference files.
 
     Returns:
-        tuple[NodeApply, ConversionIssue]: A tuple containing the converted NodeApply and any ConversionIssue encountered.
+        tuple[NodeApply | EdgeApply, ConversionIssue]: A tuple containing the converted NodeApply and any ConversionIssue encountered.
     """
-    cache = DirectRelationCache(asset=asset_instance_id_by_id, source=source_instance_id_by_external_id)
-    resource_type = _lookup_resource_type(type(resource))
+    cache = DirectRelationCache(
+        asset=asset_instance_id_by_id, source=source_instance_id_by_external_id, file=file_instance_id_by_id
+    )
+    resource_type = _lookup_resource_type(resource)
     dumped = resource.dump()
     try:
         id_ = dumped.pop("id")
@@ -117,37 +164,54 @@ def asset_centric_to_dm(
     sources: list[NodeOrEdgeData] = []
     if properties:
         sources.append(NodeOrEdgeData(source=view_source.view_id, properties=properties))
-    instance_source_properties = {
-        "resourceType": resource_type,
-        "id": id_,
-        "dataSetId": data_set_id,
-        "classicExternalId": external_id,
-    }
-    sources.append(NodeOrEdgeData(source=INSTANCE_SOURCE_VIEW_ID, properties=instance_source_properties))
 
-    node = NodeApply(
-        space=instance_id.space,
-        external_id=instance_id.external_id,
-        sources=sources,
-    )
+    if resource_type != "fileAnnotation":
+        instance_source_properties = {
+            "resourceType": resource_type,
+            "id": id_,
+            "dataSetId": data_set_id,
+            "classicExternalId": external_id,
+        }
+        sources.append(NodeOrEdgeData(source=INSTANCE_SOURCE_VIEW_ID, properties=instance_source_properties))
 
-    return node, issue
+    instance: NodeApply | EdgeApply
+    if isinstance(instance_id, EdgeId):
+        edge_properties = create_edge_properties(
+            dumped, view_source.property_mapping, resource_type, issue, cache, instance_id.space
+        )
+        if any(key not in edge_properties for key in ("start_node", "end_node", "type")):
+            # Failed conversion of edge properties
+            return None, issue
+        instance = EdgeApply(
+            space=instance_id.space,
+            external_id=instance_id.external_id,
+            sources=sources,
+            **edge_properties,  # type: ignore[arg-type]
+        )
+    elif isinstance(instance_id, NodeId):
+        instance = NodeApply(space=instance_id.space, external_id=instance_id.external_id, sources=sources)
+    else:
+        raise RuntimeError(f"Unexpected instance_id type {type(instance_id)}")
 
+    return instance, issue
 
-def _lookup_resource_type(resource_type: type[CogniteResource]) -> AssetCentricType:
-    resource_type_map: dict[type[CogniteResource], AssetCentricType] = {
-        Asset: "asset",
-        FileMetadata: "file",
-        Event: "event",
-        TimeSeries: "timeseries",
-        Sequence: "sequence",
-        ExtendedFileMetadata: "file",
-        ExtendedTimeSeries: "timeseries",
-    }
-    try:
-        return resource_type_map[resource_type]
-    except KeyError as e:
-        raise ValueError(f"Unsupported resource type: {resource_type}") from e
+
+def _lookup_resource_type(resource_type: AssetCentricResourceExtended) -> AssetCentricType:
+    if isinstance(resource_type, Asset):
+        return "asset"
+    elif isinstance(resource_type, FileMetadata):
+        return "file"
+    elif isinstance(resource_type, Event):
+        return "event"
+    elif isinstance(resource_type, TimeSeries):
+        return "timeseries"
+    elif isinstance(resource_type, Annotation):
+        if resource_type.annotated_resource_type == "file" and resource_type.annotation_type in (
+            "diagrams.AssetLink",
+            "diagrams.FileLink",
+        ):
+            return "fileAnnotation"
+    raise ValueError(f"Unsupported resource type: {resource_type}")
 
 
 def create_properties(
@@ -210,5 +274,58 @@ def create_properties(
         (set(flatten_dump.keys()) - set(property_mapping.keys())) | ignored_asset_centric_properties
     )
     issue.missing_asset_centric_properties = sorted(set(property_mapping.keys()) - set(flatten_dump.keys()))
-    issue.missing_instance_properties = sorted(set(property_mapping.values()) - set(view_properties.keys()))
+    # Node and edge properties are handled separately
+    issue.missing_instance_properties = sorted(
+        {
+            prop_id
+            for prop_id in property_mapping.values()
+            if not (prop_id.startswith("edge.") or prop_id.startswith("node."))
+        }
+        - set(view_properties.keys())
+    )
     return properties
+
+
+def create_edge_properties(
+    dumped: dict[str, Any],
+    property_mapping: dict[str, str],
+    resource_type: AssetCentricType,
+    issue: ConversionIssue,
+    direct_relation_cache: DirectRelationCache,
+    default_instance_space: str,
+) -> dict[str, DirectRelationReference]:
+    flatten_dump = flatten_dict_json_path(dumped)
+    edge_properties: dict[str, DirectRelationReference] = {}
+    for prop_json_path, prop_id in property_mapping.items():
+        if not prop_id.startswith("edge."):
+            continue
+        if prop_json_path not in flatten_dump:
+            continue
+        edge_prop_id = prop_id.removeprefix("edge.")
+        if edge_prop_id in ("startNode", "endNode", "type"):
+            # DirectRelation lookup.
+            try:
+                value = convert_to_primary_property(
+                    flatten_dump[prop_json_path],
+                    DirectRelation(),
+                    False,
+                    direct_relation_lookup=direct_relation_cache.get(resource_type, prop_json_path),
+                )
+            except (ValueError, TypeError, NotImplementedError) as e:
+                issue.failed_conversions.append(
+                    FailedConversion(property_id=prop_json_path, value=flatten_dump[prop_json_path], error=str(e))
+                )
+                continue
+        elif edge_prop_id.endswith(".externalId"):
+            # Just an external ID string.
+            edge_prop_id = edge_prop_id.removesuffix(".externalId")
+            value = DirectRelationReference(default_instance_space, str(flatten_dump[prop_json_path]))
+        else:
+            issue.invalid_instance_property_types.append(
+                InvalidPropertyDataType(property_id=prop_id, expected_type="EdgeProperty")
+            )
+            continue
+        # We know that value is DirectRelationReference here
+        edge_properties[edge_prop_id.replace("Node", "_node")] = value  # type: ignore[assignment]
+
+    return edge_properties

cognite_toolkit/_cdf_tk/commands/_migrate/data_classes.py

@@ -3,7 +3,6 @@ from pathlib import Path
 from typing import Any, Generic, Literal
 
 from cognite.client.data_classes._base import (
-    T_WritableCogniteResource,
     WriteableCogniteResource,
     WriteableCogniteResourceList,
 )
@@ -15,11 +14,14 @@ from cognite_toolkit._cdf_tk.client.data_classes.instances import InstanceApplyL
 from cognite_toolkit._cdf_tk.client.data_classes.migration import AssetCentricId
 from cognite_toolkit._cdf_tk.client.data_classes.pending_instances_ids import PendingInstanceId
 from cognite_toolkit._cdf_tk.commands._migrate.default_mappings import create_default_mappings
-from cognite_toolkit._cdf_tk.exceptions import (
-    ToolkitValueError,
-)
+from cognite_toolkit._cdf_tk.exceptions import ToolkitValueError
 from cognite_toolkit._cdf_tk.storageio._data_classes import ModelList
-from cognite_toolkit._cdf_tk.utils.useful_types import AssetCentricKind, AssetCentricType, JsonVal
+from cognite_toolkit._cdf_tk.utils.useful_types import (
+    AssetCentricKind,
+    AssetCentricType,
+    JsonVal,
+    T_AssetCentricResource,
+)
 
 
 class MigrationMapping(BaseModel, alias_generator=to_camel_case, extra="ignore", populate_by_name=True):
@@ -186,9 +188,9 @@ class TimeSeriesMigrationMappingList(MigrationMappingList):
 
 
 @dataclass
-class AssetCentricMapping(Generic[T_WritableCogniteResource], WriteableCogniteResource[InstanceApply]):
+class AssetCentricMapping(Generic[T_AssetCentricResource], WriteableCogniteResource[InstanceApply]):
     mapping: MigrationMapping
-    resource: T_WritableCogniteResource
+    resource: T_AssetCentricResource
 
     def as_write(self) -> InstanceApply:
         raise NotImplementedError()
@@ -203,9 +205,7 @@ class AssetCentricMapping(Generic[T_WritableCogniteResource], WriteableCogniteRe
         }
 
 
-class AssetCentricMappingList(
-    WriteableCogniteResourceList[InstanceApply, AssetCentricMapping[T_WritableCogniteResource]]
-):
+class AssetCentricMappingList(WriteableCogniteResourceList[InstanceApply, AssetCentricMapping[T_AssetCentricResource]]):
     _RESOURCE: type = AssetCentricMapping
 
     def as_write(self) -> InstanceApplyList:

cognite_toolkit/_cdf_tk/commands/_migrate/data_mapper.py

@@ -16,6 +16,7 @@ from cognite_toolkit._cdf_tk.constants import MISSING_INSTANCE_SPACE
 from cognite_toolkit._cdf_tk.exceptions import ToolkitValueError
 from cognite_toolkit._cdf_tk.storageio._base import T_Selector, T_WriteCogniteResource
 from cognite_toolkit._cdf_tk.utils import humanize_collection
+from cognite_toolkit._cdf_tk.utils.useful_types import T_AssetCentricResource
 
 
 class DataMapper(Generic[T_Selector, T_CogniteResource, T_WriteCogniteResource], ABC):
@@ -30,7 +31,7 @@ class DataMapper(Generic[T_Selector, T_CogniteResource, T_WriteCogniteResource],
         pass
 
     @abstractmethod
-    def map(self, source: T_CogniteResource) -> tuple[T_WriteCogniteResource, MigrationIssue]:
+    def map(self, source: T_CogniteResource) -> tuple[T_WriteCogniteResource | None, MigrationIssue]:
         """Map a chunk of source data to the target format.
 
         Args:
@@ -43,7 +44,9 @@ class DataMapper(Generic[T_Selector, T_CogniteResource, T_WriteCogniteResource],
         raise NotImplementedError("Subclasses must implement this method.")
 
 
-class AssetCentricMapper(DataMapper[AssetCentricMigrationSelector, AssetCentricMapping, InstanceApply]):
+class AssetCentricMapper(
+    DataMapper[AssetCentricMigrationSelector, AssetCentricMapping[T_AssetCentricResource], InstanceApply]
+):
     def __init__(self, client: ToolkitClient) -> None:
         self.client = client
         self._ingestion_view_by_id: dict[ViewId, View] = {}
@@ -84,7 +87,7 @@ class AssetCentricMapper(DataMapper[AssetCentricMigrationSelector, AssetCentricM
         asset_mappings = self.client.migration.instance_source.list(resource_type="asset", limit=-1)
         self._asset_mapping_by_id = {mapping.id_: mapping.as_direct_relation_reference() for mapping in asset_mappings}
 
-    def map(self, source: AssetCentricMapping) -> tuple[InstanceApply, ConversionIssue]:
+    def map(self, source: AssetCentricMapping[T_AssetCentricResource]) -> tuple[InstanceApply | None, ConversionIssue]:
         """Map a chunk of asset-centric data to InstanceApplyList format."""
         mapping = source.mapping
         ingestion_view = mapping.get_ingestion_view()
@@ -102,6 +105,7 @@ class AssetCentricMapper(DataMapper[AssetCentricMigrationSelector, AssetCentricM
             view_properties=view_properties,
             asset_instance_id_by_id=self._asset_mapping_by_id,
             source_instance_id_by_external_id=self._source_system_mapping_by_id,
+            file_instance_id_by_id={},  # Todo implement file direct relations
         )
         if mapping.instance_id.space == MISSING_INSTANCE_SPACE:
             conversion_issue.missing_instance_space = f"Missing instance space for dataset ID {mapping.data_set_id!r}"

cognite_toolkit/_cdf_tk/commands/_migrate/migration_io.py

@@ -1,9 +1,6 @@
 from collections.abc import Iterator, Mapping, Sequence
 from typing import ClassVar, cast
 
-from cognite.client.data_classes._base import (
-    T_WritableCogniteResource,
-)
 from cognite.client.data_classes.data_modeling import InstanceApply, NodeId
 
 from cognite_toolkit._cdf_tk.client import ToolkitClient
@@ -22,6 +19,7 @@ from cognite_toolkit._cdf_tk.utils.useful_types import (
     AssetCentricKind,
     AssetCentricType,
     JsonVal,
+    T_AssetCentricResource,
 )
 
 from .data_classes import AssetCentricMapping, AssetCentricMappingList, MigrationMapping, MigrationMappingList
@@ -30,7 +28,7 @@ from .selectors import AssetCentricMigrationSelector, MigrateDataSetSelector, Mi
 
 
 class AssetCentricMigrationIO(
-    UploadableStorageIO[AssetCentricMigrationSelector, AssetCentricMapping[T_WritableCogniteResource], InstanceApply]
+    UploadableStorageIO[AssetCentricMigrationSelector, AssetCentricMapping[T_AssetCentricResource], InstanceApply]
 ):
     KIND = "AssetCentricMigration"
     SUPPORTED_DOWNLOAD_FORMATS = frozenset({".parquet", ".csv", ".ndjson"})
@@ -63,11 +61,11 @@ class AssetCentricMigrationIO(
 
     def _stream_from_csv(
         self, selector: MigrationCSVFileSelector, limit: int | None = None
-    ) -> Iterator[Sequence[AssetCentricMapping[T_WritableCogniteResource]]]:
+    ) -> Iterator[Sequence[AssetCentricMapping[T_AssetCentricResource]]]:
         items = selector.items
         if limit is not None:
             items = MigrationMappingList(items[:limit])
-        chunk: list[AssetCentricMapping[T_WritableCogniteResource]] = []
+        chunk: list[AssetCentricMapping[T_AssetCentricResource]] = []
         for current_batch in chunker_sequence(items, self.CHUNK_SIZE):
             resources = self.hierarchy.get_resource_io(selector.kind).retrieve(current_batch.get_ids())
             for mapping, resource in zip(current_batch, resources, strict=True):
@@ -86,12 +84,12 @@ class AssetCentricMigrationIO(
 
     def _stream_given_dataset(
         self, selector: MigrateDataSetSelector, limit: int | None = None
-    ) -> Iterator[Sequence[AssetCentricMapping[T_WritableCogniteResource]]]:
+    ) -> Iterator[Sequence[AssetCentricMapping[T_AssetCentricResource]]]:
         asset_centric_selector = selector.as_asset_centric_selector()
         for data_chunk in self.hierarchy.stream_data(asset_centric_selector, limit):
-            mapping_list = AssetCentricMappingList[T_WritableCogniteResource]([])
+            mapping_list = AssetCentricMappingList[T_AssetCentricResource]([])
             for resource in data_chunk.items:
-                # We know data_set_id is here as we are using a DataSetSelector
+                # We got the resource from a dataset selector, so we know it is there
                 data_set_id = cast(int, resource.data_set_id)
                 space_source = self.client.migration.space_source.retrieve(data_set_id=data_set_id)
                 instance_space = space_source.instance_space if space_source else None
@@ -129,7 +127,7 @@ class AssetCentricMigrationIO(
 
     def data_to_json_chunk(
         self,
-        data_chunk: Sequence[AssetCentricMapping[T_WritableCogniteResource]],
+        data_chunk: Sequence[AssetCentricMapping[T_AssetCentricResource]],
         selector: AssetCentricMigrationSelector | None = None,
     ) -> list[dict[str, JsonVal]]:
         return [item.dump() for item in data_chunk]

cognite_toolkit/_cdf_tk/storageio/__init__.py

@@ -3,6 +3,7 @@ from pathlib import Path
 from cognite_toolkit._cdf_tk.utils._auxiliary import get_concrete_subclasses
 from cognite_toolkit._cdf_tk.utils.fileio import COMPRESSION_BY_SUFFIX
 
+from ._annotations import FileAnnotationIO
 from ._applications import CanvasIO, ChartIO
 from ._asset_centric import AssetIO, BaseAssetCentricIO, EventIO, FileMetadataIO, HierarchyIO, TimeSeriesIO
 from ._base import (
@@ -50,6 +51,7 @@ __all__ = [
     "ChartIO",
     "ConfigurableStorageIO",
     "EventIO",
+    "FileAnnotationIO",
     "FileMetadataIO",
     "HierarchyIO",
     "InstanceIO",

cognite_toolkit/_cdf_tk/storageio/_annotations.py

@@ -0,0 +1,102 @@
+from collections.abc import Iterable, Sequence
+from typing import Any
+
+from cognite.client.data_classes import Annotation, AnnotationFilter
+
+from cognite_toolkit._cdf_tk.utils.collection import chunker_sequence
+from cognite_toolkit._cdf_tk.utils.useful_types import JsonVal
+
+from ._asset_centric import FileMetadataIO
+from ._base import Page, StorageIO
+from .selectors import AssetCentricSelector
+
+
+class FileAnnotationIO(StorageIO[AssetCentricSelector, Annotation]):
+    SUPPORTED_DOWNLOAD_FORMATS = frozenset({".ndjson"})
+    SUPPORTED_COMPRESSIONS = frozenset({".gz"})
+    CHUNK_SIZE = 1000
+    BASE_SELECTOR = AssetCentricSelector
+
+    MISSING_ID = "<MISSING_RESOURCE_ID>"
+
+    def as_id(self, item: Annotation) -> str:
+        project = item._cognite_client.config.project
+        return f"INTERNAL_ID_project_{project}_{item.id!s}"
+
+    def stream_data(self, selector: AssetCentricSelector, limit: int | None = None) -> Iterable[Page]:
+        total = 0
+        for file_chunk in FileMetadataIO(self.client).stream_data(selector, None):
+            # Todo Support pagination. This is missing in the SDK.
+            results = self.client.annotations.list(
+                filter=AnnotationFilter(
+                    annotated_resource_type="file",
+                    annotated_resource_ids=[{"id": file_metadata.id} for file_metadata in file_chunk.items],
+                )
+            )
+            if limit is not None and total + len(results) > limit:
+                results = results[: limit - total]
+
+            for chunk in chunker_sequence(results, self.CHUNK_SIZE):
+                yield Page(worker_id="main", items=chunk)
+                total += len(chunk)
+            if limit is not None and total >= limit:
+                break
+
+    def count(self, selector: AssetCentricSelector) -> int | None:
+        """There is no efficient way to count annotations in CDF."""
+        return None
+
+    def data_to_json_chunk(
+        self, data_chunk: Sequence[Annotation], selector: AssetCentricSelector | None = None
+    ) -> list[dict[str, JsonVal]]:
+        files_ids: set[int] = set()
+        for item in data_chunk:
+            if item.annotated_resource_type == "file" and item.annotated_resource_id is not None:
+                files_ids.add(item.annotated_resource_id)
+            if file_id := self._get_file_id(item.data):
+                files_ids.add(file_id)
+        self.client.lookup.files.external_id(list(files_ids))  # Preload file external IDs
+        asset_ids = {asset_id for item in data_chunk if (asset_id := self._get_asset_id(item.data))}
+        self.client.lookup.assets.external_id(list(asset_ids))  # Preload asset external IDs
+        return [self.dump_annotation_to_json(item) for item in data_chunk]
+
+    def dump_annotation_to_json(self, annotation: Annotation) -> dict[str, JsonVal]:
+        """Dump annotations to a list of JSON serializable dictionaries.
+
+        Args:
+            annotation: The annotations to dump.
+
+        Returns:
+            A list of JSON serializable dictionaries representing the annotations.
+        """
+        dumped = annotation.as_write().dump()
+        if isinstance(annotated_resource_id := dumped.pop("annotatedResourceId", None), int):
+            external_id = self.client.lookup.files.external_id(annotated_resource_id)
+            dumped["annotatedResourceExternalId"] = self.MISSING_ID if external_id is None else external_id
+
+        if isinstance(data := dumped.get("data"), dict):
+            if isinstance(file_ref := data.get("fileRef"), dict) and isinstance(file_ref.get("id"), int):
+                external_id = self.client.lookup.files.external_id(file_ref.pop("id"))
+                file_ref["externalId"] = self.MISSING_ID if external_id is None else external_id
+            if isinstance(asset_ref := data.get("assetRef"), dict) and isinstance(asset_ref.get("id"), int):
+                external_id = self.client.lookup.assets.external_id(asset_ref.pop("id"))
+                asset_ref["externalId"] = self.MISSING_ID if external_id is None else external_id
+        return dumped
+
+    @classmethod
+    def _get_file_id(cls, data: dict[str, Any]) -> int | None:
+        file_ref = data.get("fileRef")
+        if isinstance(file_ref, dict):
+            id_ = file_ref.get("id")
+            if isinstance(id_, int):
+                return id_
+        return None
+
+    @classmethod
+    def _get_asset_id(cls, data: dict[str, Any]) -> int | None:
+        asset_ref = data.get("assetRef")
+        if isinstance(asset_ref, dict):
+            id_ = asset_ref.get("id")
+            if isinstance(id_, int):
+                return id_
+        return None

cognite_toolkit/_cdf_tk/tracker.py

@@ -58,7 +58,7 @@ class Tracker:
             warning_details[f"warningMostCommon{no}Count"] = count
             warning_details[f"warningMostCommon{no}Name"] = warning
 
-        positional_args, optional_args = self._parse_sys_args()
+        subcommands, optional_args = self._parse_sys_args()
         event_information = {
             "userInput": self.user_command,
             "toolkitVersion": __version__,
@@ -69,7 +69,7 @@ class Tracker:
             **warning_details,
             "result": type(result).__name__ if isinstance(result, Exception) else result,
             "error": str(result) if isinstance(result, Exception) else "",
-            **positional_args,
+            "subcommands": subcommands,
             **optional_args,
             "alphaFlags": [name for name, value in self._cdf_toml.alpha_flags.items() if value],
             "plugins": [name for name, value in self._cdf_toml.plugins.items() if value],
@@ -128,9 +128,9 @@ class Tracker:
         return distinct_id
 
     @staticmethod
-    def _parse_sys_args() -> tuple[dict[str, str], dict[str, str | bool]]:
+    def _parse_sys_args() -> tuple[list[str], dict[str, str | bool]]:
         optional_args: dict[str, str | bool] = {}
-        positional_args: dict[str, str] = {}
+        subcommands: list[str] = []
         last_key: str | None = None
         if sys.argv and len(sys.argv) > 1:
             for arg in sys.argv[1:]:
@@ -147,11 +147,11 @@ class Tracker:
                     optional_args[last_key] = arg
                     last_key = None
                 else:
-                    positional_args[f"positionalArg{len(positional_args)}"] = arg
+                    subcommands.append(arg)
 
             if last_key:
                 optional_args[last_key] = True
-        return positional_args, optional_args
+        return subcommands, optional_args
 
     @property
     def _cicd(self) -> str:

cognite_toolkit/_cdf_tk/utils/fileio/_readers.py

@@ -7,6 +7,7 @@ from dataclasses import dataclass
 from functools import partial
 from io import TextIOWrapper
 from pathlib import Path
+from typing import Any
 
 import yaml
 
@@ -87,26 +88,20 @@ class FailedParsing:
     error: str
 
 
-class TableReader(FileReader, ABC): ...
-
-
-class CSVReader(TableReader):
-    """Reads CSV files and yields each row as a dictionary.
+class TableReader(FileReader, ABC):
+    """Reads table-like files and yields each row as a dictionary.
 
     Args:
-        input_file (Path): The path to the CSV file to read.
+        input_file (Path): The path to the table file to read.
         sniff_rows (int | None): Optional number of rows to sniff for
             schema detection. If None, no schema is detected. If a schema is sniffed
-            from the first `sniff_rows` rows, it will be used to parse the CSV.
+            from the first `sniff_rows` rows, it will be used to parse the table.
         schema (Sequence[SchemaColumn] | None): Optional schema to use for parsing.
             You can either provide a schema or use `sniff_rows` to detect it.
         keep_failed_cells (bool): If True, failed cells will be kept in the
             `failed_cell` attribute. If False, they will be ignored.
-
     """
 
-    format = ".csv"
-
     def __init__(
         self,
         input_file: Path,
@@ -152,18 +147,19 @@ class CSVReader(TableReader):
     @classmethod
     def sniff_schema(cls, input_file: Path, sniff_rows: int = 100) -> list[SchemaColumn]:
         """
-        Sniff the schema from the first `sniff_rows` rows of the CSV file.
+        Sniff the schema from the first `sniff_rows` rows of the file.
 
         Args:
-            input_file (Path): The path to the CSV file.
+            input_file (Path): The path to the tabular file.
             sniff_rows (int): The number of rows to read for sniffing the schema.
 
         Returns:
             list[SchemaColumn]: The inferred schema as a list of SchemaColumn objects.
+
         Raises:
             ValueError: If `sniff_rows` is not a positive integer.
             ToolkitFileNotFoundError: If the file does not exist.
-            ToolkitValueError: If the file is not a CSV file or if there are issues with the content.
+            ToolkitValueError: If the file is not the correct format or if there are issues with the content.
 
         """
         if sniff_rows <= 0:
@@ -171,43 +167,50 @@ class CSVReader(TableReader):
 
         if not input_file.exists():
             raise ToolkitFileNotFoundError(f"File not found: {input_file.as_posix()!r}.")
-        if input_file.suffix != ".csv":
-            raise ToolkitValueError(f"Expected a .csv file got a {input_file.suffix!r} file instead.")
+        if input_file.suffix != cls.format:
+            raise ToolkitValueError(f"Expected a {cls.format} file got a {input_file.suffix!r} file instead.")
 
-        with input_file.open("r", encoding="utf-8-sig") as file:
-            reader = csv.DictReader(file)
-            column_names = Counter(reader.fieldnames)
-            if duplicated := [name for name, count in column_names.items() if count > 1]:
-                raise ToolkitValueError(f"CSV file contains duplicate headers: {humanize_collection(duplicated)}")
-            sample_rows: list[dict[str, str]] = []
-            for no, row in enumerate(reader):
-                if no >= sniff_rows:
-                    break
-                sample_rows.append(row)
+        column_names, sample_rows = cls._read_sample_rows(input_file, sniff_rows)
+        cls._check_column_names(column_names)
+        return cls._infer_schema(sample_rows, column_names)
 
-            if not sample_rows:
-                raise ToolkitValueError(f"No data found in the file: {input_file.as_posix()!r}.")
+    @classmethod
+    @abstractmethod
+    def _read_sample_rows(cls, input_file: Path, sniff_rows: int) -> tuple[Sequence[str], list[dict[str, str]]]: ...
 
-            schema = []
-            for column_name in reader.fieldnames or []:
-                sample_values = [row[column_name] for row in sample_rows if column_name in row]
-                if not sample_values:
-                    column = SchemaColumn(name=column_name, type="string")
+    @classmethod
+    def _infer_schema(cls, sample_rows: list[dict[str, Any]], column_names: Sequence[str]) -> list[SchemaColumn]:
+        schema: list[SchemaColumn] = []
+        for column_name in column_names:
+            sample_values = [row[column_name] for row in sample_rows if column_name in row]
+            if not sample_values:
+                column = SchemaColumn(name=column_name, type="string")
+            else:
+                data_types = Counter(
+                    infer_data_type_from_value(value, dtype="Json")[0] for value in sample_values if value is not None
+                )
+                if not data_types:
+                    inferred_type = "string"
                 else:
-                    data_types = Counter(
-                        infer_data_type_from_value(value, dtype="Json")[0]
-                        for value in sample_values
-                        if value is not None
-                    )
-                    if not data_types:
-                        inferred_type = "string"
-                    else:
-                        inferred_type = data_types.most_common()[0][0]
-                    # Json dtype is a subset of Datatype that SchemaColumn accepts
-                    column = SchemaColumn(name=column_name, type=inferred_type)  # type: ignore[arg-type]
-                schema.append(column)
+                    inferred_type = data_types.most_common()[0][0]
+                # Json dtype is a subset of Datatype that SchemaColumn accepts
+                column = SchemaColumn(name=column_name, type=inferred_type)  # type: ignore[arg-type]
+            schema.append(column)
         return schema
 
+    @classmethod
+    def _check_column_names(cls, column_names: Sequence[str]) -> None:
+        """Check for duplicate column names."""
+        duplicates = [name for name, count in Counter(column_names).items() if count > 1]
+        if duplicates:
+            raise ToolkitValueError(f"Duplicate column names found: {humanize_collection(duplicates)}.")
+
+
+class CSVReader(TableReader):
+    """Reads CSV files and yields each row as a dictionary."""
+
+    format = ".csv"
+
     def _read_chunks_from_file(self, file: TextIOWrapper) -> Iterator[dict[str, JsonVal]]:
         if self.keep_failed_cells and self.failed_cell:
             self.failed_cell.clear()
@@ -231,10 +234,31 @@ class CSVReader(TableReader):
         with compression.open("r") as file:
             yield from csv.DictReader(file)
 
+    @classmethod
+    def _read_sample_rows(cls, input_file: Path, sniff_rows: int) -> tuple[Sequence[str], list[dict[str, str]]]:
+        column_names: Sequence[str] = []
+        compression = Compression.from_filepath(input_file)
+        with compression.open("r") as file:
+            reader = csv.DictReader(file)
+            column_names = reader.fieldnames or []
+            sample_rows: list[dict[str, str]] = []
+            for no, row in enumerate(reader):
+                if no >= sniff_rows:
+                    break
+                sample_rows.append(row)
+
+            if not sample_rows:
+                raise ToolkitValueError(f"No data found in the file: {input_file.as_posix()!r}.")
+        return column_names, sample_rows
+
 
 class ParquetReader(TableReader):
     format = ".parquet"
 
+    def __init__(self, input_file: Path) -> None:
+        # Parquet files have their own schema, so we don't need to sniff or provide one.
+        super().__init__(input_file, sniff_rows=None, schema=None, keep_failed_cells=False)
+
     def read_chunks(self) -> Iterator[dict[str, JsonVal]]:
         import pyarrow.parquet as pq
 
@@ -258,6 +282,28 @@ class ParquetReader(TableReader):
                 return value
         return value
 
+    @classmethod
+    def _read_sample_rows(cls, input_file: Path, sniff_rows: int) -> tuple[Sequence[str], list[dict[str, str]]]:
+        import pyarrow.parquet as pq
+
+        column_names: Sequence[str] = []
+        sample_rows: list[dict[str, str]] = []
+        with pq.ParquetFile(input_file) as parquet_file:
+            column_names = parquet_file.schema.names
+            row_count = min(sniff_rows, parquet_file.metadata.num_rows)
+            row_iter = parquet_file.iter_batches(batch_size=row_count)
+            try:
+                batch = next(row_iter)
+                for row in batch.to_pylist():
+                    str_row = {key: (str(value) if value is not None else "") for key, value in row.items()}
+                    sample_rows.append(str_row)
+            except StopIteration:
+                pass
+
+            if not sample_rows:
+                raise ToolkitValueError(f"No data found in the file: {input_file.as_posix()!r}.")
+        return column_names, sample_rows
+
 
 FILE_READ_CLS_BY_FORMAT: Mapping[str, type[FileReader]] = {}
 TABLE_READ_CLS_BY_FORMAT: Mapping[str, type[TableReader]] = {}

cognite_toolkit/_cdf_tk/utils/http_client/_client.py

@@ -147,13 +147,15 @@ class HTTPClient:
             timeout=self.config.timeout,
         )
 
-    def _create_headers(self, api_version: str | None = None) -> MutableMapping[str, str]:
+    def _create_headers(
+        self, api_version: str | None = None, content_type: str = "application/json", accept: str = "application/json"
+    ) -> MutableMapping[str, str]:
         headers: MutableMapping[str, str] = {}
         headers["User-Agent"] = f"httpx/{httpx.__version__} {get_user_agent()}"
         auth_name, auth_value = self.config.credentials.authorization_header()
         headers[auth_name] = auth_value
-        headers["content-type"] = "application/json"
-        headers["accept"] = "application/json"
+        headers["content-type"] = content_type
+        headers["accept"] = accept
         headers["x-cdp-sdk"] = f"CogniteToolkit:{get_current_toolkit_version()}"
         headers["x-cdp-app"] = self.config.client_name
         headers["cdf-version"] = api_version or self.config.api_subversion
@@ -162,7 +164,7 @@ class HTTPClient:
         return headers
 
     def _make_request(self, item: RequestMessage) -> httpx.Response:
-        headers = self._create_headers(item.api_version)
+        headers = self._create_headers(item.api_version, item.content_type, item.accept)
         params: dict[str, PrimitiveType] | None = None
         if isinstance(item, ParamRequest):
             params = item.parameters

cognite_toolkit/_cdf_tk/utils/http_client/_data_classes.py

@@ -92,6 +92,8 @@ class RequestMessage(HTTPMessage):
     read_attempt: int = 0
     status_attempt: int = 0
     api_version: str | None = None
+    content_type: str = "application/json"
+    accept: str = "application/json"
 
     @property
     def total_attempts(self) -> int:

cognite_toolkit/_cdf_tk/utils/useful_types.py

@@ -2,15 +2,17 @@ from collections.abc import Hashable
 from datetime import date, datetime
 from typing import Any, Literal, TypeAlias, TypeVar, get_args
 
-from cognite.client.data_classes import Asset, Event, FileMetadata, Sequence, TimeSeries
+from cognite.client.data_classes import Annotation, Asset, Event, FileMetadata, TimeSeries
 from cognite.client.data_classes._base import CogniteObject, WriteableCogniteResourceList
 
 JsonVal: TypeAlias = None | str | int | float | bool | dict[str, "JsonVal"] | list["JsonVal"]
 
 AssetCentricDestinationType: TypeAlias = Literal["assets", "files", "events", "timeseries", "sequences"]
-AssetCentricType: TypeAlias = Literal["asset", "file", "event", "timeseries", "sequence"]
-AssetCentricResource: TypeAlias = Asset | FileMetadata | Event | TimeSeries | Sequence
-AssetCentricKind: TypeAlias = Literal["Assets", "Events", "TimeSeries", "FileMetadata"]
+AssetCentricType: TypeAlias = Literal["asset", "file", "event", "timeseries", "sequence", "fileAnnotation"]
+AssetCentricResource: TypeAlias = Asset | FileMetadata | Event | TimeSeries
+AssetCentricResourceExtended: TypeAlias = Asset | FileMetadata | Event | TimeSeries | Annotation
+AssetCentricKind: TypeAlias = Literal["Assets", "Events", "TimeSeries", "FileMetadata", "FileAnnotations"]
+
 DataType: TypeAlias = Literal["string", "integer", "float", "boolean", "json", "date", "timestamp"]
 PythonTypes: TypeAlias = str | int | float | bool | datetime | date | dict[str, Any] | list[Any]
 
@@ -22,3 +24,4 @@ T_Value = TypeVar("T_Value")
 PrimitiveType: TypeAlias = str | int | float | bool
 
 T_WriteCogniteResource = TypeVar("T_WriteCogniteResource", bound=CogniteObject)
+T_AssetCentricResource = TypeVar("T_AssetCentricResource", bound=AssetCentricResource)

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

@@ -12,7 +12,7 @@ jobs:
     environment: dev
     name: Deploy
     container:
-      image: cognite/toolkit:0.6.88
+      image: cognite/toolkit:0.6.89
       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.88
+      image: cognite/toolkit:0.6.89
       env:
         CDF_CLUSTER: ${{ vars.CDF_CLUSTER }}
         CDF_PROJECT: ${{ vars.CDF_PROJECT }}

cognite_toolkit/_resources/cdf.toml

@@ -4,7 +4,7 @@ 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.88"
+version = "0.6.89"
 
 [alpha_flags]
 external-libraries = true

cognite_toolkit/_version.py

@@ -1 +1 @@
-__version__ = "0.6.88"
+__version__ = "0.6.89"

{cognite_toolkit-0.6.88.dist-info → cognite_toolkit-0.6.89.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cognite_toolkit
-Version: 0.6.88
+Version: 0.6.89
 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

{cognite_toolkit-0.6.88.dist-info → cognite_toolkit-0.6.89.dist-info}/RECORD

@@ -1,6 +1,6 @@
 cognite_toolkit/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 cognite_toolkit/_cdf.py,sha256=1OSAvbOeuIrnsczEG2BtGqRP3L3sq0VMPthmugnqCUw,5821
-cognite_toolkit/_version.py,sha256=EgTgWpmRuWXg6fwc6D-LGpiAVF4TwDefD9KxLpZLYlU,23
+cognite_toolkit/_version.py,sha256=Zx4pn6mGv-GP7UL77bykGBRh05hpWmRUnqQ8PXKNLcc,23
 cognite_toolkit/_cdf_tk/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 cognite_toolkit/_cdf_tk/cdf_toml.py,sha256=IjmzNVLxsOV6tsMDgmJmXsy-LQru-8IEQdFzGW5DxVk,8117
 cognite_toolkit/_cdf_tk/constants.py,sha256=e9XmGvQCqGq7zYQrNoopU5e2KnYZYBPyUC5raGShK7k,6364
@@ -8,7 +8,7 @@ cognite_toolkit/_cdf_tk/exceptions.py,sha256=xG0jMwi5A20nvPvyo6sCyz_cyKycynPyIzp
 cognite_toolkit/_cdf_tk/feature_flags.py,sha256=oKvUHcNTtt8zp31eZ1eSCxfSIelm0L5B0xAQOskr1hc,2892
 cognite_toolkit/_cdf_tk/hints.py,sha256=UI1ymi2T5wCcYOpEbKbVaDnlyFReFy8TDtMVt-5E1h8,6493
 cognite_toolkit/_cdf_tk/plugins.py,sha256=yL7Q4k9UGnoHP9Ucrno02_qi1L3DrE6ggBiQI-wQKiU,783
-cognite_toolkit/_cdf_tk/tracker.py,sha256=_j8gWlqwfD0eCbRW4XKB1kAixjCwau9-o4rPvFvouNY,6016
+cognite_toolkit/_cdf_tk/tracker.py,sha256=ybazaYDMgrtmAaCEb1nlKAQzjcN352-U-om4NBGV328,5965
 cognite_toolkit/_cdf_tk/validation.py,sha256=KFdPgnNIbVM0yjFF0cqmpBB8MI8e-U-YbBYrP4IiClE,8441
 cognite_toolkit/_cdf_tk/apps/__init__.py,sha256=nNQymHhwxjXNpY9N9xDmnvSPLCMwQkn_t9oRkgDWofI,659
 cognite_toolkit/_cdf_tk/apps/_auth_app.py,sha256=ER7uYb3ViwsHMXiQEZpyhwU6TIjKaB9aEy32VI4MPpg,3397
@@ -115,17 +115,16 @@ cognite_toolkit/_cdf_tk/commands/pull.py,sha256=2Zf6IOXxSxZ-5XkNE80FlrXBuNejAWrA
 cognite_toolkit/_cdf_tk/commands/repo.py,sha256=MNy8MWphTklIZHvQOROCweq8_SYxGv6BaqnLpkFFnuk,3845
 cognite_toolkit/_cdf_tk/commands/run.py,sha256=JyX9jLEQej9eRrHVCCNlw4GuF80qETSol3-T5CCofgw,37331
 cognite_toolkit/_cdf_tk/commands/_migrate/__init__.py,sha256=i5ldcTah59K0E4fH5gHTV0GRvtDCEvVses9WQzn9Lno,226
-cognite_toolkit/_cdf_tk/commands/_migrate/base.py,sha256=aS32Wa-gd7vFNOdCTKEIpSSdjQKZ8jblzD6c3YcOmzA,4942
-cognite_toolkit/_cdf_tk/commands/_migrate/canvas.py,sha256=Tv4OG9V6tDsQbSH13YW8M0n8Ury5gU16oJB-OISVR0w,6398
-cognite_toolkit/_cdf_tk/commands/_migrate/command.py,sha256=RtC5CMQC1Wy5RmOwlCM67Mc8fZEExXubJLvzDzk2IMA,14097
-cognite_toolkit/_cdf_tk/commands/_migrate/conversion.py,sha256=Eol-0ruQ14fwS-bx2pEmbXdICodfknSJ_OsAASa6jkY,9592
+cognite_toolkit/_cdf_tk/commands/_migrate/canvas.py,sha256=R-z0yfOFcJZj-zRLhN-7z_-SLxqzSmONMgrbzNF9dGs,8843
+cognite_toolkit/_cdf_tk/commands/_migrate/command.py,sha256=059GfizsxZzsO7CJP3uEtyvfbeKt2F5eYnvq9GvHxNE,14212
+cognite_toolkit/_cdf_tk/commands/_migrate/conversion.py,sha256=ElsZeZNRPkxBm2VdTe0nCvyq75qfGLOLFY54fCuaxFY,14333
 cognite_toolkit/_cdf_tk/commands/_migrate/creators.py,sha256=FTu7w3G8KyPY8pagG3KdPpOmpLcjehaAg2auEy6iM7A,9605
-cognite_toolkit/_cdf_tk/commands/_migrate/data_classes.py,sha256=eF3Fv-tm_A8iaRUAOVOdg1DKPN7OYVPwrLQMvawsO14,8572
-cognite_toolkit/_cdf_tk/commands/_migrate/data_mapper.py,sha256=L9HJGmcTiAkC8Sftrngp8FyHBtAbkZtb9HrpLFg-ED0,6006
+cognite_toolkit/_cdf_tk/commands/_migrate/data_classes.py,sha256=S5MWfkaYAHfV6rc6RA4KR7WRETdcAKdpe6s4ZdwB0F0,8562
+cognite_toolkit/_cdf_tk/commands/_migrate/data_mapper.py,sha256=7m9uj_W11iokGmbWD979z5UK5KwOYv7L9HiY5PcpIN4,6231
 cognite_toolkit/_cdf_tk/commands/_migrate/data_model.py,sha256=i1eUsNX6Dueol9STIEwyksBnBsWUk13O8qHIjW964pM,7860
 cognite_toolkit/_cdf_tk/commands/_migrate/default_mappings.py,sha256=KkSq_4R6hQ15ccG-jHy7vVgPwC5IDd5OaXZLvz5mIZs,5547
 cognite_toolkit/_cdf_tk/commands/_migrate/issues.py,sha256=lWSnuS3CfRDbA7i1g12gJ2reJnQcLmZWxHDK19-Wxkk,5772
-cognite_toolkit/_cdf_tk/commands/_migrate/migration_io.py,sha256=J0Cb_TI_G5UoJIxpKGlCOLwjZ5VeDU93G_QP1-84lqU,9878
+cognite_toolkit/_cdf_tk/commands/_migrate/migration_io.py,sha256=obtdne1XqzOhlWuw2QnFs7vub5_FTaYOf086INB5d4I,9813
 cognite_toolkit/_cdf_tk/commands/_migrate/prepare.py,sha256=RfqaNoso5CyBwc-p6ckwcYqBfZXKhdJgdGIyd0TATaI,2635
 cognite_toolkit/_cdf_tk/commands/_migrate/selectors.py,sha256=CYle1Gz69HHnKF4onTIFxrpiOoDApvVK01SFuQuHzP0,2130
 cognite_toolkit/_cdf_tk/cruds/__init__.py,sha256=j0yxDCwd4Cl9KG7SvGTDQg4Y2bHfYVEDv8CBxbFTWUM,6070
@@ -226,7 +225,8 @@ cognite_toolkit/_cdf_tk/resource_classes/robotics/data_postprocessing.py,sha256=
 cognite_toolkit/_cdf_tk/resource_classes/robotics/frame.py,sha256=XmDqJ0pAxe_vAP0Dhktba1f9o2zg_ORCJ3Hz8cyJMrk,899
 cognite_toolkit/_cdf_tk/resource_classes/robotics/location.py,sha256=dbc9HT-bc2Qt15hHoR63SM7pg321BhNuTNjI7HHCwSA,468
 cognite_toolkit/_cdf_tk/resource_classes/robotics/map.py,sha256=j77z7CzCMiMj8r94BdUKCum9EuZRUjaSlUAy9K9DL_Q,942
-cognite_toolkit/_cdf_tk/storageio/__init__.py,sha256=aM-skaPnKTH1B7HG0faeTUNf7u1b-sT8l7hh5JRZ1E8,2288
+cognite_toolkit/_cdf_tk/storageio/__init__.py,sha256=-OhPPhl9z1ynYcFlRVYfOPD246HhQydmn4VzByJh7C0,2355
+cognite_toolkit/_cdf_tk/storageio/_annotations.py,sha256=wAMkgM-IpgXuY7_1KbtiTv8VdA555ywKjntD_boOBPk,4647
 cognite_toolkit/_cdf_tk/storageio/_applications.py,sha256=bhyG1d2_9duPkX-otC2brVcpChvdXSPkYhBHS5T_72g,4343
 cognite_toolkit/_cdf_tk/storageio/_asset_centric.py,sha256=Rhy64zUW4oxacq_vYomDeTRPmF6Vx-1mkYAFAqJE9vk,28312
 cognite_toolkit/_cdf_tk/storageio/_base.py,sha256=NWXPdgzUnpBiav5Hi8XGHkWU9QiMjNzBQTxMcuxF-LA,11017
@@ -268,16 +268,16 @@ cognite_toolkit/_cdf_tk/utils/sql_parser.py,sha256=jernu2amPQ54cQZ4vFZm1gEhFZfGc
 cognite_toolkit/_cdf_tk/utils/table_writers.py,sha256=Rxp_CZDDWrNPERNq6u1xsAX1OvzownwMMnpwNu8KdH0,17861
 cognite_toolkit/_cdf_tk/utils/text.py,sha256=EpIXjaQ5C5q5fjbUjAW7tncXpdJfiQeV7CYSbr70Bl0,3106
 cognite_toolkit/_cdf_tk/utils/thread_safe_dict.py,sha256=NbRHcZvWpF9xHP5OkOMGFpxrPNbi0Q3Eea6PUNbGlt4,3426
-cognite_toolkit/_cdf_tk/utils/useful_types.py,sha256=tPZOcK1PR1hNogMCgF863APMK6p3528t5kKaKbVl0-s,1352
+cognite_toolkit/_cdf_tk/utils/useful_types.py,sha256=ITiaJYcJrQZTR-CejtnSRPdseg86TxtWv8TqIkLi09c,1565
 cognite_toolkit/_cdf_tk/utils/validate_access.py,sha256=1puswcpgEDNCwdk91dhLqCBSu_aaUAd3Hsw21d-YVFs,21955
 cognite_toolkit/_cdf_tk/utils/fileio/__init__.py,sha256=ts5kYu_1Ks7xjnM6pIrVUrZe0nkYI6euYXeE4ox34xk,1199
 cognite_toolkit/_cdf_tk/utils/fileio/_base.py,sha256=MpWaD3lR9vrJ-kGzTiDOtChXhvFD7-xrP-Pzp7vjnLY,756
 cognite_toolkit/_cdf_tk/utils/fileio/_compression.py,sha256=8BAPgg5OKc3vkEEkqOvYsuyh12iXVNuEmC0omWwyJNQ,2355
-cognite_toolkit/_cdf_tk/utils/fileio/_readers.py,sha256=77Uq5r0pnD8gXLDVPzLT-1VPfzyOK7TZk85BtWxW7DQ,11613
+cognite_toolkit/_cdf_tk/utils/fileio/_readers.py,sha256=nGfsSfpXDYUncncsFuJD9-xYPJ5635mSFUJfuCyQ3no,13724
 cognite_toolkit/_cdf_tk/utils/fileio/_writers.py,sha256=ghNGBZjkISAlbxe8o5YWWloLXG9QKOtF_qGA9JkvYss,17712
 cognite_toolkit/_cdf_tk/utils/http_client/__init__.py,sha256=H1T-cyIoVaPL4MvN1IuG-cHgj-cqB7eszu2kIN939lw,813
-cognite_toolkit/_cdf_tk/utils/http_client/_client.py,sha256=zsN5eP1spgMkIRN6qeQ-djAohJVVuacpD2fnQh5QYx0,10916
-cognite_toolkit/_cdf_tk/utils/http_client/_data_classes.py,sha256=PZzcZb_1rnSBs0MeMaBhwwbIP1uyTg3fbRSiT8E3djE,12531
+cognite_toolkit/_cdf_tk/utils/http_client/_client.py,sha256=sN1Sizxv_rAcOSRd-2lKSs0p-SO1rA4eHL8pHMTBx54,11018
+cognite_toolkit/_cdf_tk/utils/http_client/_data_classes.py,sha256=7mr1goZKx5oR877EVu5MWu1qplm-lfCKCAMVZZ5Tvm0,12611
 cognite_toolkit/_cdf_tk/utils/http_client/_exception.py,sha256=fC9oW6BN0HbUe2AkYABMP7Kj0-9dNYXVFBY5RQztq2c,126
 cognite_toolkit/_cdf_tk/utils/http_client/_tracker.py,sha256=EBBnd-JZ7nc_jYNFJokCHN2UZ9sx0McFLZvlceUYYic,1215
 cognite_toolkit/_repo_files/.env.tmpl,sha256=UmgKZVvIp-OzD8oOcYuwb_6c7vSJsqkLhuFaiVgK7RI,972
@@ -285,13 +285,13 @@ cognite_toolkit/_repo_files/.gitignore,sha256=ip9kf9tcC5OguF4YF4JFEApnKYw0nG0vPi
 cognite_toolkit/_repo_files/AzureDevOps/.devops/README.md,sha256=OLA0D7yCX2tACpzvkA0IfkgQ4_swSd-OlJ1tYcTBpsA,240
 cognite_toolkit/_repo_files/AzureDevOps/.devops/deploy-pipeline.yml,sha256=brULcs8joAeBC_w_aoWjDDUHs3JheLMIR9ajPUK96nc,693
 cognite_toolkit/_repo_files/AzureDevOps/.devops/dry-run-pipeline.yml,sha256=OBFDhFWK1mlT4Dc6mDUE2Es834l8sAlYG50-5RxRtHk,723
-cognite_toolkit/_repo_files/GitHub/.github/workflows/deploy.yaml,sha256=3sYDml0HnE0_ryoemce1AIr0skHCm5i4SthSM_3BRDA,667
-cognite_toolkit/_repo_files/GitHub/.github/workflows/dry-run.yaml,sha256=2IMERh0DGL37Fs0X4P4wuiEN2y29gGqtDX704zitNi0,2430
-cognite_toolkit/_resources/cdf.toml,sha256=mpGJ2GDQJB63ZDiAsHxrRrrgpNjD-HLIUIf8zJRJjaw,487
+cognite_toolkit/_repo_files/GitHub/.github/workflows/deploy.yaml,sha256=4lnerKrn16H3_STBcvu5fODE5cVtxaZEloaTtvCpxXc,667
+cognite_toolkit/_repo_files/GitHub/.github/workflows/dry-run.yaml,sha256=ppSA3aVmcoucLy8pNBfEKscrIvsPgyFwnb0Wv09dLNY,2430
+cognite_toolkit/_resources/cdf.toml,sha256=YVocnddLa-tvLk7NRmUshNUS1v8EbDyHHYYJY4or2WU,487
 cognite_toolkit/demo/__init__.py,sha256=-m1JoUiwRhNCL18eJ6t7fZOL7RPfowhCuqhYFtLgrss,72
 cognite_toolkit/demo/_base.py,sha256=6xKBUQpXZXGQ3fJ5f7nj7oT0s2n7OTAGIa17ZlKHZ5U,8052
-cognite_toolkit-0.6.88.dist-info/METADATA,sha256=Ubje2tHMnCHTWlxk6Vdg2WJP8ocjbOGL4Tkoy-44QDY,4501
-cognite_toolkit-0.6.88.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-cognite_toolkit-0.6.88.dist-info/entry_points.txt,sha256=JlR7MH1_UMogC3QOyN4-1l36VbrCX9xUdQoHGkuJ6-4,83
-cognite_toolkit-0.6.88.dist-info/licenses/LICENSE,sha256=CW0DRcx5tL-pCxLEN7ts2S9g2sLRAsWgHVEX4SN9_Mc,752
-cognite_toolkit-0.6.88.dist-info/RECORD,,
+cognite_toolkit-0.6.89.dist-info/METADATA,sha256=BctpEwbWImcRv7ez5PS2z0R9qS4SmL8VRc4Ryv_GVoQ,4501
+cognite_toolkit-0.6.89.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+cognite_toolkit-0.6.89.dist-info/entry_points.txt,sha256=JlR7MH1_UMogC3QOyN4-1l36VbrCX9xUdQoHGkuJ6-4,83
+cognite_toolkit-0.6.89.dist-info/licenses/LICENSE,sha256=CW0DRcx5tL-pCxLEN7ts2S9g2sLRAsWgHVEX4SN9_Mc,752
+cognite_toolkit-0.6.89.dist-info/RECORD,,

cognite_toolkit/_cdf_tk/commands/_migrate/base.py

@@ -1,106 +0,0 @@
-from abc import ABC
-from typing import TypeVar
-
-from cognite.client.data_classes import Asset, Event, FileMetadata, Sequence, TimeSeries
-from cognite.client.data_classes.capabilities import (
-    Capability,
-    DataModelInstancesAcl,
-    DataModelsAcl,
-    SpaceIDScope,
-)
-
-from cognite_toolkit._cdf_tk.client import ToolkitClient
-from cognite_toolkit._cdf_tk.commands._base import ToolkitCommand
-from cognite_toolkit._cdf_tk.constants import DMS_INSTANCE_LIMIT_MARGIN
-from cognite_toolkit._cdf_tk.exceptions import (
-    AuthenticationError,
-    ToolkitMigrationError,
-    ToolkitValueError,
-)
-from cognite_toolkit._cdf_tk.utils import humanize_collection
-
-from .data_model import INSTANCE_SOURCE_VIEW_ID, MODEL_ID, RESOURCE_VIEW_MAPPING_VIEW_ID
-
-T_AssetCentricResource = TypeVar("T_AssetCentricResource", bound=Asset | Event | FileMetadata | TimeSeries | Sequence)
-
-
-class BaseMigrateCommand(ToolkitCommand, ABC):
-    def source_acl(self, data_set_id: list[int]) -> Capability:
-        """Return the source ACL for the given data set IDs."""
-        # This method should be implemented in subclasses that needs access to a specific source ACL.
-        # such as TimeSeries, Files, Assets, and so on.
-        raise ValueError(
-            "Bug in Toolkit: the source ACL is not defined for this migration command. "
-            "Please implement the source_acl method."
-        )
-
-    def validate_access(
-        self,
-        client: ToolkitClient,
-        instance_spaces: list[str] | None = None,
-        schema_spaces: list[str] | None = None,
-        data_set_ids: list[int] | None = None,
-    ) -> None:
-        required_capabilities: list[Capability] = []
-        if instance_spaces is not None:
-            required_capabilities.append(
-                DataModelInstancesAcl(
-                    actions=[
-                        DataModelInstancesAcl.Action.Read,
-                        DataModelInstancesAcl.Action.Write,
-                        DataModelInstancesAcl.Action.Write_Properties,
-                    ],
-                    scope=SpaceIDScope(instance_spaces),
-                )
-            )
-        if schema_spaces is not None:
-            required_capabilities.append(
-                DataModelsAcl(actions=[DataModelsAcl.Action.Read], scope=SpaceIDScope(schema_spaces)),
-            )
-
-        if data_set_ids is not None:
-            source_acl = self.source_acl(data_set_ids)
-            required_capabilities.append(source_acl)
-        if missing := client.iam.verify_capabilities(required_capabilities):
-            raise AuthenticationError(f"Missing required capabilities: {humanize_collection(missing)}.", missing)
-
-    @staticmethod
-    def validate_migration_model_available(client: ToolkitClient) -> None:
-        models = client.data_modeling.data_models.retrieve([MODEL_ID], inline_views=False)
-        if not models:
-            raise ToolkitMigrationError(
-                f"The migration data model {MODEL_ID!r} does not exist. "
-                "Please run the `cdf migrate prepare` command to deploy the migration data model."
-            )
-        elif len(models) > 1:
-            raise ToolkitMigrationError(
-                f"Multiple migration models {MODEL_ID!r}. "
-                "Please delete the duplicate models before proceeding with the migration."
-            )
-        model = models[0]
-        missing_views = {INSTANCE_SOURCE_VIEW_ID, RESOURCE_VIEW_MAPPING_VIEW_ID} - set(model.views or [])
-        if missing_views:
-            raise ToolkitMigrationError(
-                f"Invalid migration model. Missing views {humanize_collection(missing_views)}. "
-                f"Please run the `cdf migrate prepare` command to deploy the migration data model."
-            )
-
-    def validate_available_capacity(self, client: ToolkitClient, instance_count: int) -> None:
-        """Validate that the project has enough capacity to accommodate the migration."""
-
-        stats = client.data_modeling.statistics.project()
-
-        available_capacity = stats.instances.instances_limit - stats.instances.instances
-        available_capacity_after = available_capacity - instance_count
-
-        if available_capacity_after < DMS_INSTANCE_LIMIT_MARGIN:
-            raise ToolkitValueError(
-                "Cannot proceed with migration, not enough instance capacity available. Total capacity after migration"
-                f" would be {available_capacity_after:,} instances, which is less than the required margin of"
-                f" {DMS_INSTANCE_LIMIT_MARGIN:,} instances. Please increase the instance capacity in your CDF project"
-                f" or delete some existing instances before proceeding with the migration of {instance_count:,} assets."
-            )
-        total_instances = stats.instances.instances + instance_count
-        self.console(
-            f"Project has enough capacity for migration. Total instances after migration: {total_instances:,}."
-        )