cognite-toolkit 0.5.44__py3-none-any.whl → 0.7.34__py3-none-any.whl

cognite_toolkit/_cdf_tk/commands/_migrate/creators.py

@@ -0,0 +1,232 @@
+from abc import ABC, abstractmethod
+from collections.abc import Iterable
+from dataclasses import dataclass
+from functools import cached_property
+from typing import Any, Generic
+
+from cognite.client.data_classes import DataSetList, filters
+from cognite.client.data_classes._base import T_CogniteResourceList
+from cognite.client.data_classes.aggregations import UniqueResult
+from cognite.client.data_classes.assets import AssetProperty
+from cognite.client.data_classes.data_modeling import (
+    NodeApply,
+    NodeApplyList,
+    NodeOrEdgeData,
+    SpaceApply,
+    SpaceApplyList,
+    ViewId,
+)
+from cognite.client.data_classes.documents import SourceFileProperty
+from cognite.client.data_classes.events import EventProperty
+
+from cognite_toolkit._cdf_tk.client import ToolkitClient
+from cognite_toolkit._cdf_tk.client.data_classes.apm_config_v1 import APMConfig, APMConfigList
+from cognite_toolkit._cdf_tk.cruds import NodeCRUD, ResourceCRUD, SpaceCRUD
+from cognite_toolkit._cdf_tk.exceptions import ToolkitMissingResourceError, ToolkitRequiredValueError
+from cognite_toolkit._cdf_tk.utils import humanize_collection
+
+from .data_model import CREATED_SOURCE_SYSTEM_VIEW_ID, SPACE, SPACE_SOURCE_VIEW_ID
+
+
+@dataclass
+class ResourceConfig:
+    filestem: str
+    data: dict[str, Any]
+
+
+class MigrationCreator(ABC, Generic[T_CogniteResourceList]):
+    """Base class for migration resources configurations that are created resources."""
+
+    CRUD: type[ResourceCRUD]
+    DISPLAY_NAME: str
+    HAS_LINEAGE: bool = True
+
+    def __init__(self, client: ToolkitClient) -> None:
+        self.client = client
+
+    @abstractmethod
+    def create_resources(self) -> T_CogniteResourceList:
+        raise NotImplementedError("Subclasses should implement this method")
+
+    @abstractmethod
+    def resource_configs(self, resources: T_CogniteResourceList) -> list[ResourceConfig]:
+        raise NotImplementedError("Subclasses should implement this method")
+
+    def store_lineage(self, resources: T_CogniteResourceList) -> int:
+        """Store lineage information for the created resources.
+
+        Args:
+            resources: The list of created resources.
+        """
+        raise NotImplementedError("Subclasses should implement this method")
+
+
+class InstanceSpaceCreator(MigrationCreator[SpaceApplyList]):
+    """Creates instance spaces for migration."""
+
+    CRUD = SpaceCRUD
+    DISPLAY_NAME = "Instance Space"
+    HAS_LINEAGE = True
+
+    def __init__(
+        self, client: ToolkitClient, datasets: DataSetList | None = None, data_set_external_ids: list[str] | None = None
+    ) -> None:
+        super().__init__(client)
+        if sum([datasets is not None, data_set_external_ids is not None]) != 1:
+            raise ValueError("Exactly one of datasets or data_set_external_ids must be provided.")
+        self.data_set_external_ids = data_set_external_ids
+        self.datasets = datasets or DataSetList([])
+
+    def create_resources(self) -> SpaceApplyList:
+        if self.data_set_external_ids is not None:
+            self.datasets = self.client.data_sets.retrieve_multiple(external_ids=self.data_set_external_ids)
+
+        if missing_external_ids := [ds.id for ds in self.datasets if ds.external_id is None]:
+            raise ToolkitRequiredValueError(
+                f"Cannot create instance spaces for datasets with missing external IDs: {humanize_collection(missing_external_ids)}"
+            )
+
+        return SpaceApplyList(
+            [
+                SpaceApply(
+                    # This is checked above
+                    space=dataset.external_id,  # type: ignore[arg-type]
+                    name=dataset.name,
+                    description=dataset.description,
+                )
+                for dataset in self.datasets
+            ]
+        )
+
+    def resource_configs(self, resources: SpaceApplyList) -> list[ResourceConfig]:
+        return [ResourceConfig(filestem=space.space, data=space.dump()) for space in resources]
+
+    def store_lineage(self, resources: SpaceApplyList) -> int:
+        data_set_by_external_id = {ds.external_id: ds for ds in self.datasets}
+        nodes = NodeApplyList(
+            [
+                NodeApply(
+                    space=SPACE.space,
+                    external_id=space.space,
+                    sources=[
+                        NodeOrEdgeData(
+                            source=SPACE_SOURCE_VIEW_ID,
+                            properties={
+                                "instanceSpace": space.space,
+                                "dataSetId": data_set_by_external_id[space.space].id,
+                                "dataSetExternalId": data_set_by_external_id[space.space].external_id,
+                            },
+                        )
+                    ],
+                )
+                for space in resources
+            ]
+        )
+        res = self.client.data_modeling.instances.apply(nodes)
+        return len(res.nodes)
+
+
+class SourceSystemCreator(MigrationCreator[NodeApplyList]):
+    CRUD = NodeCRUD
+    DISPLAY_NAME = "Source System"
+    COGNITE_SOURCE_SYSTEM_VIEW_ID = ViewId("cdf_cdm", "CogniteSourceSystem", "v1")
+
+    def __init__(
+        self,
+        client: ToolkitClient,
+        instance_space: str,
+        data_set_external_id: str | None = None,
+        hierarchy: str | None = None,
+    ) -> None:
+        super().__init__(client)
+        if sum([data_set_external_id is not None, hierarchy is not None]) != 1:
+            raise ValueError("Exactly one of data_set_external_id or hierarchy must be provided.")
+        self._instance_space = instance_space
+        self.data_set_external_id = data_set_external_id
+        self.hierarchy = hierarchy
+
+    def create_resources(self) -> NodeApplyList:
+        seen: set[str] = set()
+        nodes = NodeApplyList([])
+        for source in self._lookup_sources():
+            source_str = source.value
+            if not isinstance(source_str, str) or source_str in seen:
+                continue
+            seen.add(source_str)
+            nodes.append(
+                NodeApply(
+                    space=self._instance_space,
+                    external_id=source_str,
+                    sources=[
+                        NodeOrEdgeData(source=self.COGNITE_SOURCE_SYSTEM_VIEW_ID, properties={"name": source_str}),
+                        NodeOrEdgeData(source=CREATED_SOURCE_SYSTEM_VIEW_ID, properties={"source": source_str}),
+                    ],
+                )
+            )
+        return nodes
+
+    def resource_configs(self, resources: NodeApplyList) -> list[ResourceConfig]:
+        output: list[ResourceConfig] = []
+        for node in resources:
+            copy = NodeApply(
+                space=node.space,
+                external_id=node.external_id,
+                # We remove the lineage source as this is not expected to be part of the governed
+                # SourceSystem.
+                sources=[s for s in node.sources if s.source != CREATED_SOURCE_SYSTEM_VIEW_ID],
+            )
+            output.append(ResourceConfig(filestem=node.external_id, data=copy.dump()))
+        return output
+
+    def _lookup_sources(self) -> Iterable[UniqueResult]:
+        yield from self.client.assets.aggregate_unique_values(AssetProperty.source, filter=self._simple_filter)
+        yield from self.client.events.aggregate_unique_values(property=EventProperty.source, filter=self._simple_filter)
+        yield from self.client.documents.aggregate_unique_values(
+            SourceFileProperty.source, filter=self._advanced_filter, limit=1000
+        )
+
+    @cached_property
+    def _simple_filter(self) -> dict[str, Any] | None:
+        if self.data_set_external_id is not None:
+            return {"dataSetIds": [{"externalId": self.data_set_external_id}]}
+        if self.hierarchy is not None:
+            return {"assetSubtreeIds": [{"externalId": self.hierarchy}]}
+        return None
+
+    @cached_property
+    def _advanced_filter(self) -> filters.Filter:
+        if self.data_set_external_id is not None:
+            data_set_id = self.client.lookup.data_sets.id(self.data_set_external_id)
+            if data_set_id is None:
+                raise ToolkitMissingResourceError(f"Data set with external ID '{self.data_set_external_id}' not found.")
+            return filters.Equals(SourceFileProperty.data_set_id, data_set_id)
+        if self.hierarchy is not None:
+            return filters.InAssetSubtree(SourceFileProperty.asset_external_ids, [self.hierarchy])
+        else:
+            raise ValueError("This should not happen.")
+
+    def store_lineage(self, resources: NodeApplyList) -> int:
+        # We already store lineage when creating the resources.
+        return len(resources)
+
+
+class InfieldV2ConfigCreator(MigrationCreator[NodeApplyList]):
+    CRUD = NodeCRUD
+    DISPLAY_NAME = "Infield V2 Configuration"
+    HAS_LINEAGE = False
+
+    def create_resources(self) -> NodeApplyList:
+        apm_config_nodes = self.client.data_modeling.instances.list(instance_type="node", sources=APMConfig.view_id)
+        apm_config = APMConfigList.from_nodes(apm_config_nodes)
+
+        new_config_nodes = NodeApplyList([])
+        for config in apm_config:
+            new_config = self._create_infield_v2_config(config)
+            new_config_nodes.append(new_config)
+        return new_config_nodes
+
+    def resource_configs(self, resources: NodeApplyList) -> list[ResourceConfig]:
+        return [ResourceConfig(filestem=node.external_id, data=node.dump()) for node in resources]
+
+    def _create_infield_v2_config(self, config: APMConfig) -> NodeApply:
+        raise NotImplementedError("To be implemented")

