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

cognite_toolkit/_cdf_tk/client/api/three_d.py

@@ -1,14 +1,23 @@
-from collections.abc import Sequence
+from collections import defaultdict
+from collections.abc import Iterable, Sequence
+from typing import Any, TypeVar
 
+from pydantic import TypeAdapter
 from rich.console import Console
 
-from cognite_toolkit._cdf_tk.client.data_classes.api_classes import PagedResponse
-from cognite_toolkit._cdf_tk.client.data_classes.three_d import ThreeDModelClassicRequest, ThreeDModelResponse
+from cognite_toolkit._cdf_tk.client.data_classes.api_classes import InternalIdRequest, PagedResponse
+from cognite_toolkit._cdf_tk.client.data_classes.three_d import (
+    AssetMappingClassicRequest,
+    AssetMappingDMRequest,
+    AssetMappingResponse,
+    ThreeDModelClassicRequest,
+    ThreeDModelResponse,
+)
+from cognite_toolkit._cdf_tk.utils.collection import chunker_sequence
 from cognite_toolkit._cdf_tk.utils.http_client import (
     HTTPClient,
-    ItemsRequest,
+    ItemsRequest2,
     RequestMessage2,
-    SimpleBodyRequest,
 )
 from cognite_toolkit._cdf_tk.utils.useful_types import PrimitiveType
 
@@ -37,16 +46,15 @@ class ThreeDModelAPI:
             return []
         if len(models) > self.MAX_CLASSIC_MODELS_PER_CREATE_REQUEST:
             raise ValueError("Cannot create more than 1000 3D models in a single request.")
