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

cognite_toolkit/_cdf_tk/commands/_migrate/conversion.py

@@ -1,6 +1,5 @@
-from collections.abc import Mapping, Set
-from dataclasses import dataclass
-from typing import Any, ClassVar, overload
+from collections.abc import Iterable, Mapping, Set
+from typing import Any, ClassVar, cast
 
 from cognite.client.data_classes import Annotation, Asset, Event, FileMetadata, TimeSeries
 from cognite.client.data_classes.data_modeling import (
@@ -13,8 +12,13 @@ from cognite.client.data_classes.data_modeling import (
 )
 from cognite.client.data_classes.data_modeling.instances import EdgeApply, NodeOrEdgeData, PropertyValueWrite
 from cognite.client.data_classes.data_modeling.views import ViewProperty
+from cognite.client.utils._identifier import InstanceId
 
-from cognite_toolkit._cdf_tk.client.data_classes.migration import AssetCentricId, ResourceViewMapping
+from cognite_toolkit._cdf_tk.client import ToolkitClient
+from cognite_toolkit._cdf_tk.client.data_classes.migration import (
+    AssetCentricId,
+    ResourceViewMappingApply,
+)
 from cognite_toolkit._cdf_tk.utils.collection import flatten_dict_json_path
 from cognite_toolkit._cdf_tk.utils.dtype_conversion import (
     asset_centric_convert_to_primary_property,
@@ -22,124 +26,162 @@ from cognite_toolkit._cdf_tk.utils.dtype_conversion import (
 )
 from cognite_toolkit._cdf_tk.utils.useful_types import (
     AssetCentricResourceExtended,
-    AssetCentricType,
+    AssetCentricTypeExtended,
 )
 
 from .data_model import INSTANCE_SOURCE_VIEW_ID
 from .issues import ConversionIssue, FailedConversion, InvalidPropertyDataType
 
 
-@dataclass
 class DirectRelationCache:
     """Cache for direct relation references to look up target of direct relations.
 
-    This is used when creating direct relations from asset-centric resources to CogniteAsset and CogniteSourceSystem.
-
-    Attributes:
-        asset: Mapping[int, DirectRelationReference]
-            A mapping from asset IDs to their corresponding DirectRelationReference in the data model.
-        source: Mapping[str, DirectRelationReference]
-            A mapping from source strings to their corresponding DirectRelationReference in the data model.
-
-    Methods:
-        get(resource_type: AssetCentric, property_id: str) -> Mapping[str | int, DirectRelationReference]:
-            Get the appropriate mapping based on the resource type and property ID.
-
-    Note:
-        The ASSET_REFERENCE_PROPERTIES and SOURCE_REFERENCE_PROPERTIES class variables define which properties
-        in asset-centric resources reference assets and sources, respectively.
-
+    This is used when creating direct relations from asset-centric resources to assets, files, and source systems.
     """
 
-    ASSET_REFERENCE_PROPERTIES: ClassVar[Set[tuple[AssetCentricType, str]]] = {
+    class TableName:
+        ASSET_ID = "assetId"
+        SOURCE_NAME = "source"
+        FILE_ID = "fileId"
+        ASSET_EXTERNAL_ID = "assetExternalId"
+        FILE_EXTERNAL_ID = "fileExternalId"
+
+    ASSET_ID_PROPERTIES: ClassVar[Set[tuple[str, str]]] = {
         ("timeseries", "assetId"),
         ("file", "assetIds"),
         ("event", "assetIds"),
         ("sequence", "assetId"),
+        ("annotation", "data.assetRef.id"),
         ("asset", "parentId"),
-        ("fileAnnotation", "data.assetRef.id"),
     }
-    SOURCE_REFERENCE_PROPERTIES: ClassVar[Set[tuple[AssetCentricType, str]]] = {
+    SOURCE_NAME_PROPERTIES: ClassVar[Set[tuple[str, str]]] = {
         ("asset", "source"),
         ("event", "source"),
         ("file", "source"),
     }
-    FILE_REFERENCE_PROPERTIES: ClassVar[Set[tuple[AssetCentricType, str]]] = {
-        ("fileAnnotation", "data.fileRef.id"),
-        ("fileAnnotation", "annotatedResourceId"),
+    FILE_ID_PROPERTIES: ClassVar[Set[tuple[str, str]]] = {
+        ("annotation", "data.fileRef.id"),
+        ("annotation", "annotatedResourceId"),
     }
+    ASSET_EXTERNAL_ID_PROPERTIES: ClassVar[Set[tuple[str, str]]] = {("annotation", "data.assetRef.externalId")}
+    FILE_EXTERNAL_ID_PROPERTIES: ClassVar[Set[tuple[str, str]]] = {("annotation", "data.fileRef.externalId")}
+
+    def __init__(self, client: ToolkitClient) -> None:
+        self._client = client
+        self._cache_map: dict[
+            tuple[str, str] | str, dict[str, DirectRelationReference] | dict[int, DirectRelationReference]
+        ] = {}
+        # Constructing the cache map to be accessed by both table name and property id
+        for table_name, properties in [
+            (self.TableName.ASSET_ID, self.ASSET_ID_PROPERTIES),
+            (self.TableName.SOURCE_NAME, self.SOURCE_NAME_PROPERTIES),
+            (self.TableName.FILE_ID, self.FILE_ID_PROPERTIES),
+            (self.TableName.ASSET_EXTERNAL_ID, self.ASSET_EXTERNAL_ID_PROPERTIES),
+            (self.TableName.FILE_EXTERNAL_ID, self.FILE_EXTERNAL_ID_PROPERTIES),
+        ]:
+            cache: dict[str, DirectRelationReference] | dict[int, DirectRelationReference] = {}
+            self._cache_map[table_name] = cache
+            for key in properties:
+                self._cache_map[key] = cache
+
+    def update(self, resources: Iterable[AssetCentricResourceExtended]) -> None:
+        """Update the cache with direct relation references for the given asset-centric resources.
+
+        This is used to bulk update the cache for a chunk of resources before converting them to data model instances.
+        """
+        asset_ids: set[int] = set()
+        source_ids: set[str] = set()
+        file_ids: set[int] = set()
+        asset_external_ids: set[str] = set()
+        file_external_ids: set[str] = set()
+        for resource in resources:
+            if isinstance(resource, Annotation):
+                if resource.annotated_resource_type == "file" and resource.annotated_resource_id:
+                    file_ids.add(resource.annotated_resource_id)
+                if "assetRef" in resource.data:
+                    asset_ref = resource.data["assetRef"]
+                    if isinstance(asset_id := asset_ref.get("id"), int):
+                        asset_ids.add(asset_id)
+                    if isinstance(asset_external_id := asset_ref.get("externalId"), str):
+                        asset_external_ids.add(asset_external_id)
+                if "fileRef" in resource.data:
+                    file_ref = resource.data["fileRef"]
+                    if isinstance(file_id := file_ref.get("id"), int):
+                        file_ids.add(file_id)
+                    if isinstance(file_external_id := file_ref.get("externalId"), str):
+                        file_external_ids.add(file_external_id)
+            elif isinstance(resource, Asset):
+                if resource.source:
+                    source_ids.add(resource.source)
+                if resource.parent_id is not None:
+                    asset_ids.add(resource.parent_id)
+            elif isinstance(resource, FileMetadata):
+                if resource.source:
+                    source_ids.add(resource.source)
+                if resource.asset_ids:
+                    asset_ids.update(resource.asset_ids)
+            elif isinstance(resource, Event):
+                if resource.source:
+                    source_ids.add(resource.source)
+                if resource.asset_ids:
+                    asset_ids.update(resource.asset_ids)
+            elif isinstance(resource, TimeSeries):
+                if resource.asset_id is not None:
+                    asset_ids.add(resource.asset_id)
+        if asset_ids:
+            self._update_cache(self._client.migration.lookup.assets(id=list(asset_ids)), self.TableName.ASSET_ID)
+        if source_ids:
+            # SourceSystems are not cached in the client, so we have to handle the caching ourselves.
+            cache = cast(dict[str, DirectRelationReference], self._cache_map[self.TableName.SOURCE_NAME])
+            missing = set(source_ids) - set(cache.keys())
+            if missing:
+                source_systems = self._client.migration.created_source_system.retrieve(list(missing))
+                for source_system in source_systems:
+                    cache[source_system.source] = source_system.as_direct_relation_reference()
+        if file_ids:
+            self._update_cache(self._client.migration.lookup.files(id=list(file_ids)), self.TableName.FILE_ID)
+        if asset_external_ids:
+            self._update_cache(
+                self._client.migration.lookup.assets(external_id=list(asset_external_ids)),
+                self.TableName.ASSET_EXTERNAL_ID,
+            )
+        if file_external_ids:
+            self._update_cache(
+                self._client.migration.lookup.files(external_id=list(file_external_ids)),
+                self.TableName.FILE_EXTERNAL_ID,
+            )
 
-    asset: Mapping[int, DirectRelationReference]
-    source: Mapping[str, DirectRelationReference]
-    file: Mapping[int, DirectRelationReference]
-
-    def get(self, resource_type: AssetCentricType, property_id: str) -> Mapping[str | int, DirectRelationReference]:
-        key = resource_type, property_id
-        if key in self.ASSET_REFERENCE_PROPERTIES:
-            return self.asset  # type: ignore[return-value]
-        if key in self.SOURCE_REFERENCE_PROPERTIES:
-            return self.source  # type: ignore[return-value]
-        if key in self.FILE_REFERENCE_PROPERTIES:
-            return self.file  # type: ignore[return-value]
-        return {}
-
-
-@overload
-def asset_centric_to_dm(
-    resource: AssetCentricResourceExtended,
-    instance_id: NodeId,
-    view_source: ResourceViewMapping,
-    view_properties: dict[str, ViewProperty],
-    asset_instance_id_by_id: Mapping[int, DirectRelationReference],
-    source_instance_id_by_external_id: Mapping[str, DirectRelationReference],
-    file_instance_id_by_id: Mapping[int, DirectRelationReference],
-) -> tuple[NodeApply | None, ConversionIssue]: ...
-
+    def _update_cache(self, instance_id_by_id: dict[int, NodeId] | dict[str, NodeId], table_name: str) -> None:
+        cache = self._cache_map[table_name]
+        for identifier, instance_id in instance_id_by_id.items():
+            cache[identifier] = DirectRelationReference(space=instance_id.space, external_id=instance_id.external_id)  # type: ignore[index]
 
-@overload
-def asset_centric_to_dm(
-    resource: AssetCentricResourceExtended,
-    instance_id: EdgeId,
-    view_source: ResourceViewMapping,
-    view_properties: dict[str, ViewProperty],
-    asset_instance_id_by_id: Mapping[int, DirectRelationReference],
-    source_instance_id_by_external_id: Mapping[str, DirectRelationReference],
-    file_instance_id_by_id: Mapping[int, DirectRelationReference],
-) -> tuple[EdgeApply | None, ConversionIssue]: ...
+    def get_cache(
+        self, resource_type: AssetCentricTypeExtended, property_id: str
+    ) -> Mapping[str | int, DirectRelationReference] | None:
+        """Get the cache for the given resource type and property ID."""
+        return self._cache_map.get((resource_type, property_id))  # type: ignore[return-value]
 
 
 def asset_centric_to_dm(
     resource: AssetCentricResourceExtended,
-    instance_id: NodeId | EdgeId,
-    view_source: ResourceViewMapping,
+    instance_id: InstanceId,
+    view_source: ResourceViewMappingApply,
     view_properties: dict[str, ViewProperty],
-    asset_instance_id_by_id: Mapping[int, DirectRelationReference],
-    source_instance_id_by_external_id: Mapping[str, DirectRelationReference],
-    file_instance_id_by_id: Mapping[int, DirectRelationReference],
+    direct_relation_cache: DirectRelationCache,
 ) -> tuple[NodeApply | EdgeApply | None, ConversionIssue]:
     """Convert an asset-centric resource to a data model instance.
 
     Args:
         resource (CogniteResource): The asset-centric resource to convert.
         instance_id (NodeId | EdgeApply): The ID of the instance to create or update.
-        view_source (ResourceViewMapping): The view source defining how to map the resource to the data model.
+        view_source (ResourceViewMappingApply): The view source defining how to map the resource to the data model.
         view_properties (dict[str, ViewProperty]): The defined properties referenced in the view source mapping.
-        asset_instance_id_by_id (dict[int, DirectRelationReference]): A mapping from asset IDs to their corresponding
-            DirectRelationReference in the data model. This is used to create direct relations for resources that
-            reference assets.
-        source_instance_id_by_external_id (dict[str, DirectRelationReference]): A mapping from source strings to their
-            corresponding DirectRelationReference in the data model. This is used to create direct relations for resources
-            that reference sources.
-        file_instance_id_by_id (dict[int, DirectRelationReference]): A mapping from file IDs to their corresponding
-            DirectRelationReference in the data model. This is used to create direct relations for resources that
-            reference files.
+        direct_relation_cache (DirectRelationCache): Cache for direct relation references.
 
     Returns:
         tuple[NodeApply | EdgeApply, ConversionIssue]: A tuple containing the converted NodeApply and any ConversionIssue encountered.
     """
-    cache = DirectRelationCache(
-        asset=asset_instance_id_by_id, source=source_instance_id_by_external_id, file=file_instance_id_by_id
-    )
     resource_type = _lookup_resource_type(resource)
     dumped = resource.dump()
     try:
@@ -159,13 +201,13 @@ def asset_centric_to_dm(
         view_source.property_mapping,
         resource_type,
         issue=issue,
-        direct_relation_cache=cache,
+        direct_relation_cache=direct_relation_cache,
     )
     sources: list[NodeOrEdgeData] = []
     if properties:
         sources.append(NodeOrEdgeData(source=view_source.view_id, properties=properties))
 
-    if resource_type != "fileAnnotation":
+    if resource_type != "annotation":
         instance_source_properties = {
             "resourceType": resource_type,
             "id": id_,
@@ -177,7 +219,7 @@ def asset_centric_to_dm(
     instance: NodeApply | EdgeApply
     if isinstance(instance_id, EdgeId):
         edge_properties = create_edge_properties(
-            dumped, view_source.property_mapping, resource_type, issue, cache, instance_id.space
+            dumped, view_source.property_mapping, resource_type, issue, direct_relation_cache, instance_id.space
         )
         if any(key not in edge_properties for key in ("start_node", "end_node", "type")):
             # Failed conversion of edge properties
@@ -196,7 +238,7 @@ def asset_centric_to_dm(
     return instance, issue
 
 
-def _lookup_resource_type(resource_type: AssetCentricResourceExtended) -> AssetCentricType:
+def _lookup_resource_type(resource_type: AssetCentricResourceExtended) -> AssetCentricTypeExtended:
     if isinstance(resource_type, Asset):
         return "asset"
     elif isinstance(resource_type, FileMetadata):
@@ -210,7 +252,7 @@ def _lookup_resource_type(resource_type: AssetCentricResourceExtended) -> AssetC
             "diagrams.AssetLink",
             "diagrams.FileLink",
         ):
-            return "fileAnnotation"
+            return "annotation"
     raise ValueError(f"Unsupported resource type: {resource_type}")
 
 
@@ -218,7 +260,7 @@ def create_properties(
     dumped: dict[str, Any],
     view_properties: dict[str, ViewProperty],
     property_mapping: dict[str, str],
-    resource_type: AssetCentricType,
+    resource_type: AssetCentricTypeExtended,
     issue: ConversionIssue,
     direct_relation_cache: DirectRelationCache,
 ) -> dict[str, PropertyValueWrite]:
@@ -261,7 +303,7 @@ def create_properties(
                 dm_prop.nullable,
                 destination_container_property=(dm_prop.container, dm_prop.container_property_identifier),
                 source_property=(resource_type, prop_json_path),
-                direct_relation_lookup=direct_relation_cache.get(resource_type, prop_json_path),
+                direct_relation_lookup=direct_relation_cache.get_cache(resource_type, prop_json_path),
             )
         except (ValueError, TypeError, NotImplementedError) as e:
             issue.failed_conversions.append(
@@ -289,7 +331,7 @@ def create_properties(
 def create_edge_properties(
     dumped: dict[str, Any],
     property_mapping: dict[str, str],
-    resource_type: AssetCentricType,
+    resource_type: AssetCentricTypeExtended,
     issue: ConversionIssue,
     direct_relation_cache: DirectRelationCache,
     default_instance_space: str,
@@ -309,7 +351,7 @@ def create_edge_properties(
                     flatten_dump[prop_json_path],
                     DirectRelation(),
                     False,
-                    direct_relation_lookup=direct_relation_cache.get(resource_type, prop_json_path),
+                    direct_relation_lookup=direct_relation_cache.get_cache(resource_type, prop_json_path),
                 )
             except (ValueError, TypeError, NotImplementedError) as e:
                 issue.failed_conversions.append(

cognite_toolkit/_cdf_tk/commands/_migrate/data_classes.py

@@ -1,26 +1,30 @@
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Generic, Literal
+from typing import Annotated, Any, Generic, Literal
 
 from cognite.client.data_classes._base import (
     WriteableCogniteResource,
     WriteableCogniteResourceList,
 )
-from cognite.client.data_classes.data_modeling import InstanceApply, NodeId, ViewId
+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, field_validator, model_validator
+from pydantic import BaseModel, BeforeValidator, field_validator, model_validator
 
 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 create_default_mappings
+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 (
-    AssetCentricKind,
-    AssetCentricType,
+    AssetCentricKindExtended,
     JsonVal,
-    T_AssetCentricResource,
+    T_AssetCentricResourceExtended,
 )
 
 
@@ -37,8 +41,8 @@ class MigrationMapping(BaseModel, alias_generator=to_camel_case, extra="ignore",
            for example, the Canvas migration to determine which view to use for the resource.
     """
 
-    resource_type: AssetCentricType
-    instance_id: NodeId
+    resource_type: str
+    instance_id: InstanceId
     id: int
     data_set_id: int | None = None
     ingestion_view: str | None = None
@@ -56,7 +60,8 @@ class MigrationMapping(BaseModel, alias_generator=to_camel_case, extra="ignore",
         raise ToolkitValueError(f"No default ingestion view specified for resource type '{self.resource_type}'")
 
     def as_asset_centric_id(self) -> AssetCentricId:
-        return AssetCentricId(resource_type=self.resource_type, id_=self.id)
+        # 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:
@@ -87,12 +92,6 @@ class MigrationMapping(BaseModel, alias_generator=to_camel_case, extra="ignore",
             return ViewId.load(v)
         return v
 
-    @field_validator("instance_id", mode="before")
-    def _validate_instance_id(cls, v: Any) -> Any:
-        if isinstance(v, dict):
-            return NodeId.load(v)
-        return v
-
 
 class MigrationMappingList(ModelList[MigrationMapping]):
     @classmethod
@@ -113,14 +112,22 @@ class MigrationMappingList(ModelList[MigrationMapping]):
 
     def as_node_ids(self) -> list[NodeId]:
         """Return a list of NodeIds from the migration mappings."""
-        return [mapping.instance_id for mapping in self]
+        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]
+        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."""
@@ -131,7 +138,9 @@ class MigrationMappingList(ModelList[MigrationMapping]):
         return {mapping.id: mapping for mapping in self}
 
     @classmethod
-    def read_csv_file(cls, filepath: Path, resource_type: AssetCentricKind | None = None) -> "MigrationMappingList":
+    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]] = {
@@ -139,6 +148,7 @@ class MigrationMappingList(ModelList[MigrationMapping]):
             "TimeSeries": TimeSeriesMigrationMappingList,
             "FileMetadata": FileMigrationMappingList,
             "Events": EventMigrationMappingList,
+            "Annotations": AnnotationMigrationMappingList,
         }
         if resource_type not in cls_by_resource_type:
             raise ToolkitValueError(
@@ -147,20 +157,53 @@ class MigrationMappingList(ModelList[MigrationMapping]):
         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):
@@ -187,10 +230,16 @@ class TimeSeriesMigrationMappingList(MigrationMappingList):
         return TimeSeriesMapping
 
 
+class AnnotationMigrationMappingList(MigrationMappingList):
+    @classmethod
+    def _get_base_model_cls(cls) -> type[AnnotationMapping]:
+        return AnnotationMapping
+
+
 @dataclass
-class AssetCentricMapping(Generic[T_AssetCentricResource], WriteableCogniteResource[InstanceApply]):
+class AssetCentricMapping(Generic[T_AssetCentricResourceExtended], WriteableCogniteResource[InstanceApply]):
     mapping: MigrationMapping
-    resource: T_AssetCentricResource
+    resource: T_AssetCentricResourceExtended
 
     def as_write(self) -> InstanceApply:
         raise NotImplementedError()
@@ -205,7 +254,9 @@ class AssetCentricMapping(Generic[T_AssetCentricResource], WriteableCogniteResou
         }
 
 
-class AssetCentricMappingList(WriteableCogniteResourceList[InstanceApply, AssetCentricMapping[T_AssetCentricResource]]):
+class AssetCentricMappingList(
+    WriteableCogniteResourceList[InstanceApply, AssetCentricMapping[T_AssetCentricResourceExtended]]
+):
     _RESOURCE: type = AssetCentricMapping
 
     def as_write(self) -> InstanceApplyList: