cognite-toolkit 0.6.101__tar.gz → 0.6.102__tar.gz

{cognite_toolkit-0.6.101 → cognite_toolkit-0.6.102}/.gemini/styleguide.md

@@ -64,6 +64,8 @@ def load_config(path: Path) -> dict[str, Any]:
 - **Absolute imports**: Always use absolute imports for clarity
 - **Sort alphabetically** within groups
 - **Type checking imports**: Use `TYPE_CHECKING` for type-only imports
+- **`_cdf_tk` imports**: Always use the full path `cognite_toolkit._cdf_tk`
+when importing from `_cdf_tk`, never use `from _cdf_tk import ...`
 
 ```python
 import json
@@ -76,6 +78,13 @@ from pydantic import BaseModel
 
 from cog_ai.common.types import ModelConfig
 
+# Good - full path for _cdf_tk imports
+from cognite_toolkit._cdf_tk.feature_flags import Flags
+from cognite_toolkit._cdf_tk.commands.repo import RepoCommand
+
+# Bad - don't use relative _cdf_tk imports
+# from _cdf_tk.feature_flags import Flags  # ❌
+
 if TYPE_CHECKING:
     from cog_ai.tools.query.types import QueryCompletion
 ```

{cognite_toolkit-0.6.101 → cognite_toolkit-0.6.102}/.pre-commit-config.yaml

@@ -2,7 +2,7 @@
 repos:
   - repo: https://github.com/astral-sh/uv-pre-commit
     # uv version.
-    rev: 0.9.8
+    rev: 0.9.9
     hooks:
       - id: uv-lock
 
@@ -18,7 +18,7 @@ repos:
           - --fixable=E,W,F,I,T,RUF,TID,UP
           - --target-version=py310
       - id: ruff-format
-    rev: v0.14.4
+    rev: v0.14.5
 
   - repo: https://github.com/igorshubovych/markdownlint-cli
     rev: v0.45.0

{cognite_toolkit-0.6.101 → cognite_toolkit-0.6.102}/.vscode/launch.json

@@ -23,6 +23,18 @@
       "console": "integratedTerminal",
       "justMyCode": false
     },
+    {
+      "name": "cdf init",
+      "type": "debugpy",
+      "request": "launch",
+      "program": "./cdf-tk-dev.py",
+      "args": [
+        "init",
+        "--dry-run"
+      ],
+      "console": "integratedTerminal",
+      "justMyCode": false
+    },
     {
       "name": "cdf collect opt-in",
       "type": "debugpy",
@@ -73,7 +85,6 @@
         "--override-env",
         "modules",
         "init",
-        "cognite_toolkit"
       ],
       "console": "integratedTerminal",
       "justMyCode": false

{cognite_toolkit-0.6.101 → cognite_toolkit-0.6.102}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cognite_toolkit
-Version: 0.6.101
+Version: 0.6.102
 Summary: Official Cognite Data Fusion tool for project templates and configuration deployment
 Project-URL: Homepage, https://docs.cognite.com/cdf/deploy/cdf_toolkit/
 Project-URL: Changelog, https://github.com/cognitedata/toolkit/releases

{cognite_toolkit-0.6.101 → cognite_toolkit-0.6.102}/cdf.toml

@@ -32,7 +32,7 @@ dev = true
 [modules]
 # This is the version of the modules. It should not be changed manually.
 # It will be updated by the 'cdf modules upgrade' command.
-version = "0.6.101"
+version = "0.6.102"
 
 
 [library.toolkit-data]

cognite_toolkit-0.6.102/cognite_toolkit/_cdf_tk/apps/_landing_app.py

@@ -0,0 +1,42 @@
+from typing import Annotated
+
+import typer
+
+from cognite_toolkit._cdf_tk.cdf_toml import CDFToml
+from cognite_toolkit._cdf_tk.commands import InitCommand
+from cognite_toolkit._cdf_tk.feature_flags import Flags
+
+
+class LandingApp(typer.Typer):
+    def __init__(self, *args, **kwargs) -> None:  # type: ignore
+        super().__init__(*args, **kwargs)
+
+    def main_init(
+        self,
+        dry_run: Annotated[
+            bool,
+            typer.Option(
+                "--dry-run",
+                "-r",
+                help="Whether to do a dry-run, do dry-run if present.",
+            ),
+        ] = False,
+        # TODO: this is a temporary solution to be able to test the functionality
+        # in a new environment, assuming that the toml file doesn't exist yet.
+        # remove this once v.07 is released
+        v7: Annotated[
+            bool,
+            typer.Option(
+                "--seven",
+                "-s",
+                help="Emulate v0.7",
+                hidden=(Flags.v07.is_enabled() or not CDFToml.load().is_loaded_from_file),
+            ),
+        ] = False,
+    ) -> None:
+        """Getting started checklist"""
+        cmd = InitCommand()
+        # Tracking  command with the usual lambda run construct
+        # is intentionally left out because we don't want to expose the user to the warning
+        # before they've had the chance to opt in (which is something they'll do later using this command).
+        cmd.execute(dry_run=dry_run, emulate_dot_seven=v7)

