cognite-toolkit 0.6.97__py3-none-any.whl → 0.7.30__py3-none-any.whl

cognite_toolkit/_cdf_tk/commands/_migrate/data_mapper.py

@@ -1,25 +1,55 @@
 from abc import ABC, abstractmethod
-from typing import Generic
+from collections import defaultdict
+from collections.abc import Callable, Sequence
+from typing import Generic, cast
+from uuid import uuid4
 
-from cognite.client.data_classes._base import (
-    T_CogniteResource,
+from cognite.client.data_classes.data_modeling import (
+    EdgeApply,
+    InstanceApply,
+    NodeApply,
+    NodeId,
+    View,
+    ViewId,
 )
-from cognite.client.data_classes.data_modeling import DirectRelationReference, InstanceApply, View, ViewId
+from cognite.client.exceptions import CogniteException
 
 from cognite_toolkit._cdf_tk.client import ToolkitClient
-from cognite_toolkit._cdf_tk.client.data_classes.migration import ResourceViewMapping
-from cognite_toolkit._cdf_tk.commands._migrate.conversion import asset_centric_to_dm
+from cognite_toolkit._cdf_tk.client.data_classes.canvas import (
+    ContainerReferenceApply,
+    FdmInstanceContainerReferenceApply,
+    IndustrialCanvas,
+    IndustrialCanvasApply,
+)
+from cognite_toolkit._cdf_tk.client.data_classes.charts import Chart, ChartWrite
+from cognite_toolkit._cdf_tk.client.data_classes.charts_data import (
+    ChartCoreTimeseries,
+    ChartSource,
+    ChartTimeseries,
+)
+from cognite_toolkit._cdf_tk.client.data_classes.migration import ResourceViewMappingApply
+from cognite_toolkit._cdf_tk.commands._migrate.conversion import DirectRelationCache, asset_centric_to_dm
 from cognite_toolkit._cdf_tk.commands._migrate.data_classes import AssetCentricMapping
-from cognite_toolkit._cdf_tk.commands._migrate.issues import ConversionIssue, MigrationIssue
+from cognite_toolkit._cdf_tk.commands._migrate.default_mappings import create_default_mappings
+from cognite_toolkit._cdf_tk.commands._migrate.issues import (
+    CanvasMigrationIssue,
+    ChartMigrationIssue,
+    ConversionIssue,
+    MigrationIssue,
+)
 from cognite_toolkit._cdf_tk.commands._migrate.selectors import AssetCentricMigrationSelector
 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.exceptions import ToolkitMigrationError, ToolkitValueError
+from cognite_toolkit._cdf_tk.protocols import T_ResourceRequest, T_ResourceResponse
+from cognite_toolkit._cdf_tk.storageio._base import T_Selector
+from cognite_toolkit._cdf_tk.storageio.selectors import CanvasSelector, ChartSelector
 from cognite_toolkit._cdf_tk.utils import humanize_collection
-from cognite_toolkit._cdf_tk.utils.useful_types import T_AssetCentricResource
+from cognite_toolkit._cdf_tk.utils.useful_types import (
+    T_AssetCentricResourceExtended,
+)
 
 
-class DataMapper(Generic[T_Selector, T_CogniteResource, T_WriteCogniteResource], ABC):
+class DataMapper(Generic[T_Selector, T_ResourceResponse, T_ResourceRequest], ABC):
     def prepare(self, source_selector: T_Selector) -> None:
         """Prepare the data mapper with the given source selector.
 
@@ -31,7 +61,7 @@ class DataMapper(Generic[T_Selector, T_CogniteResource, T_WriteCogniteResource],
         pass
 
     @abstractmethod
-    def map(self, source: T_CogniteResource) -> tuple[T_WriteCogniteResource | None, MigrationIssue]:
+    def map(self, source: Sequence[T_ResourceResponse]) -> Sequence[tuple[T_ResourceRequest | None, MigrationIssue]]:
         """Map a chunk of source data to the target format.
 
         Args:
@@ -45,28 +75,27 @@ class DataMapper(Generic[T_Selector, T_CogniteResource, T_WriteCogniteResource],
 
 
 class AssetCentricMapper(
-    DataMapper[AssetCentricMigrationSelector, AssetCentricMapping[T_AssetCentricResource], InstanceApply]
+    DataMapper[AssetCentricMigrationSelector, AssetCentricMapping[T_AssetCentricResourceExtended], InstanceApply]
 ):
     def __init__(self, client: ToolkitClient) -> None:
         self.client = client
         self._ingestion_view_by_id: dict[ViewId, View] = {}
-        self._view_mapping_by_id: dict[str, ResourceViewMapping] = {}
-        # This is used to keep track of already mapped assets, such that we can creat direct relations
-        # to them from files, events, and time series.
-        self._asset_mapping_by_id: dict[int, DirectRelationReference] = {}
-        self._source_system_mapping_by_id: dict[str, DirectRelationReference] = {}
+        self._view_mapping_by_id: dict[str, ResourceViewMappingApply] = {}
+        self._direct_relation_cache = DirectRelationCache(client)
 
     def prepare(self, source_selector: AssetCentricMigrationSelector) -> None:
         ingestion_view_ids = source_selector.get_ingestion_mappings()
         ingestion_views = self.client.migration.resource_view_mapping.retrieve(ingestion_view_ids)
-        self._view_mapping_by_id = {view.external_id: view for view in ingestion_views}
+        defaults = {mapping.external_id: mapping for mapping in create_default_mappings()}
+        # Custom mappings from CDF override the default mappings
+        self._view_mapping_by_id = defaults | {view.external_id: view.as_write() for view in ingestion_views}
         missing_mappings = set(ingestion_view_ids) - set(self._view_mapping_by_id.keys())
         if missing_mappings:
             raise ToolkitValueError(
                 f"The following ingestion views were not found: {humanize_collection(missing_mappings)}"
             )
 
-        view_ids = list({view.view_id for view in ingestion_views})
+        view_ids = list({mapping.view_id for mapping in self._view_mapping_by_id.values()})
         views = self.client.data_modeling.views.retrieve(view_ids)
         self._ingestion_view_by_id = {view.as_id(): view for view in views}
         missing_views = set(view_ids) - set(self._ingestion_view_by_id.keys())
@@ -74,22 +103,23 @@ class AssetCentricMapper(
             raise ToolkitValueError(
                 f"The following ingestion views were not found in Data Modeling: {humanize_collection(missing_views)}"
             )
-        # We just look-up all source system for now. This can be optimized to only
-        # look-up the ones that are actually used in the ingestion views. However, SourceSystem is typically in
-        # the order ~10 instances, so this is not a big deal for now. See task CDF-25974.
-        source_systems = self.client.migration.created_source_system.list(limit=-1)
-        self._source_system_mapping_by_id = {
-            source_system.source: source_system.as_direct_relation_reference() for source_system in source_systems
-        }
-
-        # We look-up all existing asset mappings to be able to create direct relations to already mapped assets.
-        # This is needed in case the migration is run multiple times, or if assets are mapped
-        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[T_AssetCentricResource]) -> tuple[InstanceApply | None, ConversionIssue]:
+    def map(
+        self, source: Sequence[AssetCentricMapping[T_AssetCentricResourceExtended]]
+    ) -> Sequence[tuple[InstanceApply | None, ConversionIssue]]:
         """Map a chunk of asset-centric data to InstanceApplyList format."""
-        mapping = source.mapping
+        # We update the direct relation cache in bulk for all resources in the chunk.
+        self._direct_relation_cache.update(item.resource for item in source)
+        output: list[tuple[InstanceApply | None, ConversionIssue]] = []
+        for item in source:
+            instance, conversion_issue = self._map_single_item(item)
+            output.append((instance, conversion_issue))
+        return output
+
+    def _map_single_item(
+        self, item: AssetCentricMapping[T_AssetCentricResourceExtended]
+    ) -> tuple[NodeApply | EdgeApply | None, ConversionIssue]:
+        mapping = item.mapping
         ingestion_view = mapping.get_ingestion_view()
         try:
             view_source = self._view_mapping_by_id[ingestion_view]
@@ -99,19 +129,257 @@ class AssetCentricMapper(
                 f"Failed to lookup mapping or view for ingestion view '{ingestion_view}'. Did you forget to call .prepare()?"
             ) from e
         instance, conversion_issue = asset_centric_to_dm(
-            source.resource,
+            item.resource,
             instance_id=mapping.instance_id,
             view_source=view_source,
             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
+            direct_relation_cache=self._direct_relation_cache,
         )
         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}"
-
-        if mapping.resource_type == "asset":
-            self._asset_mapping_by_id[mapping.id] = DirectRelationReference(
-                space=mapping.instance_id.space, external_id=mapping.instance_id.external_id
-            )
         return instance, conversion_issue
+
+
+class ChartMapper(DataMapper[ChartSelector, Chart, ChartWrite]):
+    def __init__(self, client: ToolkitClient) -> None:
+        self.client = client
+
+    def map(self, source: Sequence[Chart]) -> Sequence[tuple[ChartWrite | None, MigrationIssue]]:
+        self._populate_cache(source)
+        output: list[tuple[ChartWrite | None, MigrationIssue]] = []
+        for item in source:
+            mapped_item, issue = self._map_single_item(item)
+            output.append((mapped_item, issue))
+        return output
+
+    def _populate_cache(self, source: Sequence[Chart]) -> None:
+        """Populate the internal cache with timeseries from the source charts.
+
+        Note that the consumption views are also cached as part of the timeseries lookup.
+        """
+        timeseries_ids: set[int] = set()
+        timeseries_external_ids: set[str] = set()
+        for chart in source:
+            for item in chart.data.time_series_collection or []:
+                if item.ts_id:
+                    timeseries_ids.add(item.ts_id)
+                if item.ts_external_id:
+                    timeseries_external_ids.add(item.ts_external_id)
+        if timeseries_ids:
+            self.client.migration.lookup.time_series(list(timeseries_ids))
+        if timeseries_external_ids:
+            self.client.migration.lookup.time_series(external_id=list(timeseries_external_ids))
+
+    def _map_single_item(self, item: Chart) -> tuple[ChartWrite | None, ChartMigrationIssue]:
+        issue = ChartMigrationIssue(chart_external_id=item.external_id)
+        time_series_collection = item.data.time_series_collection or []
+        timeseries_core_collection = self._create_timeseries_core_collection(time_series_collection, issue)
+        if issue.has_issues:
+            return None, issue
+
+        updated_source_collection = self._update_source_collection(
+            item.data.source_collection or [], time_series_collection, timeseries_core_collection
+        )
+
+        mapped_chart = item.as_write()
+        mapped_chart.data.core_timeseries_collection = timeseries_core_collection
+        mapped_chart.data.time_series_collection = None
+        mapped_chart.data.source_collection = updated_source_collection
+        return mapped_chart, issue
+
+    def _create_timeseries_core_collection(
+        self, time_series_collection: list[ChartTimeseries], issue: ChartMigrationIssue
+    ) -> list[ChartCoreTimeseries]:
+        timeseries_core_collection: list[ChartCoreTimeseries] = []
+        for ts_item in time_series_collection or []:
+            node_id, consumer_view_id = self._get_node_id_consumer_view_id(ts_item)
+
+            if node_id is None:
+                if ts_item.ts_id is not None:
+                    issue.missing_timeseries_ids.append(ts_item.ts_id)
+                elif ts_item.ts_external_id is not None:
+                    issue.missing_timeseries_external_ids.append(ts_item.ts_external_id)
+                else:
+                    issue.missing_timeseries_identifier.append(ts_item.id or "unknown")
+                continue
+
+            core_timeseries = self._create_new_timeseries_core(ts_item, node_id, consumer_view_id)
+            timeseries_core_collection.append(core_timeseries)
+        return timeseries_core_collection
+
+    def _create_new_timeseries_core(
+        self, ts_item: ChartTimeseries, node_id: NodeId, consumer_view_id: ViewId | None
+    ) -> ChartCoreTimeseries:
+        dumped = ts_item.model_dump(mode="json", by_alias=True, exclude_unset=True)
+        dumped["nodeReference"] = node_id
+        dumped["viewReference"] = consumer_view_id
+        new_uuid = str(uuid4())
+        dumped["id"] = new_uuid
+        dumped["type"] = "coreTimeseries"
+        # We ignore extra here to only include the fields that are shared between ChartTimeseries and ChartCoreTimeseries
+        core_timeseries = ChartCoreTimeseries.model_validate(dumped, extra="ignore")
+        return core_timeseries
+
+    def _get_node_id_consumer_view_id(self, ts_item: ChartTimeseries) -> tuple[NodeId | None, ViewId | None]:
+        """Look up the node ID and consumer view ID for a given timeseries item.
+
+        Prioritizes lookup by internal ID, then by external ID.
+
+        Args:
+            ts_item: The ChartTimeseries item to look up.
+
+        Returns:
+            A tuple containing the consumer view ID and node ID, or None if not found.
+        """
+        node_id: NodeId | None = None
+        consumer_view_id: ViewId | None = None
+        for id_name, id_value in [("id", ts_item.ts_id), ("external_id", ts_item.ts_external_id)]:
+            if id_value is None:
+                continue
+            arg = {id_name: id_value}
+            node_id = self.client.migration.lookup.time_series(**arg)  # type: ignore[arg-type]
+            consumer_view_id = self.client.migration.lookup.time_series.consumer_view(**arg)  # type: ignore[arg-type]
+            if node_id is not None:
+                break
+        return node_id, consumer_view_id
+
+    def _update_source_collection(
+        self,
+        source_collection: list[ChartSource],
+        time_series_collection: list[ChartTimeseries],
+        timeseries_core_collection: list[ChartCoreTimeseries],
+    ) -> list[ChartSource]:
+        remove_ids = {ts_item.id for ts_item in time_series_collection if ts_item.id is not None}
+        updated_source_collection = [ts_item for ts_item in source_collection if ts_item.id not in remove_ids]
+        for core_ts_item in timeseries_core_collection:
+            # We cast there two as we set them in the _create_timeseries_core_collection method
+            new_source_item = ChartSource(id=cast(str, core_ts_item.id), type=cast(str, core_ts_item.type))
+            updated_source_collection.append(new_source_item)
+        return updated_source_collection
+
+
+class CanvasMapper(DataMapper[CanvasSelector, IndustrialCanvas, IndustrialCanvasApply]):
+    # Note sequences are not supported in Canvas, so we do not include them here.
+    asset_centric_resource_types = tuple(("asset", "event", "timeseries", "file"))
+    DEFAULT_ASSET_VIEW = ViewId("cdf_cdm", "CogniteAsset", "v1")
+    DEFAULT_EVENT_VIEW = ViewId("cdf_cdm", "CogniteActivity", "v1")
+    DEFAULT_FILE_VIEW = ViewId("cdf_cdm", "CogniteFile", "v1")
+    DEFAULT_TIMESERIES_VIEW = ViewId("cdf_cdm", "CogniteTimeSeries", "v1")
+
+    def __init__(self, client: ToolkitClient, dry_run: bool, skip_on_missing_ref: bool = False) -> None:
+        self.client = client
+        self.dry_run = dry_run
+        self.skip_on_missing_ref = skip_on_missing_ref
+
+    def map(self, source: Sequence[IndustrialCanvas]) -> Sequence[tuple[IndustrialCanvasApply | None, MigrationIssue]]:
+        self._populate_cache(source)
+        output: list[tuple[IndustrialCanvasApply | None, MigrationIssue]] = []
+        for item in source:
+            mapped_item, issue = self._map_single_item(item)
+            output.append((mapped_item, issue))
+        return output
+
+    @property
+    def lookup_methods(self) -> dict[str, Callable]:
+        return {
+            "asset": self.client.migration.lookup.assets,
+            "event": self.client.migration.lookup.events,
+            "timeseries": self.client.migration.lookup.time_series,
+            "file": self.client.migration.lookup.files,
+        }
+
+    def _populate_cache(self, source: Sequence[IndustrialCanvas]) -> None:
+        """Populate the internal cache with references from the source canvases.
+
+        Note that the consumption views are also cached as part of the timeseries lookup.
+        """
+        ids_by_type: dict[str, set[int]] = defaultdict(set)
+        for canvas in source:
+            for ref in canvas.container_references or []:
+                if ref.container_reference_type in self.asset_centric_resource_types:
+                    ids_by_type[ref.container_reference_type].add(ref.resource_id)
+
+        for resource_type, lookup_method in self.lookup_methods.items():
+            ids = ids_by_type.get(resource_type)
+            if ids:
+                lookup_method(list(ids))
+
+    def _get_node_id(self, resource_id: int, resource_type: str) -> NodeId | None:
+        """Look up the node ID for a given resource ID and type."""
+        try:
+            return self.lookup_methods[resource_type](resource_id)
+        except KeyError:
+            raise ToolkitValueError(f"Unsupported resource type '{resource_type}' for container reference migration.")
+
+    def _get_consumer_view_id(self, resource_id: int, resource_type: str) -> ViewId:
+        """Look up the consumer view ID for a given resource ID and type."""
+        lookup_map = {
+            "asset": (self.client.migration.lookup.assets.consumer_view, self.DEFAULT_ASSET_VIEW),
+            "event": (self.client.migration.lookup.events.consumer_view, self.DEFAULT_EVENT_VIEW),
+            "timeseries": (self.client.migration.lookup.time_series.consumer_view, self.DEFAULT_TIMESERIES_VIEW),
+            "file": (self.client.migration.lookup.files.consumer_view, self.DEFAULT_FILE_VIEW),
+        }
+        if lookup_tuple := lookup_map.get(resource_type):
+            method, fallback = lookup_tuple
+            return method(resource_id) or fallback
+
+        raise ToolkitValueError(f"Unsupported resource type '{resource_type}' for container reference migration.")
+
+    def _map_single_item(self, canvas: IndustrialCanvas) -> tuple[IndustrialCanvasApply | None, CanvasMigrationIssue]:
+        update = canvas.as_write()
+        issue = CanvasMigrationIssue(canvas_external_id=canvas.canvas.external_id, canvas_name=canvas.canvas.name)
+
+        remaining_container_references: list[ContainerReferenceApply] = []
+        new_fdm_references: list[FdmInstanceContainerReferenceApply] = []
+        for ref in update.container_references or []:
+            if ref.container_reference_type not in self.asset_centric_resource_types:
+                remaining_container_references.append(ref)
+                continue
+            node_id = self._get_node_id(ref.resource_id, ref.container_reference_type)
+            if node_id is None:
+                issue.missing_reference_ids.append(ref.as_asset_centric_id())
+            else:
+                consumer_view = self._get_consumer_view_id(ref.resource_id, ref.container_reference_type)
+                fdm_ref = self.migrate_container_reference(ref, canvas.canvas.external_id, node_id, consumer_view)
+                new_fdm_references.append(fdm_ref)
+        if issue.missing_reference_ids and self.skip_on_missing_ref:
+            return None, issue
+
+        update.container_references = remaining_container_references
+        update.fdm_instance_container_references.extend(new_fdm_references)
+        if not self.dry_run:
+            backup = canvas.as_write().create_backup()
+            try:
+                self.client.canvas.industrial.create(backup)
+            except CogniteException as e:
+                raise ToolkitMigrationError(
+                    f"Failed to create backup for canvas '{canvas.canvas.name}': {e!s}. "
+                ) from e
+
+        return update, issue
+
+    @classmethod
+    def migrate_container_reference(
+        cls, reference: ContainerReferenceApply, canvas_external_id: str, instance_id: NodeId, consumer_view: ViewId
+    ) -> FdmInstanceContainerReferenceApply:
+        """Migrate a single container reference by replacing the asset-centric ID with the data model instance ID."""
+        new_id = str(uuid4())
+        new_external_id = f"{canvas_external_id}_{new_id}"
+        return FdmInstanceContainerReferenceApply(
+            external_id=new_external_id,
+            id_=new_id,
+            container_reference_type="fdmInstance",
+            instance_space=instance_id.space,
+            instance_external_id=instance_id.external_id,
+            view_space=consumer_view.space,
+            view_external_id=consumer_view.external_id,
+            view_version=consumer_view.version,
+            label=reference.label,
+            properties_=reference.properties_,
+            x=reference.x,
+            y=reference.y,
+            width=reference.width,
+            height=reference.height,
+            max_width=reference.max_width,
+            max_height=reference.max_height,
+        )

cognite_toolkit/_cdf_tk/commands/_migrate/default_mappings.py

@@ -93,7 +93,7 @@ def create_default_mappings() -> list[ResourceViewMappingApply]:
             view_id=ViewId("cdf_cdm", "CogniteDiagramAnnotation", "v1"),
             property_mapping={
                 # We are ignoring the symbol region in the default mapping.
-                "annotatedResource.id": "edge.startNode",
+                "annotatedResourceId": "edge.startNode",
                 "annotationType": "edge.type.externalId",
                 "creatingUser": "sourceCreatedUser",
                 "creatingApp": "sourceId",
@@ -114,16 +114,16 @@ def create_default_mappings() -> list[ResourceViewMappingApply]:
         ResourceViewMappingApply(
             external_id=FILE_ANNOTATIONS_ID,
             resource_type="fileAnnotation",
-            view_id=ViewId("cdf_cdm", "CogniteFileAnnotation", "v1"),
+            view_id=ViewId("cdf_cdm", "CogniteDiagramAnnotation", "v1"),
             property_mapping={
-                "annotatedResource.id": "edge.startNode",
+                "annotatedResourceId": "edge.startNode",
                 "annotationType": "edge.type.externalId",
                 "creatingUser": "sourceCreatedUser",
                 "creatingApp": "sourceId",
                 "creatingAppVersion": "sourceContext",
                 "status": "status",
-                "data.fileRef.id": "edge.startNode",
-                "data.fileRef.externalId": "edge.startNode",
+                "data.fileRef.id": "edge.endNode",
+                "data.fileRef.externalId": "edge.endNode",
                 "data.description": "description",
                 "data.pageNumber": "startNodePageNumber",
                 "data.textRegion.confidence": "confidence",

cognite_toolkit/_cdf_tk/commands/_migrate/issues.py

@@ -30,6 +30,39 @@ class MigrationIssue(MigrationObject):
         return True
 
 
+class ChartMigrationIssue(MigrationIssue):
+    """Represents a chart migration issue encountered during migration.
+
+    Attributes:
+        chart_external_id (str): The external ID of the chart that could not be migrated.
+    """
+
+    type: ClassVar[str] = "chartMigration"
+    chart_external_id: str
+    missing_timeseries_ids: list[int] = Field(default_factory=list)
+    missing_timeseries_external_ids: list[str] = Field(default_factory=list)
+    missing_timeseries_identifier: list[str] = Field(default_factory=list)
+
+    @property
+    def has_issues(self) -> bool:
+        """Check if there are any issues recorded in this ChartMigrationIssue."""
+        return bool(
+            self.missing_timeseries_ids or self.missing_timeseries_external_ids or self.missing_timeseries_identifier
+        )
+
+
+class CanvasMigrationIssue(MigrationIssue):
+    type: ClassVar[str] = "canvasMigration"
+    canvas_external_id: str
+    canvas_name: str
+    missing_reference_ids: list[AssetCentricId] = Field(default_factory=list)
+
+    @property
+    def has_issues(self) -> bool:
+        """Check if there are any issues recorded in this CanvasMigrationIssue."""
+        return bool(self.missing_reference_ids)
+
+
 class ReadIssue(MigrationIssue):
     """Represents a read issue encountered during migration."""