-        responses = self._http_client.request_with_retries(
-            ItemsRequest(
+        responses = self._http_client.request_items_retries(
+            ItemsRequest2(
                 endpoint_url=self._config.create_api_url(self.ENDPOINT),
                 method="POST",
-                items=list(models),
+                items=models,
             )
         )
         responses.raise_for_status()
-        body = responses.get_first_body()
-        return PagedResponse[ThreeDModelResponse].model_validate(body).items
+        return TypeAdapter(list[ThreeDModelResponse]).validate_python(responses.get_items())
 
     def delete(self, ids: Sequence[int]) -> None:
         """Delete 3D models by their IDs.
@@ -58,11 +66,11 @@ class ThreeDModelAPI:
             return None
         if len(ids) > self.MAX_MODELS_PER_DELETE_REQUEST:
             raise ValueError("Cannot delete more than 1000 3D models in a single request.")
-        responses = self._http_client.request_with_retries(
-            SimpleBodyRequest(
+        responses = self._http_client.request_items_retries(
+            ItemsRequest2(
                 endpoint_url=self._config.create_api_url(self.ENDPOINT + "/delete"),
                 method="POST",
-                body_content={"items": [{"id": id_} for id_ in ids]},
+                items=InternalIdRequest.from_ids(list(ids)),
             )
         )
         responses.raise_for_status()
@@ -126,6 +134,251 @@ class ThreeDModelAPI:
         return results
 
 
+T_RequestMapping = TypeVar("T_RequestMapping", bound=AssetMappingClassicRequest | AssetMappingDMRequest)
+
+
+class ThreeDAssetMappingAPI:
+    ENDPOINT = "/3d/models/{modelId}/revisions/{revisionId}/mappings"
+    CREATE_CLASSIC_MAX_MAPPINGS_PER_REQUEST = 1000
+    CREATE_DM_MAX_MAPPINGS_PER_REQUEST = 100
+    DELETE_CLASSIC_MAX_MAPPINGS_PER_REQUEST = 1000
+    DELETE_DM_MAX_MAPPINGS_PER_REQUEST = 100
+    LIST_REQUEST_MAX_LIMIT = 1000
+
+    def __init__(self, http_client: HTTPClient, console: Console) -> None:
+        self._http_client = http_client
+        self._console = console
+        self._config = http_client.config
+
+    def create(self, mappings: Sequence[AssetMappingClassicRequest]) -> list[AssetMappingResponse]:
+        """Create 3D asset mappings.
+
+        Args:
+            mappings (Sequence[AssetMappingClassicRequest]):
+                The 3D asset mapping(s) to create.
+
+        Returns:
+            list[AssetMappingResponse]: The created 3D asset mapping(s).
+        """
+        results: list[AssetMappingResponse] = []
+        for endpoint, model_id, revision_id, revision_mappings in self._chunk_mappings_by_endpoint(
+            mappings, self.CREATE_CLASSIC_MAX_MAPPINGS_PER_REQUEST
+        ):
+            responses = self._http_client.request_items_retries(
+                ItemsRequest2(
+                    endpoint_url=self._config.create_api_url(endpoint),
+                    method="POST",
+                    items=revision_mappings,
+                )
+            )
+            responses.raise_for_status()
+            items = responses.get_items()
+            for item in items:
+                # We append modelId and revisionId to each item since the API does not return them
+                # this is needed to fully populate the AssetMappingResponse data class
+                item["modelId"] = model_id
+                item["revisionId"] = revision_id
+            results.extend(TypeAdapter(list[AssetMappingResponse]).validate_python(items))
+        return results
+
+    def create_dm(
+        self, mappings: Sequence[AssetMappingDMRequest], object_3d_space: str, cad_node_space: str
+    ) -> list[AssetMappingResponse]:
+        """Create 3D asset mappings in Data Modeling format.
+
+        Args:
+            mappings (Sequence[AssetMappingDMRequest]):
+                The 3D asset mapping(s) to create
+            object_3d_space (str):
+                The instance space where the Cognite3DObject are located.
+            cad_node_space (str):
+                The instance space where the CogniteCADNode are located.
+        Returns:
+            list[AssetMappingResponse]: The created 3D asset mapping(s).
+        """
+        results: list[AssetMappingResponse] = []
+        for endpoint, model_id, revision_id, revision_mappings in self._chunk_mappings_by_endpoint(
+            mappings, self.CREATE_DM_MAX_MAPPINGS_PER_REQUEST
+        ):
+            responses = self._http_client.request_items_retries(
+                ItemsRequest2(
+                    endpoint_url=self._config.create_api_url(endpoint),
+                    method="POST",
+                    items=revision_mappings,
+                    extra_body_fields={
+                        "dmsContextualizationConfig": {
+                            "object3DSpace": object_3d_space,
+                            "cadNodeSpace": cad_node_space,
+                        }
+                    },
+                )
+            )
+            responses.raise_for_status()
+            items = responses.get_items()
+            for item in items:
+                # We append modelId and revisionId to each item since the API does not return them
+                # this is needed to fully populate the AssetMappingResponse data class
+                item["modelId"] = model_id
+                item["revisionId"] = revision_id
+            results.extend(TypeAdapter(list[AssetMappingResponse]).validate_python(items))
+        return results
+
+    @classmethod
+    def _chunk_mappings_by_endpoint(
+        cls, mappings: Sequence[T_RequestMapping], chunk_size: int
+    ) -> Iterable[tuple[str, int, int, list[T_RequestMapping]]]:
+        chunked_mappings: dict[tuple[int, int], list[T_RequestMapping]] = defaultdict(list)
+        for mapping in mappings:
+            key = mapping.model_id, mapping.revision_id
+            chunked_mappings[key].append(mapping)
+        for (model_id, revision_id), revision_mappings in chunked_mappings.items():
+            endpoint = cls.ENDPOINT.format(modelId=model_id, revisionId=revision_id)
+            for chunk in chunker_sequence(revision_mappings, chunk_size):
+                yield endpoint, model_id, revision_id, chunk
+
+    def delete(self, mappings: Sequence[AssetMappingClassicRequest]) -> None:
+        """Delete 3D asset mappings.
+
+        Args:
+            mappings (Sequence[AssetMappingClassicRequest]):
+                The 3D asset mapping(s) to delete.
+        """
+        for endpoint, *_, revision_mappings in self._chunk_mappings_by_endpoint(
+            mappings, self.DELETE_CLASSIC_MAX_MAPPINGS_PER_REQUEST
+        ):
+            responses = self._http_client.request_items_retries(
+                ItemsRequest2(
+                    endpoint_url=self._config.create_api_url(f"{endpoint}/delete"),
+                    method="DELETE",
+                    items=revision_mappings,
+                )
+            )
+            responses.raise_for_status()
+        return None
+
+    def delete_dm(self, mappings: Sequence[AssetMappingDMRequest], object_3d_space: str, cad_node_space: str) -> None:
+        """Delete 3D asset mappings in Data Modeling format.
+
+        Args:
+            mappings (Sequence[AssetMappingDMRequest]):
+                The 3D asset mapping(s) to delete.
+            object_3d_space (str):
+                The instance space where the Cognite3DObject are located.
+            cad_node_space (str):
+                The instance space where the CogniteCADNode are located.
+        """
+        for endpoint, *_, revision_mappings in self._chunk_mappings_by_endpoint(
+            mappings, self.DELETE_DM_MAX_MAPPINGS_PER_REQUEST
+        ):
+            responses = self._http_client.request_items_retries(
+                ItemsRequest2(
+                    endpoint_url=self._config.create_api_url(f"{endpoint}/delete"),
+                    method="DELETE",
+                    items=revision_mappings,
+                    extra_body_fields={
+                        "dmsContextualizationConfig": {
+                            "object3DSpace": object_3d_space,
+                            "cadNodeSpace": cad_node_space,
+                        }
+                    },
+                )
+            )
+            responses.raise_for_status()
+        return None
+
+    def iterate(
+        self,
+        model_id: int,
+        revision_id: int,
+        asset_ids: list[int] | None = None,
+        asset_instance_ids: list[str] | None = None,
+        node_ids: list[int] | None = None,
+        tree_indexes: list[int] | None = None,
+        get_dms_instances: bool = False,
+        limit: int = 100,
+        cursor: str | None = None,
+    ) -> PagedResponse[AssetMappingResponse]:
+        if not (0 < limit <= self.LIST_REQUEST_MAX_LIMIT):
+            raise ValueError(f"Limit must be between 1 and {self.LIST_REQUEST_MAX_LIMIT}, got {limit}.")
+        if sum(param is not None for param in [asset_ids, asset_instance_ids, node_ids, tree_indexes]) > 1:
+            raise ValueError("Only one of asset_ids, asset_instance_ids, node_ids, or tree_indexes can be provided.")
+        body: dict[str, Any] = {
+            "getDmsInstances": get_dms_instances,
+            "limit": limit,
+        }
+        if asset_ids is not None:
+            if not (0 < len(asset_ids) <= 100):
+                raise ValueError("asset_ids must contain between 1 and 100 IDs.")
+            body["filter"] = {"assetIds": asset_ids}
+        elif asset_instance_ids is not None:
+            if not (0 < len(asset_instance_ids) <= 100):
+                raise ValueError("asset_instance_ids must contain between 1 and 100 IDs.")
+            body["filter"] = {"assetInstanceIds": asset_instance_ids}
+        elif node_ids is not None:
+            if not (0 < len(node_ids) <= 100):
+                raise ValueError("node_ids must contain between 1 and 100 IDs.")
+            body["filter"] = {"nodeIds": node_ids}
+        elif tree_indexes is not None:
+            if not (0 < len(tree_indexes) <= 100):
+                raise ValueError("tree_indexes must contain between 1 and 100 indexes.")
+            body["filter"] = {"treeIndexes": tree_indexes}
+        if cursor is not None:
+            body["cursor"] = cursor
+
+        endpoint = self.ENDPOINT.format(modelId=model_id, revisionId=revision_id)
+        responses = self._http_client.request_single_retries(
+            RequestMessage2(
+                endpoint_url=self._config.create_api_url(f"{endpoint}/list"),
+                method="POST",
+                body_content=body,
+            )
+        )
+        success_response = responses.get_success_or_raise()
+        body_json = success_response.body_json
+        # Add modelId and revisionId to items since the API does not return them
+        for item in body_json.get("items", []):
+            item["modelId"] = model_id
+            item["revisionId"] = revision_id
+        return PagedResponse[AssetMappingResponse].model_validate(body_json)
+
+    def list(
+        self,
+        model_id: int,
+        revision_id: int,
+        asset_ids: list[int] | None = None,
+        asset_instance_ids: list[str] | None = None,
+        node_ids: list[int] | None = None,
+        tree_indexes: list[int] | None = None,
+        get_dms_instances: bool = False,
+        limit: int | None = 100,
+    ) -> list[AssetMappingResponse]:
+        results: list[AssetMappingResponse] = []
+        cursor: str | None = None
+        while True:
+            request_limit = (
+                self.LIST_REQUEST_MAX_LIMIT if limit is None else min(limit - len(results), self.LIST_REQUEST_MAX_LIMIT)
+            )
+            if request_limit <= 0:
+                break
+            page = self.iterate(
+                model_id=model_id,
+                revision_id=revision_id,
+                asset_ids=asset_ids,
+                asset_instance_ids=asset_instance_ids,
+                node_ids=node_ids,
+                tree_indexes=tree_indexes,
+                get_dms_instances=get_dms_instances,
+                limit=request_limit,
+                cursor=cursor,
+            )
+            results.extend(page.items)
+            if page.next_cursor is None:
+                break
+            cursor = page.next_cursor
+        return results
+
+
 class ThreeDAPI:
     def __init__(self, http_client: HTTPClient, console: Console) -> None:
         self.models = ThreeDModelAPI(http_client, console)
+        self.asset_mappings = ThreeDAssetMappingAPI(http_client, console)

cognite_toolkit/_cdf_tk/client/data_classes/api_classes.py

@@ -2,6 +2,8 @@ from typing import Generic, TypeVar
 
 from pydantic import BaseModel, Field, JsonValue
 
+from cognite_toolkit._cdf_tk.utils.http_client._data_classes2 import RequestResource
+
 T = TypeVar("T", bound=BaseModel)
 
 
@@ -15,3 +17,14 @@ class QueryResponse(BaseModel, Generic[T]):
     typing: dict[str, JsonValue] | None = None
     next_cursor: dict[str, str] = Field(alias="nextCursor")
     debug: dict[str, JsonValue] | None = None
+
+
+class InternalIdRequest(RequestResource):
+    id: int
+
+    def as_id(self) -> int:
+        return self.id
+
+    @classmethod
+    def from_ids(cls, ids: list[int]) -> list["InternalIdRequest"]:
+        return [cls(id=id_) for id_ in ids]

cognite_toolkit/_cdf_tk/client/data_classes/base.py

@@ -3,7 +3,7 @@ from abc import ABC, abstractmethod
 from collections import UserList
 from typing import TYPE_CHECKING, Any, Generic, TypeVar
 
-from pydantic import BaseModel, ConfigDict
+from pydantic import ConfigDict
 from pydantic.alias_generators import to_camel
 
 from cognite_toolkit._cdf_tk.utils.http_client._data_classes2 import BaseModelObject, RequestResource
@@ -31,21 +31,17 @@ class ResponseResource(BaseModelObject, Generic[T_RequestResource], ABC):
         return self.as_request_resource()
 
 
-class Identifier(BaseModel):
+class Identifier(RequestResource, ABC):
     """Base class for all identifier classes."""
 
     model_config = ConfigDict(alias_generator=to_camel, extra="ignore", populate_by_name=True, frozen=True)
 
-    def dump(self, include_type: bool = True) -> dict[str, Any]:
-        """Dump the identifier to a dictionary.
+    def dump(self, camel_case: bool = True, include_type: bool = True) -> dict[str, Any]:
+        """Dump the resource to a dictionary.
 
-        Args:
-            include_type (bool): Whether to include the type of the identifier in the output.
-
-        Returns:
-            dict[str, Any]: The dumped identifier.
+        This is the default serialization method for request resources.
         """
-        return self.model_dump(mode="json", by_alias=True, exclude_defaults=not include_type)
+        return self.model_dump(mode="json", by_alias=camel_case, exclude_unset=not include_type)
 
     def as_id(self) -> Self:
         return self

cognite_toolkit/_cdf_tk/client/data_classes/instance_api.py

@@ -15,19 +15,19 @@ class TypedInstanceIdentifier(Identifier):
     external_id: str
 
 
-class InstanceIdentifier(Identifier):
-    space: str
-    external_id: str
-
-
-class NodeIdentifier(TypedInstanceIdentifier):
+class TypedNodeIdentifier(TypedInstanceIdentifier):
     instance_type: Literal["node"] = "node"
 
 
-class EdgeIdentifier(TypedInstanceIdentifier):
+class TypedEdgeIdentifier(TypedInstanceIdentifier):
     instance_type: Literal["edge"] = "edge"
 
 
+class InstanceIdentifier(Identifier):
+    space: str
+    external_id: str
+
+
 class InstanceResult(BaseModelObject):
     instance_type: InstanceType
     version: int

cognite_toolkit/_cdf_tk/client/data_classes/{canvas.py → legacy/canvas.py}

@@ -30,7 +30,7 @@ from cognite.client.data_classes.data_modeling.instances import (
     TypedNodeApply,
 )
 
-from cognite_toolkit._cdf_tk.client.data_classes.migration import AssetCentricId
+from cognite_toolkit._cdf_tk.client.data_classes.legacy.migration import AssetCentricId
 from cognite_toolkit._cdf_tk.utils.useful_types import JsonVal
 
 if sys.version_info >= (3, 11):

cognite_toolkit/_cdf_tk/client/data_classes/three_d.py

@@ -1,5 +1,8 @@
+from collections.abc import Hashable
 from typing import Literal
 
+from pydantic import Field
+
 from .base import BaseModelObject, RequestResource, ResponseResource
 
 
@@ -51,3 +54,59 @@ class ThreeDModelResponse(ResponseResource[ThreeDModelRequest]):
             return ThreeDModelClassicRequest._load(self.dump())
         else:
             return ThreeDModelDMSRequest._load(self.dump())
+
+
+class AssetMappingDMRequest(RequestResource):
+    node_id: int
+    asset_instance_id: NodeReference
+    # These fields are part of the path request and not the body schema.
+    model_id: int = Field(exclude=True)
+    revision_id: int = Field(exclude=True)
+
+    def as_id(self) -> Hashable:
+        return (
+            self.model_id,
+            self.revision_id,
+            self.node_id,
+            self.asset_instance_id.space,
+            self.asset_instance_id.external_id,
+        )
+
+
+class AssetMappingClassicRequest(RequestResource):
+    node_id: int
+    asset_id: int | None = None
+    asset_instance_id: NodeReference | None = None
+    # These fields are part of the path request and not the body schema.
+    model_id: int = Field(exclude=True)
+    revision_id: int = Field(exclude=True)
+
+    def as_id(self) -> Hashable:
+        if self.asset_id:
+            return self.model_id, self.revision_id, self.node_id, self.asset_id
+        elif self.asset_instance_id:
+            return (
+                self.model_id,
+                self.revision_id,
+                self.node_id,
+                self.asset_instance_id.space,
+                self.asset_instance_id.external_id,
+            )
+        else:
+            raise AttributeError("asset_id or asset_instance_id is required")
+
+
+class AssetMappingResponse(ResponseResource[AssetMappingClassicRequest]):
+    node_id: int
+    asset_id: int | None = None
+    asset_instance_id: NodeReference | None = None
+    tree_index: int | None = None
+    subtree_size: int | None = None
+    # These fields are part of the path request and response, but they are included here for convenience.
+    model_id: int = Field(exclude=True)
+    revision_id: int = Field(exclude=True)
+
+    def as_request_resource(self) -> AssetMappingClassicRequest:
+        return AssetMappingClassicRequest.model_validate(
+            {**self.dump(), "modelId": self.model_id, "revisionId": self.revision_id}
+        )

cognite_toolkit/_cdf_tk/client/testing.py

@@ -12,18 +12,27 @@ from cognite.client.testing import CogniteClientMock
 from rich.console import Console
 
 from cognite_toolkit._cdf_tk.client._toolkit_client import ToolkitClient
+from cognite_toolkit._cdf_tk.client.api.legacy.canvas import CanvasAPI, IndustrialCanvasAPI
+from cognite_toolkit._cdf_tk.client.api.legacy.charts import ChartsAPI
+from cognite_toolkit._cdf_tk.client.api.legacy.dml import DMLAPI
+from cognite_toolkit._cdf_tk.client.api.legacy.extended_data_modeling import ExtendedInstancesAPI
+from cognite_toolkit._cdf_tk.client.api.legacy.extended_files import ExtendedFileMetadataAPI
+from cognite_toolkit._cdf_tk.client.api.legacy.extended_functions import ExtendedFunctionsAPI
+from cognite_toolkit._cdf_tk.client.api.legacy.extended_raw import ExtendedRawAPI
+from cognite_toolkit._cdf_tk.client.api.legacy.extended_timeseries import ExtendedTimeSeriesAPI
+from cognite_toolkit._cdf_tk.client.api.legacy.location_filters import LocationFiltersAPI
+from cognite_toolkit._cdf_tk.client.api.legacy.robotics import (
+    CapabilitiesAPI,
+    DataPostProcessingAPI,
+    FramesAPI,
+    MapsAPI,
+    RoboticsAPI,
+)
+from cognite_toolkit._cdf_tk.client.api.legacy.robotics import LocationsAPI as RoboticsLocationsAPI
+from cognite_toolkit._cdf_tk.client.api.legacy.search_config import SearchConfigurationsAPI
 
 from ._toolkit_client import ToolAPI
-from .api.canvas import CanvasAPI, IndustrialCanvasAPI
-from .api.charts import ChartsAPI
-from .api.dml import DMLAPI
-from .api.extended_data_modeling import ExtendedInstancesAPI
-from .api.extended_files import ExtendedFileMetadataAPI
-from .api.extended_functions import ExtendedFunctionsAPI
-from .api.extended_raw import ExtendedRawAPI
-from .api.extended_timeseries import ExtendedTimeSeriesAPI
 from .api.infield import InfieldAPI, InFieldCDMConfigAPI, InfieldConfigAPI
-from .api.location_filters import LocationFiltersAPI
 from .api.lookup import (
     AssetLookUpAPI,
     DataSetLookUpAPI,
@@ -45,14 +54,7 @@ from .api.migration import (
     ResourceViewMappingAPI,
 )
 from .api.project import ProjectAPI
-from .api.robotics import RoboticsAPI
-from .api.robotics.capabilities import CapabilitiesAPI
-from .api.robotics.data_postprocessing import DataPostProcessingAPI
-from .api.robotics.frames import FramesAPI
-from .api.robotics.locations import LocationsAPI as RoboticsLocationsAPI
-from .api.robotics.maps import MapsAPI
 from .api.search import SearchAPI
-from .api.search_config import SearchConfigurationsAPI
 from .api.streams import StreamsAPI
 from .api.three_d import ThreeDAPI, ThreeDModelAPI
 from .api.token import TokenAPI

cognite_toolkit/_cdf_tk/commands/_migrate/conversion.py

@@ -15,7 +15,7 @@ from cognite.client.data_classes.data_modeling.views import ViewProperty
 from cognite.client.utils._identifier import InstanceId
 
 from cognite_toolkit._cdf_tk.client import ToolkitClient
-from cognite_toolkit._cdf_tk.client.data_classes.migration import (
+from cognite_toolkit._cdf_tk.client.data_classes.legacy.migration import (
     AssetCentricId,
     ResourceViewMappingApply,
 )

cognite_toolkit/_cdf_tk/commands/_migrate/creators.py

@@ -20,7 +20,7 @@ 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.client.data_classes.legacy.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

cognite_toolkit/_cdf_tk/commands/_migrate/data_classes.py

@@ -13,9 +13,9 @@ from pydantic import BaseModel, BeforeValidator, Field, field_validator, model_v
 
 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.client.data_classes.legacy.instances import InstanceApplyList
+from cognite_toolkit._cdf_tk.client.data_classes.legacy.migration import AssetCentricId
+from cognite_toolkit._cdf_tk.client.data_classes.legacy.pending_instances_ids import PendingInstanceId
 from cognite_toolkit._cdf_tk.commands._migrate.default_mappings import (
     ASSET_ANNOTATIONS_ID,
     FILE_ANNOTATIONS_ID,

cognite_toolkit/_cdf_tk/commands/_migrate/data_mapper.py

@@ -15,21 +15,27 @@ from cognite.client.data_classes.data_modeling import (
 from cognite.client.exceptions import CogniteException
 
 from cognite_toolkit._cdf_tk.client import ToolkitClient
-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.instance_api import InstanceIdentifier
-from cognite_toolkit._cdf_tk.client.data_classes.migration import ResourceViewMappingApply
-from cognite_toolkit._cdf_tk.client.data_classes.three_d import RevisionStatus, ThreeDModelResponse
+from cognite_toolkit._cdf_tk.client.data_classes.legacy.canvas import (
+    ContainerReferenceApply,
+    FdmInstanceContainerReferenceApply,
+    IndustrialCanvas,
+    IndustrialCanvasApply,
+)
+from cognite_toolkit._cdf_tk.client.data_classes.legacy.charts import Chart, ChartWrite
+from cognite_toolkit._cdf_tk.client.data_classes.legacy.migration import ResourceViewMappingApply
+from cognite_toolkit._cdf_tk.client.data_classes.three_d import (
+    AssetMappingDMRequest,
+    AssetMappingResponse,
+    NodeReference,
+    RevisionStatus,
+    ThreeDModelResponse,
+)
 from cognite_toolkit._cdf_tk.commands._migrate.conversion import DirectRelationCache, asset_centric_to_dm
 from cognite_toolkit._cdf_tk.commands._migrate.data_classes import (
     Model,
@@ -443,15 +449,15 @@ class ThreeDMapper(DataMapper[ThreeDSelector, ThreeDModelResponse, ThreeDMigrati
             return None, issue
 
         mapped_request = ThreeDMigrationRequest(
-            model_id=item.id,
+            modelId=item.id,
             type=model_type,
             space=instance_space,
             revision=ThreeDRevisionMigrationRequest(
                 space=instance_space,
                 type=model_type,
-                revision_id=last_revision_id,
+                revisionId=last_revision_id,
                 model=Model(
-                    instance_id=InstanceIdentifier(
+                    instanceId=InstanceIdentifier(
                         space=instance_space,
                         external_id=f"cog_3d_model_{item.id!s}",
                     )
@@ -469,3 +475,45 @@ class ThreeDMapper(DataMapper[ThreeDSelector, ThreeDModelResponse, ThreeDMigrati
             return "PointCloud"
         else:
             return None
+
+
+class ThreeDAssetMapper(DataMapper[ThreeDSelector, AssetMappingResponse, AssetMappingDMRequest]):
+    def __init__(self, client: ToolkitClient) -> None:
+        self.client = client
+
+    def map(
+        self, source: Sequence[AssetMappingResponse]
+    ) -> Sequence[tuple[AssetMappingDMRequest | None, MigrationIssue]]:
+        output: list[tuple[AssetMappingDMRequest | None, MigrationIssue]] = []
+        self._populate_cache(source)
+        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[AssetMappingResponse]) -> None:
+        asset_ids: set[int] = set()
+        for mapping in source:
+            if mapping.asset_id is not None:
+                asset_ids.add(mapping.asset_id)
+        self.client.migration.lookup.assets(list(asset_ids))
+
+    def _map_single_item(
+        self, item: AssetMappingResponse
+    ) -> tuple[AssetMappingDMRequest | None, ThreeDModelMigrationIssue]:
+        issue = ThreeDModelMigrationIssue(model_name=f"AssetMapping_{item.model_id}", model_id=item.model_id)
+        asset_instance_id = item.asset_instance_id
+        if item.asset_id and asset_instance_id is None:
+            asset_node_id = self.client.migration.lookup.assets(item.asset_id)
+            if asset_node_id is None:
+                issue.error_message.append(f"Missing asset instance for asset ID {item.asset_id!r}")
+                return None, issue
+            asset_instance_id = NodeReference(space=asset_node_id.space, externalId=asset_node_id.external_id)
+
+        if asset_instance_id is None:
+            issue.error_message.append("Neither assetInstanceId nor assetId provided for mapping.")
+            return None, issue
+        mapped_request = AssetMappingDMRequest(
+            modelId=item.model_id, revisionId=item.revision_id, nodeId=item.node_id, assetInstanceId=asset_instance_id
+        )
+        return mapped_request, issue