cognite_toolkit/_cdf_tk/commands/_migrate/data_classes.py

@@ -0,0 +1,294 @@
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Annotated, Any, Generic, Literal
+
+from cognite.client.data_classes._base import (
+    WriteableCogniteResource,
+    WriteableCogniteResourceList,
+)
+from cognite.client.data_classes.data_modeling import EdgeId, InstanceApply, NodeId, ViewId
+from cognite.client.utils._identifier import InstanceId
+from cognite.client.utils._text import to_camel_case
+from pydantic import BaseModel, BeforeValidator, Field, field_validator, model_validator
+
+from cognite_toolkit._cdf_tk.client.data_classes.base import BaseModelObject, RequestResource
+from cognite_toolkit._cdf_tk.client.data_classes.instance_api import InstanceIdentifier
+from cognite_toolkit._cdf_tk.client.data_classes.instances import InstanceApplyList
+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 (
+    ASSET_ANNOTATIONS_ID,
+    FILE_ANNOTATIONS_ID,
+    create_default_mappings,
+)
+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 (
+    AssetCentricKindExtended,
+    JsonVal,
+    T_AssetCentricResourceExtended,
+)
+
+
+class MigrationMapping(BaseModel, alias_generator=to_camel_case, extra="ignore", populate_by_name=True):
+    """The mapping between an asset-centric ID and a data modeling instance ID.
+    Args
+        resource_type (str): The asset-centric type of the resource (e.g., "asset", "event", "timeseries").
+        instance_id (NodeId): The target NodeId in data modeling.
+        id (int): The asset-centric ID of the resource.
+        data_set_id (int | None): The data set ID of the resource. This is used to validate access to the resource.
+        ingestion_view (str | None): The ingestion view name. This is the view mapping that will be used to
+            ingest the resource into data modeling.
+        preferred_consumer_view (ViewId | None): The preferred consumer view for the resource. This is used in
+           for example, the Canvas migration to determine which view to use for the resource.
+    """
+
+    resource_type: str
+    instance_id: InstanceId
+    id: int
+    data_set_id: int | None = None
+    ingestion_view: str | None = None
+    preferred_consumer_view: ViewId | None = None
+
+    def get_ingestion_view(self) -> str:
+        """Get the ingestion view for the mapping. If not specified, return the default ingestion view."""
+        if self.ingestion_view:
+            return self.ingestion_view
+
+        default_mappings = create_default_mappings()
+        for mapping in default_mappings:
+            if mapping.resource_type == self.resource_type:
+                return mapping.external_id
+        raise ToolkitValueError(f"No default ingestion view specified for resource type '{self.resource_type}'")
+
+    def as_asset_centric_id(self) -> AssetCentricId:
+        # MyPy fails to understand that resource_type is AssetCentricKindExtended in subclasses
+        return AssetCentricId(resource_type=self.resource_type, id_=self.id)  # type: ignore[arg-type]
+
+    @model_validator(mode="before")
+    def _handle_flat_dict(cls, values: Any) -> Any:
+        if not isinstance(values, dict):
+            return values
+
+        if "space" in values and "externalId" in values:
+            values["instanceId"] = {"space": values.pop("space"), "externalId": values.pop("externalId")}
+        if "consumerViewSpace" in values and "consumerViewExternalId" in values:
+            consumer_view = {
+                "space": values.pop("consumerViewSpace"),
+                "externalId": values.pop("consumerViewExternalId"),
+            }
+            if "consumerViewVersion" in values:
+                consumer_view["version"] = values.pop("consumerViewVersion")
+            values["preferredConsumerView"] = consumer_view
+        return values
+
+    @field_validator("data_set_id", "ingestion_view", mode="before")
+    def _empty_string_to_none(cls, v: Any) -> Any:
+        if isinstance(v, str) and not v.strip():
+            return None
+        return v
+
+    @field_validator("preferred_consumer_view", mode="before")
+    def _validate_preferred_consumer_view(cls, v: Any) -> Any:
+        if isinstance(v, dict):
+            return ViewId.load(v)
+        return v
+
+
+class MigrationMappingList(ModelList[MigrationMapping]):
+    @classmethod
+    def _get_base_model_cls(cls) -> type[MigrationMapping]:
+        return MigrationMapping
+
+    @classmethod
+    def _required_header_names(cls) -> set[str]:
+        return {"id", "space", "externalId"}
+
+    @classmethod
+    def _optional_header_names(cls) -> set[str]:
+        return {"dataSetId", "ingestionView", "consumerViewSpace", "consumerViewExternalId", "consumerViewVersion"}
+
+    def get_ids(self) -> list[int]:
+        """Return a list of IDs from the migration mappings."""
+        return [mapping.id for mapping in self]
+
+    def as_node_ids(self) -> list[NodeId]:
+        """Return a list of NodeIds from the migration mappings."""
+        return [mapping.instance_id for mapping in self if isinstance(mapping.instance_id, NodeId)]
+
+    def as_edge_ids(self) -> list[EdgeId]:
+        """Return a list of EdgeIds from the migration mappings."""
+        return [mapping.instance_id for mapping in self if isinstance(mapping.instance_id, EdgeId)]
+
+    def spaces(self) -> set[str]:
+        """Return a set of spaces from the migration mappings."""
+        return {mapping.instance_id.space for mapping in self}
+
+    def as_pending_ids(self) -> list[PendingInstanceId]:
+        return [
+            PendingInstanceId(pending_instance_id=mapping.instance_id, id=mapping.id)
+            for mapping in self
+            if isinstance(mapping.instance_id, NodeId)
+        ]
+
+    def get_data_set_ids(self) -> set[int]:
+        """Return a list of data set IDs from the migration mappings."""
+        return {mapping.data_set_id for mapping in self if mapping.data_set_id is not None}
+
+    def as_mapping_by_id(self) -> dict[int, MigrationMapping]:
+        """Return a mapping of IDs to MigrationMapping objects."""
+        return {mapping.id: mapping for mapping in self}
+
+    @classmethod
+    def read_csv_file(
+        cls, filepath: Path, resource_type: AssetCentricKindExtended | None = None
+    ) -> "MigrationMappingList":
+        if cls is not MigrationMappingList or resource_type is None:
+            return super().read_csv_file(filepath)
+        cls_by_resource_type: dict[str, type[MigrationMappingList]] = {
+            "Assets": AssetMigrationMappingList,
+            "TimeSeries": TimeSeriesMigrationMappingList,
+            "FileMetadata": FileMigrationMappingList,
+            "Events": EventMigrationMappingList,
+            "Annotations": AnnotationMigrationMappingList,
+        }
+        if resource_type not in cls_by_resource_type:
+            raise ToolkitValueError(
+                f"Invalid resource type '{resource_type}'. Must be one of 'asset', 'timeseries', or 'file'."
+            )
+        return cls_by_resource_type[resource_type].read_csv_file(filepath, resource_type=None)
+
+
+def _validate_node_id(value: Any) -> Any:
+    if isinstance(value, dict):
+        return NodeId.load(value)
+    return value
+
+
+class AssetMapping(MigrationMapping):
+    resource_type: Literal["asset"] = "asset"
+    instance_id: Annotated[NodeId, BeforeValidator(_validate_node_id)]
+
+
+class EventMapping(MigrationMapping):
+    resource_type: Literal["event"] = "event"
+    instance_id: Annotated[NodeId, BeforeValidator(_validate_node_id)]
+
+
+class TimeSeriesMapping(MigrationMapping):
+    resource_type: Literal["timeseries"] = "timeseries"
+    instance_id: Annotated[NodeId, BeforeValidator(_validate_node_id)]
+
+
+class FileMapping(MigrationMapping):
+    resource_type: Literal["file"] = "file"
+    instance_id: Annotated[NodeId, BeforeValidator(_validate_node_id)]
+
+
+class AnnotationMapping(MigrationMapping):
+    resource_type: Literal["annotation"] = "annotation"
+    instance_id: EdgeId
+    annotation_type: Literal["diagrams.AssetLink", "diagrams.FileLink"] | None = None
+
+    def get_ingestion_view(self) -> str:
+        """Get the ingestion view for the mapping. If not specified, return the default ingestion view."""
+        if self.ingestion_view:
+            return self.ingestion_view
+        elif self.annotation_type == "diagrams.AssetLink":
+            return ASSET_ANNOTATIONS_ID
+        elif self.annotation_type == "diagrams.FileLink":
+            return FILE_ANNOTATIONS_ID
+        else:
+            raise ToolkitValueError("Cannot determine default ingestion view for annotation without annotation_type")
+
+    @field_validator("instance_id", mode="before")
+    def _validate_instance_id(cls, v: Any) -> Any:
+        if isinstance(v, dict):
+            return EdgeId.load(v)
+        return v
+
+
+class AssetMigrationMappingList(MigrationMappingList):
+    @classmethod
+    def _get_base_model_cls(cls) -> type[AssetMapping]:
+        return AssetMapping
+
+
+class EventMigrationMappingList(MigrationMappingList):
+    @classmethod
+    def _get_base_model_cls(cls) -> type[EventMapping]:
+        return EventMapping
+
+
+class FileMigrationMappingList(MigrationMappingList):
+    @classmethod
+    def _get_base_model_cls(cls) -> type[FileMapping]:
+        return FileMapping
+
+
+class TimeSeriesMigrationMappingList(MigrationMappingList):
+    @classmethod
+    def _get_base_model_cls(cls) -> type[TimeSeriesMapping]:
+        return TimeSeriesMapping
+
+
+class AnnotationMigrationMappingList(MigrationMappingList):
+    @classmethod
+    def _get_base_model_cls(cls) -> type[AnnotationMapping]:
+        return AnnotationMapping
+
+
+@dataclass
+class AssetCentricMapping(Generic[T_AssetCentricResourceExtended], WriteableCogniteResource[InstanceApply]):
+    mapping: MigrationMapping
+    resource: T_AssetCentricResourceExtended
+
+    def as_write(self) -> InstanceApply:
+        raise NotImplementedError()
+
+    def dump(self, camel_case: bool = True) -> dict[str, JsonVal]:
+        mapping = self.mapping.model_dump(exclude_unset=True, by_alias=camel_case)
+        # Ensure that resource type is always included, even if unset.
+        mapping["resourceType" if camel_case else "resource_type"] = self.mapping.resource_type
+        return {
+            "mapping": mapping,
+            "resource": self.resource.dump(camel_case=camel_case),
+        }
+
+
+class AssetCentricMappingList(
+    WriteableCogniteResourceList[InstanceApply, AssetCentricMapping[T_AssetCentricResourceExtended]]
+):
+    _RESOURCE: type = AssetCentricMapping
+
+    def as_write(self) -> InstanceApplyList:
+        return InstanceApplyList([item.as_write() for item in self])
+
+
+class Model(BaseModelObject):
+    instance_id: InstanceIdentifier
+
+
+class Thumbnail(BaseModelObject):
+    instance_id: InstanceIdentifier
+
+
+class ThreeDRevisionMigrationRequest(RequestResource):
+    space: str
+    type: Literal["CAD", "PointCloud", "Image360"]
+    revision_id: int
+    model: Model
+
+    def as_id(self) -> int:
+        return self.revision_id
+
+
+class ThreeDMigrationRequest(RequestResource):
+    model_id: int
+    type: Literal["CAD", "PointCloud", "Image360"]
+    space: str
+    thumbnail: Thumbnail | None = None
+    revision: ThreeDRevisionMigrationRequest = Field(exclude=True)
+
+    def as_id(self) -> int:
+        return self.model_id