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

cognite_toolkit/_cdf_tk/commands/_migrate/migration_io.py

@@ -0,0 +1,350 @@
+from collections.abc import Iterable, Iterator, Mapping, Sequence
+from typing import ClassVar, cast
+
+from cognite.client.data_classes import Annotation
+from cognite.client.data_classes.data_modeling import EdgeId, InstanceApply, NodeId
+
+from cognite_toolkit._cdf_tk.client import ToolkitClient
+from cognite_toolkit._cdf_tk.client.data_classes.pending_instances_ids import PendingInstanceId
+from cognite_toolkit._cdf_tk.constants import MISSING_EXTERNAL_ID, MISSING_INSTANCE_SPACE
+from cognite_toolkit._cdf_tk.exceptions import ToolkitNotImplementedError, ToolkitValueError
+from cognite_toolkit._cdf_tk.storageio import (
+    AnnotationIO,
+    HierarchyIO,
+    InstanceIO,
+    UploadableStorageIO,
+)
+from cognite_toolkit._cdf_tk.storageio._base import Page, UploadItem
+from cognite_toolkit._cdf_tk.tk_warnings import MediumSeverityWarning
+from cognite_toolkit._cdf_tk.utils.collection import chunker_sequence
+from cognite_toolkit._cdf_tk.utils.http_client import HTTPClient, HTTPMessage, ItemsRequest, SuccessResponseItems
+from cognite_toolkit._cdf_tk.utils.useful_types import (
+    AssetCentricKindExtended,
+    AssetCentricType,
+    JsonVal,
+    T_AssetCentricResource,
+)
+
+from .data_classes import (
+    AnnotationMapping,
+    AssetCentricMapping,
+    AssetCentricMappingList,
+    MigrationMapping,
+    MigrationMappingList,
+)
+from .data_model import INSTANCE_SOURCE_VIEW_ID
+from .default_mappings import ASSET_ANNOTATIONS_ID, FILE_ANNOTATIONS_ID
+from .selectors import AssetCentricMigrationSelector, MigrateDataSetSelector, MigrationCSVFileSelector
+
+
+class AssetCentricMigrationIO(
+    UploadableStorageIO[AssetCentricMigrationSelector, AssetCentricMapping[T_AssetCentricResource], InstanceApply]
+):
+    KIND = "AssetCentricMigration"
+    SUPPORTED_DOWNLOAD_FORMATS = frozenset({".parquet", ".csv", ".ndjson"})
+    SUPPORTED_COMPRESSIONS = frozenset({".gz"})
+    SUPPORTED_READ_FORMATS = frozenset({".parquet", ".csv", ".ndjson", ".yaml", ".yml"})
+    CHUNK_SIZE = 1000
+    UPLOAD_ENDPOINT = InstanceIO.UPLOAD_ENDPOINT
+
+    PENDING_INSTANCE_ID_ENDPOINT_BY_KIND: ClassVar[Mapping[AssetCentricKindExtended, str]] = {
+        "TimeSeries": "/timeseries/set-pending-instance-ids",
+        "FileMetadata": "/files/set-pending-instance-ids",
+    }
+
+    def __init__(self, client: ToolkitClient, skip_linking: bool = True) -> None:
+        super().__init__(client)
+        self.hierarchy = HierarchyIO(client)
+        self.skip_linking = skip_linking
+
+    def as_id(self, item: AssetCentricMapping) -> str:
+        return f"{item.mapping.resource_type}_{item.mapping.id}"
+
+    def stream_data(self, selector: AssetCentricMigrationSelector, limit: int | None = None) -> Iterator[Page]:
+        if isinstance(selector, MigrationCSVFileSelector):
+            iterator = self._stream_from_csv(selector, limit)
+        elif isinstance(selector, MigrateDataSetSelector):
+            iterator = self._stream_given_dataset(selector, limit)
+        else:
+            raise ToolkitNotImplementedError(f"Selector {type(selector)} is not supported for stream_data")
+        yield from (Page(worker_id="main", items=items) for items in iterator)
+
+    def _stream_from_csv(
+        self, selector: MigrationCSVFileSelector, limit: int | None = None
+    ) -> Iterator[Sequence[AssetCentricMapping[T_AssetCentricResource]]]:
+        items = selector.items
+        if limit is not None:
+            items = MigrationMappingList(items[:limit])
+        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):
+                chunk.append(AssetCentricMapping(mapping=mapping, resource=resource))
+            if chunk:
+                yield chunk
+                chunk = []
+
+    def count(self, selector: AssetCentricMigrationSelector) -> int | None:
+        if isinstance(selector, MigrationCSVFileSelector):
+            return len(selector.items)
+        elif isinstance(selector, MigrateDataSetSelector):
+            return self.hierarchy.count(selector.as_asset_centric_selector())
+        else:
+            raise ToolkitNotImplementedError(f"Selector {type(selector)} is not supported for count")
+
+    def _stream_given_dataset(
+        self, selector: MigrateDataSetSelector, limit: int | None = None
+    ) -> 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_AssetCentricResource]([])
+            for resource in data_chunk.items:
+                # 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
+                if instance_space is None:
+                    instance_space = MISSING_INSTANCE_SPACE
+                external_id = resource.external_id
+                if external_id is None:
+                    external_id = MISSING_EXTERNAL_ID.format(project=self.client.config.project, id=resource.id)
+                mapping = MigrationMapping(
+                    resource_type=self._kind_to_resource_type(selector.kind),
+                    instance_id=NodeId(
+                        space=instance_space,
+                        external_id=external_id,
+                    ),
+                    id=resource.id,
+                    data_set_id=resource.data_set_id,
+                    ingestion_view=selector.ingestion_mapping,
+                    preferred_consumer_view=selector.preferred_consumer_view,
+                )
+                mapping_list.append(AssetCentricMapping(mapping=mapping, resource=resource))
+            yield mapping_list
+
+    @staticmethod
+    def _kind_to_resource_type(kind: AssetCentricKindExtended) -> AssetCentricType:
+        mapping: dict[AssetCentricKindExtended, AssetCentricType] = {
+            "Assets": "asset",
+            "Events": "event",
+            "TimeSeries": "timeseries",
+            "FileMetadata": "file",
+        }
+        try:
+            return mapping[kind]
+        except KeyError as e:
+            raise ToolkitNotImplementedError(f"Kind '{kind}' is not supported") from e
+
+    def data_to_json_chunk(
+        self,
+        data_chunk: Sequence[AssetCentricMapping[T_AssetCentricResource]],
+        selector: AssetCentricMigrationSelector | None = None,
+    ) -> list[dict[str, JsonVal]]:
+        return [item.dump() for item in data_chunk]
+
+    def json_to_resource(self, item_json: dict[str, JsonVal]) -> InstanceApply:
+        raise NotImplementedError()
+
+    def upload_items(
+        self,
+        data_chunk: Sequence[UploadItem[InstanceApply]],
+        http_client: HTTPClient,
+        selector: AssetCentricMigrationSelector | None = None,
+    ) -> Sequence[HTTPMessage]:
+        """Upload items by first linking them using files/set-pending-instance-ids and then uploading the instances."""
+        if self.skip_linking:
+            return list(super().upload_items(data_chunk, http_client, None))
+        elif selector is None:
+            raise ToolkitNotImplementedError(f"Selector must be provided for uploading {self.KIND} items.")
+        elif selector.kind not in self.PENDING_INSTANCE_ID_ENDPOINT_BY_KIND:
+            return list(super().upload_items(data_chunk, http_client, None))
+
+        pending_instance_id_endpoint = self.PENDING_INSTANCE_ID_ENDPOINT_BY_KIND[selector.kind]
+        results: list[HTTPMessage] = []
+        to_upload = self.link_asset_centric(data_chunk, http_client, results, pending_instance_id_endpoint)
+        if to_upload:
+            results.extend(list(super().upload_items(to_upload, http_client, None)))
+        return results
+
+    @classmethod
+    def link_asset_centric(
+        cls,
+        data_chunk: Sequence[UploadItem[InstanceApply]],
+        http_client: HTTPClient,
+        results: list[HTTPMessage],
+        pending_instance_id_endpoint: str,
+    ) -> Sequence[UploadItem[InstanceApply]]:
+        """Links asset-centric resources to their (uncreated) instances using the pending-instance-ids endpoint."""
+        config = http_client.config
+        successful_linked: set[str] = set()
+        for batch in chunker_sequence(data_chunk, cls.CHUNK_SIZE):
+            batch_results = http_client.request_with_retries(
+                message=ItemsRequest(
+                    endpoint_url=config.create_api_url(pending_instance_id_endpoint),
+                    method="POST",
+                    api_version="alpha",
+                    items=[
+                        UploadItem(source_id=item.source_id, item=cls.as_pending_instance_id(item.item))
+                        for item in batch
+                    ],
+                )
+            )
+            for res in batch_results:
+                if isinstance(res, SuccessResponseItems):
+                    successful_linked.update(res.ids)
+            results.extend(batch_results)
+        to_upload = [item for item in data_chunk if item.source_id in successful_linked]
+        return to_upload
+
+    @staticmethod
+    def as_pending_instance_id(item: InstanceApply) -> PendingInstanceId:
+        """Convert an InstanceApply to a PendingInstanceId for linking."""
+        source = next((source for source in item.sources if source.source == INSTANCE_SOURCE_VIEW_ID), None)
+        if source is None:
+            raise ValueError(f"Cannot extract ID from item of type {type(item).__name__!r}")
+        if not isinstance(source.properties["id"], int):
+            raise ValueError(f"Unexpected ID type: {type(source.properties['id']).__name__!r}")
+        id_ = source.properties["id"]
+        return PendingInstanceId(
+            pending_instance_id=NodeId(item.space, item.external_id),
+            id=id_,
+        )
+
+
+class AnnotationMigrationIO(
+    UploadableStorageIO[AssetCentricMigrationSelector, AssetCentricMapping[Annotation], InstanceApply]
+):
+    """IO class for migrating Annotations.
+
+    Args:
+        client: The ToolkitClient to use for CDF interactions.
+        instance_space: The instance space to use for the migrated annotations.
+        default_asset_annotation_mapping: The default ingestion mapping to use for asset-linked annotations.
+        default_file_annotation_mapping: The default ingestion mappingto use for file-linked annotations.
+
+    """
+
+    KIND = "AnnotationMigration"
+    SUPPORTED_DOWNLOAD_FORMATS = frozenset({".parquet", ".csv", ".ndjson"})
+    SUPPORTED_COMPRESSIONS = frozenset({".gz"})
+    SUPPORTED_READ_FORMATS = frozenset({".parquet", ".csv", ".ndjson", ".yaml", ".yml"})
+    CHUNK_SIZE = 1000
+    UPLOAD_ENDPOINT = InstanceIO.UPLOAD_ENDPOINT
+
+    SUPPORTED_ANNOTATION_TYPES = frozenset({"diagrams.AssetLink", "diagrams.FileLink"})
+
+    def __init__(
+        self,
+        client: ToolkitClient,
+        instance_space: str | None = None,
+        default_asset_annotation_mapping: str | None = None,
+        default_file_annotation_mapping: str | None = None,
+    ) -> None:
+        super().__init__(client)
+        self.annotation_io = AnnotationIO(client)
+        self.instance_space = instance_space
+        self.default_asset_annotation_mapping = default_asset_annotation_mapping or ASSET_ANNOTATIONS_ID
+        self.default_file_annotation_mapping = default_file_annotation_mapping or FILE_ANNOTATIONS_ID
+
+    def as_id(self, item: AssetCentricMapping[Annotation]) -> str:
+        return f"Annotation_{item.mapping.id}"
+
+    def count(self, selector: AssetCentricMigrationSelector) -> int | None:
+        if isinstance(selector, MigrationCSVFileSelector):
+            return len(selector.items)
+        else:
+            # There is no efficient way to count annotations in CDF.
+            return None
+
+    def stream_data(self, selector: AssetCentricMigrationSelector, limit: int | None = None) -> Iterable[Page]:
+        if isinstance(selector, MigrateDataSetSelector):
+            iterator = self._stream_from_dataset(selector, limit)
+        elif isinstance(selector, MigrationCSVFileSelector):
+            iterator = self._stream_from_csv(selector, limit)
+        else:
+            raise ToolkitNotImplementedError(f"Selector {type(selector)} is not supported for stream_data")
+        yield from (Page(worker_id="main", items=items) for items in iterator)
+
+    def _stream_from_dataset(
+        self, selector: MigrateDataSetSelector, limit: int | None = None
+    ) -> Iterator[Sequence[AssetCentricMapping[Annotation]]]:
+        if self.instance_space is None:
+            raise ToolkitValueError("Instance space must be provided for dataset-based annotation migration.")
+        asset_centric_selector = selector.as_asset_centric_selector()
+        for data_chunk in self.annotation_io.stream_data(asset_centric_selector, limit):
+            mapping_list = AssetCentricMappingList[Annotation]([])
+            for resource in data_chunk.items:
+                if resource.annotation_type not in self.SUPPORTED_ANNOTATION_TYPES:
+                    # This should not happen, as the annotation_io should already filter these out.
+                    # This is just in case.
+                    continue
+                mapping = AnnotationMapping(
+                    instance_id=EdgeId(space=self.instance_space, external_id=f"annotation_{resource.id!r}"),
+                    id=resource.id,
+                    ingestion_view=self._get_mapping(selector.ingestion_mapping, resource),
+                    preferred_consumer_view=selector.preferred_consumer_view,
+                    # The PySDK is poorly typed.
+                    annotation_type=resource.annotation_type,  # type: ignore[arg-type]
+                )
+                mapping_list.append(AssetCentricMapping(mapping=mapping, resource=resource))
+            yield mapping_list
+
+    def _stream_from_csv(
+        self, selector: MigrationCSVFileSelector, limit: int | None = None
+    ) -> Iterator[Sequence[AssetCentricMapping[Annotation]]]:
+        items = selector.items
+        if limit is not None:
+            items = MigrationMappingList(items[:limit])
+        chunk: list[AssetCentricMapping[Annotation]] = []
+        for current_batch in chunker_sequence(items, self.CHUNK_SIZE):
+            resources = self.client.annotations.retrieve_multiple(current_batch.get_ids())
+            resources_by_id = {resource.id: resource for resource in resources}
+            not_found = 0
+            incorrect_type_count = 0
+            for mapping in current_batch:
+                resource = resources_by_id.get(mapping.id)
+                if resource is None:
+                    not_found += 1
+                    continue
+                if resource.annotation_type not in self.SUPPORTED_ANNOTATION_TYPES:
+                    incorrect_type_count += 1
+                    continue
+                mapping.ingestion_view = self._get_mapping(mapping.ingestion_view, resource)
+                chunk.append(AssetCentricMapping(mapping=mapping, resource=resource))
+            if chunk:
+                yield chunk
+                chunk = []
+            if not_found:
+                MediumSeverityWarning(
+                    f"Could not find {not_found} annotations referenced in the CSV file. They will be skipped during migration."
+                ).print_warning(include_timestamp=True, console=self.client.console)
+            if incorrect_type_count:
+                MediumSeverityWarning(
+                    f"Found {incorrect_type_count} annotations with unsupported types. Only 'diagrams.AssetLink' and "
+                    "'diagrams.FileLink' are supported. These annotations will be skipped during migration."
+                ).print_warning(include_timestamp=True, console=self.client.console)
+
+    def _get_mapping(self, current_mapping: str | None, resource: Annotation) -> str:
+        try:
+            return (
+                current_mapping
+                or {
+                    "diagrams.AssetLink": self.default_asset_annotation_mapping,
+                    "diagrams.FileLink": self.default_file_annotation_mapping,
+                }[resource.annotation_type]
+            )
+        except KeyError as e:
+            raise ToolkitValueError(
+                f"Could not determine default ingestion view for annotation type '{resource.annotation_type}'. "
+                "Please specify the ingestion view explicitly in the CSV file."
+            ) from e
+
+    def json_to_resource(self, item_json: dict[str, JsonVal]) -> InstanceApply:
+        raise NotImplementedError("Deserializing Annotation Migrations from JSON is not supported.")
+
+    def data_to_json_chunk(
+        self,
+        data_chunk: Sequence[AssetCentricMapping[Annotation]],
+        selector: AssetCentricMigrationSelector | None = None,
+    ) -> list[dict[str, JsonVal]]:
+        raise NotImplementedError("Serializing Annotation Migrations to JSON is not supported.")