{cognite_toolkit-0.6.101 → cognite_toolkit-0.6.102}/cognite_toolkit/_cdf_tk/cdf_toml.py

@@ -9,7 +9,7 @@ from typing import Any, ClassVar
 from rich import print
 
 from cognite_toolkit import _version
-from cognite_toolkit._cdf_tk.constants import clean_name
+from cognite_toolkit._cdf_tk.constants import RESOURCES_PATH, EnvType, clean_name
 from cognite_toolkit._cdf_tk.exceptions import (
     ToolkitRequiredValueError,
     ToolkitTOMLFormatError,
@@ -176,6 +176,25 @@ class CDFToml:
                 is_loaded_from_file=False,
             )
 
+    @classmethod
+    def write(cls, organization_dir: Path, env: EnvType = "dev", version: str = _version.__version__) -> None:
+        destination = Path.cwd() / CDFToml.file_name
+        if destination.exists():
+            print("cdf.toml file already exists. Skipping creation.")
+            return
+        cdf_toml_content = (RESOURCES_PATH / CDFToml.file_name).read_text(encoding="utf-8")
+        cdf_toml_content = cdf_toml_content.replace("0.0.0", version)
+        if organization_dir != Path.cwd():
+            cdf_toml_content = cdf_toml_content.replace(
+                "#<PLACEHOLDER>",
+                f'''
+default_organization_dir = "{organization_dir.name}"''',
+            )
+        else:
+            cdf_toml_content = cdf_toml_content.replace("#<PLACEHOLDER>", "")
+        cdf_toml_content = cdf_toml_content.replace("<DEFAULT_ENV_PLACEHOLDER>", env)
+        destination.write_text(cdf_toml_content, encoding="utf-8")
+
 
 def _read_toml(file_path: Path) -> dict[str, Any]:
     # TOML files are required to be UTF-8 encoded

{cognite_toolkit-0.6.101 → cognite_toolkit-0.6.102}/cognite_toolkit/_cdf_tk/client/api/migration.py

@@ -1,7 +1,7 @@
 import warnings
 from collections.abc import Sequence
 from itertools import groupby
