cognite-toolkit 0.6.88__tar.gz → 0.6.89__tar.gz

{cognite_toolkit-0.6.88 → cognite_toolkit-0.6.89}/PKG-INFO

@@ -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 → cognite_toolkit-0.6.89}/cdf.toml

@@ -31,7 +31,7 @@ dump = true
 [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"
 
 
 [library.toolkit-data]

{cognite_toolkit-0.6.88 → cognite_toolkit-0.6.89}/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-0.6.88 → cognite_toolkit-0.6.89}/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-0.6.88 → cognite_toolkit-0.6.89}/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-0.6.88 → cognite_toolkit-0.6.89}/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-0.6.88 → cognite_toolkit-0.6.89}/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-0.6.88 → cognite_toolkit-0.6.89}/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-0.6.88 → cognite_toolkit-0.6.89}/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",