cognite_toolkit/_cdf_tk/commands/_migrate/prepare.py

@@ -6,12 +6,14 @@ from cognite_toolkit._cdf_tk.commands.deploy import DeployCommand
 from cognite_toolkit._cdf_tk.cruds import (
     ContainerCRUD,
     DataModelCRUD,
+    ResourceCRUD,
     ResourceViewMappingCRUD,
     ResourceWorker,
     SpaceCRUD,
     ViewCRUD,
 )
 from cognite_toolkit._cdf_tk.data_classes import DeployResults
+from cognite_toolkit._cdf_tk.tk_warnings import HighSeverityWarning
 
 from .data_model import COGNITE_MIGRATION_MODEL, CONTAINERS, MODEL_ID, SPACE, VIEWS
 from .default_mappings import create_default_mappings
@@ -27,27 +29,30 @@ class MigrationPrepareCommand(ToolkitCommand):
         verb = "Would deploy" if dry_run else "Deploying"
         print(f"{verb} {MODEL_ID!r}")
         results = DeployResults([], "deploy", dry_run=dry_run)
-        for loader_cls, resource_list in [
+        crud_cls: type[ResourceCRUD]
+        for crud_cls, resource_list in [  # type: ignore[assignment]
             (SpaceCRUD, [SPACE]),
             (ContainerCRUD, CONTAINERS),
             (ViewCRUD, VIEWS),
             (DataModelCRUD, [COGNITE_MIGRATION_MODEL]),
             (ResourceViewMappingCRUD, create_default_mappings()),
         ]:
-            # MyPy does not understand that `loader_cls` has a `create_loader` method.
-            loader = loader_cls.create_loader(client)  # type: ignore[attr-defined]
-            worker = ResourceWorker(loader, "deploy")
+            crud = crud_cls.create_loader(client)
+            if warning := crud.prerequisite_warning():
+                self.warn(HighSeverityWarning(warning))
+                continue
+            worker = ResourceWorker(crud, "deploy")
             # MyPy does not understand that `loader` has a `get_id` method.
-            dump_arg = {"context": "local"} if loader_cls is ResourceViewMappingCRUD else {}
-            local_by_id = {loader.get_id(item): (item.dump(**dump_arg), item) for item in resource_list}  # type: ignore[attr-defined]
+            dump_arg = {"context": "local"} if crud_cls is ResourceViewMappingCRUD else {}
+            local_by_id = {crud.get_id(item): (item.dump(**dump_arg), item) for item in resource_list}  # type: ignore[attr-defined]
             worker.validate_access(local_by_id, is_dry_run=dry_run)
-            cdf_resources = loader.retrieve(list(local_by_id.keys()))
+            cdf_resources = crud.retrieve(list(local_by_id.keys()))
             resources = worker.categorize_resources(local_by_id, cdf_resources, False, verbose)
 
             if dry_run:
-                result = deploy_cmd.dry_run_deploy(resources, loader, False, False)
+                result = deploy_cmd.dry_run_deploy(resources, crud, False, False)
             else:
-                result = deploy_cmd.actual_deploy(resources, loader)
+                result = deploy_cmd.actual_deploy(resources, crud)
             if result:
                 results[result.name] = result
         if results.has_counts:

cognite_toolkit/_cdf_tk/commands/_migrate/selectors.py

@@ -0,0 +1,71 @@
+from abc import ABC, abstractmethod
+from functools import cached_property
+from pathlib import Path
+from typing import Literal
+
+from cognite.client.data_classes.data_modeling import ViewId
+
+from cognite_toolkit._cdf_tk.commands._migrate.data_classes import MigrationMappingList
+from cognite_toolkit._cdf_tk.storageio import DataSelector
+from cognite_toolkit._cdf_tk.storageio.selectors import DataSetSelector
+from cognite_toolkit._cdf_tk.utils.useful_types import AssetCentricKindExtended
+
+
+class AssetCentricMigrationSelector(DataSelector, ABC):
+    kind: AssetCentricKindExtended
+
+    @abstractmethod
+    def get_ingestion_mappings(self) -> list[str]:
+        raise NotImplementedError()
+
+
+class MigrationCSVFileSelector(AssetCentricMigrationSelector):
+    type: Literal["migrationCSVFile"] = "migrationCSVFile"
+    datafile: Path
+
+    @property
+    def group(self) -> str:
+        return f"Migration_{self.kind}"
+
+    def __str__(self) -> str:
+        return f"file_{self.datafile.name}"
+
+    def get_ingestion_mappings(self) -> list[str]:
+        views = {item.get_ingestion_view() for item in self.items}
+        return sorted(views)
+
+    @cached_property
+    def items(self) -> MigrationMappingList:
+        return MigrationMappingList.read_csv_file(self.datafile, resource_type=self.kind)
+
+
+class MigrateDataSetSelector(AssetCentricMigrationSelector):
+    type: Literal["migrateDataSet"] = "migrateDataSet"
+    kind: AssetCentricKindExtended
+    data_set_external_id: str
+    ingestion_mapping: str | None = None
+    preferred_consumer_view: ViewId | None = None
+
+    @property
+    def group(self) -> str:
+        return f"DataSet_{self.data_set_external_id}"
+
+    def __str__(self) -> str:
+        return self.kind
+
+    def get_schema_spaces(self) -> list[str] | None:
+        return None
+
+    def get_instance_spaces(self) -> list[str] | None:
+        return None
+
+    def get_ingestion_mappings(self) -> list[str]:
+        return [self.ingestion_mapping] if self.ingestion_mapping else []
+
+    def as_asset_centric_selector(self) -> DataSetSelector:
+        if self.kind == "Annotations":
+            # Annotations are connected to file metadata, so we need to download the file metadata
+            # and look up the annotations connected to each file metadata.
+            return DataSetSelector(data_set_external_id=self.data_set_external_id, kind="FileMetadata")
+        else:
+            return DataSetSelector(data_set_external_id=self.data_set_external_id, kind=self.kind)