-from typing import TypeVar, overload
+from typing import Literal, TypeVar, cast, overload
 
 from cognite.client._constants import DEFAULT_LIMIT_READ
 from cognite.client.data_classes.data_modeling import (
@@ -347,9 +347,117 @@ class SpaceSourceAPI:
         return results
 
 
+class LookupAPI:
+    def __init__(self, instance_api: ExtendedInstancesAPI, resource_type: AssetCentricType) -> None:
+        self._instance_api = instance_api
+        self._resource_type = resource_type
+        self._view_id = InstanceSource.get_source()
+        self._node_id_by_id: dict[int, NodeId | None] = {}
+        self._node_id_by_external_id: dict[str, NodeId | None] = {}
+        self._RETRIEVE_LIMIT = 1000
+
+    @overload
+    def __call__(self, id: int, external_id: None = None) -> NodeId | None: ...
+
+    @overload
+    def __call__(self, id: Sequence[int], external_id: None = None) -> dict[int, NodeId]: ...
+
+    @overload
+    def __call__(self, *, external_id: str) -> NodeId | None: ...
+
+    @overload
+    def __call__(self, *, external_id: SequenceNotStr[str]) -> dict[str, NodeId]: ...
+
+    def __call__(
+        self, id: int | Sequence[int] | None = None, external_id: str | SequenceNotStr[str] | None = None
+    ) -> dict[int, NodeId] | dict[str, NodeId] | NodeId | None:
+        """Lookup NodeId by either internal ID or external ID.
+
+        Args:
+            id (int | Sequence[int] | None): The internal ID(s) to lookup.
+            external_id (str | SequenceNotStr[str] | None): The external ID(s) to lookup.
+
+        Returns:
+            NodeId | dict[int, NodeId] | dict[str, NodeId] | None: The corresponding NodeId(s) if found, otherwise None.
+
+        """
+        if id is not None and external_id is None:
+            return self._lookup_by_id(id)
+        elif external_id is not None and id is None:
+            return self._lookup_by_external_id(external_id)
+        else:
+            raise ValueError("Either id or external_id must be provided, but not both.")
+
+    def _lookup_by_id(self, id: int | Sequence[int]) -> dict[int, NodeId] | NodeId | None:
+        ids: list[int] = [id] if isinstance(id, int) else list(id)
+
+        missing = [id_ for id_ in ids if id_ not in self._node_id_by_id]
+        if missing:
+            self._fetch_and_cache(missing, by="id")
+        if isinstance(id, int):
+            return self._node_id_by_id.get(id)
+        return {id_: node_id for id_ in ids if isinstance(node_id := self._node_id_by_id.get(id_), NodeId)}
+
+    def _lookup_by_external_id(self, external_id: str | SequenceNotStr[str]) -> dict[str, NodeId] | NodeId | None:
+        external_ids: list[str] = [external_id] if isinstance(external_id, str) else list(external_id)
+
+        missing = [ext_id for ext_id in external_ids if ext_id not in self._node_id_by_external_id]
+        if missing:
+            self._fetch_and_cache(missing, by="classicExternalId")
+        if isinstance(external_id, str):
+            return self._node_id_by_external_id.get(external_id)
+        return {
+            ext_id: node_id
+            for ext_id in external_ids
+            if isinstance(node_id := self._node_id_by_external_id.get(ext_id), NodeId)
+        }
+
+    def _fetch_and_cache(self, identifiers: Sequence[int | str], by: Literal["id", "classicExternalId"]) -> None:
+        for chunk in chunker_sequence(identifiers, self._RETRIEVE_LIMIT):
+            retrieve_query = query.Query(
+                with_={
+                    "instanceSource": query.NodeResultSetExpression(
+                        filter=filters.And(
+                            filters.HasData(views=[self._view_id]),
+                            filters.Equals(self._view_id.as_property_ref("resourceType"), self._resource_type),
+                            filters.In(self._view_id.as_property_ref(by), list(chunk)),
+                        ),
+                        limit=len(chunk),
+                    ),
+                },
+                select={"instanceSource": query.Select([query.SourceSelector(self._view_id, ["*"])])},
+            )
+            chunk_response = self._instance_api.query(retrieve_query)
+            items = chunk_response.get("instanceSource", [])
+            for item in items:
+                instance_source = InstanceSource._load(item.dump())
+                node_id = instance_source.as_id()
+                self._node_id_by_id[instance_source.id_] = node_id
+                if instance_source.classic_external_id:
+                    self._node_id_by_external_id[instance_source.classic_external_id] = node_id
+            missing = set(chunk) - set(self._node_id_by_id.keys()) - set(self._node_id_by_external_id.keys())
+            if by == "id":
+                for missing_id in cast(set[int], missing):
+                    if missing_id not in self._node_id_by_id:
+                        self._node_id_by_id[missing_id] = None
+            elif by == "classicExternalId":
+                for missing_ext_id in cast(set[str], missing):
+                    if missing_ext_id not in self._node_id_by_external_id:
+                        self._node_id_by_external_id[missing_ext_id] = None
+
+
+class MigrationLookupAPI:
+    def __init__(self, instance_api: ExtendedInstancesAPI) -> None:
+        self.assets = LookupAPI(instance_api, "asset")
+        self.events = LookupAPI(instance_api, "event")
+        self.files = LookupAPI(instance_api, "file")
+        self.time_series = LookupAPI(instance_api, "timeseries")
+
+
 class MigrationAPI:
     def __init__(self, instance_api: ExtendedInstancesAPI) -> None:
         self.instance_source = InstanceSourceAPI(instance_api)
         self.resource_view_mapping = ResourceViewMappingAPI(instance_api)
         self.created_source_system = CreatedSourceSystemAPI(instance_api)
         self.space_source = SpaceSourceAPI(instance_api)
+        self.lookup = MigrationLookupAPI(instance_api)

{cognite_toolkit-0.6.101 → cognite_toolkit-0.6.102}/cognite_toolkit/_cdf_tk/client/data_classes/migration.py

@@ -16,7 +16,7 @@ from cognite.client.data_classes.data_modeling.instances import (
 
 from cognite_toolkit._cdf_tk.constants import COGNITE_MIGRATION_SPACE
 from cognite_toolkit._cdf_tk.tk_warnings import IgnoredValueWarning
-from cognite_toolkit._cdf_tk.utils.useful_types import AssetCentricType
+from cognite_toolkit._cdf_tk.utils.useful_types import AssetCentricType, AssetCentricTypeExtended
 
 if sys.version_info >= (3, 11):
     from typing import Self
@@ -26,7 +26,7 @@ else:
 
 @dataclass(frozen=True)
 class AssetCentricId(CogniteObject):
-    resource_type: AssetCentricType
+    resource_type: AssetCentricTypeExtended
     id_: int
 
     @classmethod

{cognite_toolkit-0.6.101 → cognite_toolkit-0.6.102}/cognite_toolkit/_cdf_tk/commands/_changes.py

@@ -12,6 +12,7 @@ from packaging.version import parse as parse_version
 from rich import print
 
 from cognite_toolkit._cdf_tk.builders import get_loader
+from cognite_toolkit._cdf_tk.cdf_toml import CDFToml
 from cognite_toolkit._cdf_tk.constants import DOCKER_IMAGE_NAME
 from cognite_toolkit._cdf_tk.data_classes import ModuleDirectories
 from cognite_toolkit._cdf_tk.utils import iterate_modules, read_yaml_file, safe_read, safe_write
@@ -362,7 +363,6 @@ After:
 
     def do(self) -> set[Path]:
         # Avoid circular import
-        from .modules import ModulesCommand
 
         system_yaml = self._organization_dir / "_system.yaml"
         if not system_yaml.exists():
@@ -372,11 +372,8 @@ After:
         content = read_yaml_file(system_yaml)
         current_version = content.get("cdf_toolkit_version", __version__)
 
-        cdf_toml_content = ModulesCommand(skip_tracking=True).create_cdf_toml(self._organization_dir)
-        cdf_toml_content = cdf_toml_content.replace(f'version = "{__version__}"', f'version = "{current_version}"')
-
-        cdf_toml_path = Path.cwd() / "cdf.toml"
-        cdf_toml_path.write_text(cdf_toml_content, encoding="utf-8")
+        CDFToml.write(self._organization_dir, version=current_version)
+        cdf_toml_path = Path.cwd() / CDFToml.file_name
         system_yaml.unlink()
         return {cdf_toml_path, system_yaml}
 

{cognite_toolkit-0.6.101 → cognite_toolkit-0.6.102}/cognite_toolkit/_cdf_tk/commands/_migrate/conversion.py

@@ -1,6 +1,6 @@
 from collections.abc import Mapping, Set
 from dataclasses import dataclass
-from typing import Any, ClassVar, overload
+from typing import Any, ClassVar
 
 from cognite.client.data_classes import Annotation, Asset, Event, FileMetadata, TimeSeries
 from cognite.client.data_classes.data_modeling import (
@@ -13,6 +13,7 @@ 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.utils.collection import flatten_dict_json_path
@@ -22,7 +23,7 @@ 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
@@ -51,29 +52,31 @@ class DirectRelationCache:
 
     """
 
-    ASSET_REFERENCE_PROPERTIES: ClassVar[Set[tuple[AssetCentricType, str]]] = {
+    ASSET_REFERENCE_PROPERTIES: ClassVar[Set[tuple[AssetCentricTypeExtended, str]]] = {
         ("timeseries", "assetId"),
         ("file", "assetIds"),
         ("event", "assetIds"),
         ("sequence", "assetId"),
         ("asset", "parentId"),
-        ("fileAnnotation", "data.assetRef.id"),
+        ("annotation", "data.assetRef.id"),
     }
-    SOURCE_REFERENCE_PROPERTIES: ClassVar[Set[tuple[AssetCentricType, str]]] = {
+    SOURCE_REFERENCE_PROPERTIES: ClassVar[Set[tuple[AssetCentricTypeExtended, str]]] = {
         ("asset", "source"),
         ("event", "source"),
         ("file", "source"),
     }
-    FILE_REFERENCE_PROPERTIES: ClassVar[Set[tuple[AssetCentricType, str]]] = {
-        ("fileAnnotation", "data.fileRef.id"),
-        ("fileAnnotation", "annotatedResourceId"),
+    FILE_REFERENCE_PROPERTIES: ClassVar[Set[tuple[AssetCentricTypeExtended, str]]] = {
+        ("annotation", "data.fileRef.id"),
+        ("annotation", "annotatedResourceId"),
     }
 
     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]:
+    def get(
+        self, resource_type: AssetCentricTypeExtended, 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]
@@ -84,33 +87,9 @@ class DirectRelationCache:
         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]: ...
-
-
-@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 asset_centric_to_dm(
-    resource: AssetCentricResourceExtended,
-    instance_id: NodeId | EdgeId,
+    instance_id: InstanceId,
     view_source: ResourceViewMapping,
     view_properties: dict[str, ViewProperty],
     asset_instance_id_by_id: Mapping[int, DirectRelationReference],
@@ -165,7 +144,7 @@ def asset_centric_to_dm(
     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_,
@@ -196,7 +175,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 +189,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 +197,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]:
@@ -289,7 +268,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,

{cognite_toolkit-0.6.101 → cognite_toolkit-0.6.102}/cognite_toolkit/_cdf_tk/commands/_migrate/data_classes.py

@@ -1,14 +1,15 @@
 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
@@ -17,10 +18,9 @@ from cognite_toolkit._cdf_tk.commands._migrate.default_mappings import create_de
 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 +37,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 +56,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 +88,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 +108,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 +134,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 +144,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 +153,41 @@ 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
+
+    @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 +214,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 +238,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: