cognite-toolkit 0.6.79__py3-none-any.whl → 0.6.81__py3-none-any.whl

cognite_toolkit/_builtin_modules/cdf.toml

@@ -4,7 +4,7 @@ default_env = "<DEFAULT_ENV_PLACEHOLDER>"
 [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.79"
+version = "0.6.81"
 
 
 [plugins]

cognite_toolkit/_cdf.py

@@ -29,7 +29,6 @@ from cognite_toolkit._cdf_tk.apps import (
     LandingApp,
     MigrateApp,
     ModulesApp,
-    PopulateApp,
     ProfileApp,
     PurgeApp,
     RepoApp,
@@ -100,9 +99,6 @@ if Plugins.dump.value.is_enabled():
 if Plugins.purge.value.is_enabled() and not Flags.v07.is_enabled():
     _app.add_typer(PurgeApp(**default_typer_kws), name="purge")
 
-if Flags.POPULATE.is_enabled():
-    _app.add_typer(PopulateApp(**default_typer_kws), name="populate")
-
 if Flags.PROFILE.is_enabled():
     _app.add_typer(ProfileApp(**default_typer_kws), name="profile")
 

cognite_toolkit/_cdf_tk/apps/__init__.py

@@ -6,7 +6,6 @@ from ._dump_app import DumpApp
 from ._landing_app import LandingApp
 from ._migrate_app import MigrateApp
 from ._modules_app import ModulesApp
-from ._populate_app import PopulateApp
 from ._profile_app import ProfileApp
 from ._purge import PurgeApp
 from ._repo_app import RepoApp
@@ -22,7 +21,6 @@ __all__ = [
     "LandingApp",
     "MigrateApp",
     "ModulesApp",
-    "PopulateApp",
     "ProfileApp",
     "PurgeApp",
     "RepoApp",

cognite_toolkit/_cdf_tk/apps/_migrate_app.py

@@ -12,17 +12,20 @@ from cognite_toolkit._cdf_tk.commands import (
     MigrationPrepareCommand,
 )
 from cognite_toolkit._cdf_tk.commands._migrate import MigrationCommand
-from cognite_toolkit._cdf_tk.commands._migrate.adapter import (
-    AssetCentricMigrationIOAdapter,
-    FileMetaIOAdapter,
+from cognite_toolkit._cdf_tk.commands._migrate.creators import (
+    InfieldV2ConfigCreator,
+    InstanceSpaceCreator,
+    SourceSystemCreator,
+)
+from cognite_toolkit._cdf_tk.commands._migrate.data_mapper import AssetCentricMapper
+from cognite_toolkit._cdf_tk.commands._migrate.migration_io import (
+    AssetCentricMigrationIO,
+)
+from cognite_toolkit._cdf_tk.commands._migrate.selectors import (
+    AssetCentricMigrationSelector,
     MigrateDataSetSelector,
     MigrationCSVFileSelector,
-    MigrationSelector,
-    TimeSeriesIOAdapter,
 )
-from cognite_toolkit._cdf_tk.commands._migrate.creators import InstanceSpaceCreator, SourceSystemCreator
-from cognite_toolkit._cdf_tk.commands._migrate.data_mapper import AssetCentricMapper
-from cognite_toolkit._cdf_tk.storageio import AssetIO, EventIO
 from cognite_toolkit._cdf_tk.utils.auth import EnvironmentVariables
 from cognite_toolkit._cdf_tk.utils.cli_args import parse_view_str
 from cognite_toolkit._cdf_tk.utils.interactive_select import (
@@ -47,6 +50,8 @@ class MigrateApp(typer.Typer):
         self.command("timeseries")(self.timeseries)
         self.command("files")(self.files)
         self.command("canvas")(self.canvas)
+        # Uncomment when infield v2 config migration is ready
+        # self.command("infield-configs")(self.infield_configs)
 
     def main(self, ctx: typer.Context) -> None:
         """Migrate resources from Asset-Centric to data modeling in CDF."""
@@ -317,7 +322,7 @@ class MigrateApp(typer.Typer):
         cmd.run(
             lambda: cmd.migrate(
                 selected=selected,
-                data=AssetCentricMigrationIOAdapter(client, AssetIO(client)),
+                data=AssetCentricMigrationIO(client),
                 mapper=AssetCentricMapper(client),
                 log_dir=log_dir,
                 dry_run=dry_run,
@@ -337,11 +342,11 @@ class MigrateApp(typer.Typer):
         kind: AssetCentricKind,
         resource_type: str,
         container_id: ContainerId,
-    ) -> tuple[MigrationSelector, bool, bool]:
+    ) -> tuple[AssetCentricMigrationSelector, bool, bool]:
         if data_set_id is not None and mapping_file is not None:
             raise typer.BadParameter("Cannot specify both data_set_id and mapping_file")
         elif mapping_file is not None:
-            selected: MigrationSelector = MigrationCSVFileSelector(datafile=mapping_file, kind=kind)
+            selected: AssetCentricMigrationSelector = MigrationCSVFileSelector(datafile=mapping_file, kind=kind)
         elif data_set_id is not None:
             parsed_view = parse_view_str(consumption_view) if consumption_view is not None else None
             selected = MigrateDataSetSelector(
@@ -466,7 +471,7 @@ class MigrateApp(typer.Typer):
         cmd.run(
             lambda: cmd.migrate(
                 selected=selected,
-                data=AssetCentricMigrationIOAdapter(client, EventIO(client)),
+                data=AssetCentricMigrationIO(client),
                 mapper=AssetCentricMapper(client),
                 log_dir=log_dir,
                 dry_run=dry_run,
@@ -573,7 +578,7 @@ class MigrateApp(typer.Typer):
         cmd.run(
             lambda: cmd.migrate(
                 selected=selected,
-                data=TimeSeriesIOAdapter(client, skip_linking=skip_linking),
+                data=AssetCentricMigrationIO(client, skip_linking=skip_linking),
                 mapper=AssetCentricMapper(client),
                 log_dir=log_dir,
                 dry_run=dry_run,
@@ -681,7 +686,7 @@ class MigrateApp(typer.Typer):
         cmd.run(
             lambda: cmd.migrate(
                 selected=selected,
-                data=FileMetaIOAdapter(client, skip_linking=skip_linking),
+                data=AssetCentricMigrationIO(client, skip_linking=skip_linking),
                 mapper=AssetCentricMapper(client),
                 log_dir=log_dir,
                 dry_run=dry_run,
@@ -734,3 +739,46 @@ class MigrateApp(typer.Typer):
                 verbose=verbose,
             )
         )
+
+    @staticmethod
+    def infield_configs(
+        ctx: typer.Context,
+        output_dir: Annotated[
+            Path,
+            typer.Option(
+                "--output-dir",
+                "-o",
+                help="Path to the directory where the Infield V2 configuration definitions will be dumped. It is recommended "
+                "to govern these configurations in a git repository.",
+            ),
+        ] = Path("tmp"),
+        dry_run: Annotated[
+            bool,
+            typer.Option(
+                "--dry-run",
+                "-d",
+                help="If set, the migration will not be executed, but only a report of what would be done is printed.",
+            ),
+        ] = False,
+        verbose: Annotated[
+            bool,
+            typer.Option(
+                "--verbose",
+                "-v",
+                help="Turn on to get more verbose output when running the command",
+            ),
+        ] = False,
+    ) -> None:
+        """Creates Infield V2 configurations from existing APM Configurations in CDF."""
+        client = EnvironmentVariables.create_from_environment().get_client()
+
+        cmd = MigrationCommand()
+        cmd.run(
+            lambda: cmd.create(
+                client,
+                creator=InfieldV2ConfigCreator(client),
+                output_dir=output_dir,
+                dry_run=dry_run,
+                verbose=verbose,
+            )
+        )

cognite_toolkit/_cdf_tk/apps/_modules_app.py

@@ -6,6 +6,7 @@ from rich import print
 
 from cognite_toolkit._cdf_tk.cdf_toml import CDFToml
 from cognite_toolkit._cdf_tk.commands import ModulesCommand, PullCommand
+from cognite_toolkit._cdf_tk.feature_flags import Flags
 from cognite_toolkit._cdf_tk.utils.auth import EnvironmentVariables
 from cognite_toolkit._version import __version__
 
@@ -51,6 +52,24 @@ class ModulesApp(typer.Typer):
                 help="Clean target directory if it exists",
             ),
         ] = False,
+        library_url: Annotated[
+            str | None,
+            typer.Option(
+                "--library-url",
+                "-u",
+                help="URL of the library to add to the project.",
+                hidden=not Flags.EXTERNAL_LIBRARIES.is_enabled(),
+            ),
+        ] = None,
+        library_checksum: Annotated[
+            str | None,
+            typer.Option(
+                "--library-checksum",
+                "-c",
+                help="Checksum of the library to add to the project.",
+                hidden=not Flags.EXTERNAL_LIBRARIES.is_enabled(),
+            ),
+        ] = None,
         verbose: Annotated[
             bool,
             typer.Option(
@@ -62,12 +81,20 @@ class ModulesApp(typer.Typer):
     ) -> None:
         """Initialize or upgrade a new CDF project with templates interactively."""
 
+        if library_url and not library_checksum:
+            raise typer.BadParameter(
+                "--library-checksum must be provided when --library-url is specified.",
+                param_hint="--library-checksum",
+            )
+
         with ModulesCommand() as cmd:
             cmd.run(
                 lambda: cmd.init(
                     organization_dir=organization_dir,
                     select_all=all,
                     clean=clean,
+                    library_url=library_url,
+                    library_checksum=library_checksum,
                 )
             )
 

cognite_toolkit/_cdf_tk/apps/_purge.py

@@ -16,6 +16,7 @@ from cognite_toolkit._cdf_tk.storageio.selectors import (
     SelectedView,
 )
 from cognite_toolkit._cdf_tk.utils.auth import EnvironmentVariables
+from cognite_toolkit._cdf_tk.utils.cli_args import parse_view_str
 from cognite_toolkit._cdf_tk.utils.interactive_select import DataModelingSelect
 
 
@@ -191,10 +192,10 @@ class PurgeApp(typer.Typer):
     @staticmethod
     def purge_instances(
         view: Annotated[
-            list[str] | None,
+            str | None,
             typer.Argument(
                 help="Purge instances with properties in the specified view. Expected format is "
-                "'space externalId version'. For example 'cdf_cdm CogniteTimeSeries v1' will purge all nodes"
+                "'space:externalId/version'. For example 'cdf_cdm:CogniteTimeSeries/v1' will purge all nodes"
                 "that have properties in the CogniteTimeSeries view. If not provided and no "
                 "instance list is provided, interactive mode will be used.",
             ),
@@ -291,7 +292,7 @@ class PurgeApp(typer.Typer):
         elif instance_list is not None:
             selector = InstanceFileSelector(datafile=instance_list)
         elif view is not None:
-            view_id = cmd.get_selected_view_id(view)  # Will raise if not exactly one view
+            view_id = parse_view_str(view)
             selector = InstanceViewSelector(
                 view=SelectedView(
                     space=view_id.space, external_id=view_id.external_id, version=cast(str, view_id.version)

cognite_toolkit/_cdf_tk/commands/__init__.py

@@ -1,11 +1,8 @@
 from ._download import DownloadCommand
 from ._migrate import (
-    MigrateFilesCommand,
-    MigrateTimeseriesCommand,
     MigrationCanvasCommand,
     MigrationPrepareCommand,
 )
-from ._populate import PopulateCommand
 from ._profile import ProfileAssetCentricCommand, ProfileAssetCommand, ProfileRawCommand, ProfileTransformationCommand
 from ._purge import PurgeCommand
 from ._upload import UploadCommand
@@ -34,12 +31,9 @@ __all__ = [
     "DumpResourceCommand",
     "FeatureFlagCommand",
     "InitCommand",
-    "MigrateFilesCommand",
-    "MigrateTimeseriesCommand",
     "MigrationCanvasCommand",
     "MigrationPrepareCommand",
     "ModulesCommand",
-    "PopulateCommand",
     "ProfileAssetCentricCommand",
     "ProfileAssetCommand",
     "ProfileRawCommand",

cognite_toolkit/_cdf_tk/commands/_migrate/__init__.py

@@ -1,12 +1,8 @@
 from .canvas import MigrationCanvasCommand
 from .command import MigrationCommand
-from .files import MigrateFilesCommand
 from .prepare import MigrationPrepareCommand
-from .timeseries import MigrateTimeseriesCommand
 
 __all__ = [
-    "MigrateFilesCommand",
-    "MigrateTimeseriesCommand",
     "MigrationCanvasCommand",
     "MigrationCommand",
     "MigrationPrepareCommand",

cognite_toolkit/_cdf_tk/commands/_migrate/command.py

@@ -1,9 +1,8 @@
 from collections.abc import Callable, Iterable, Sequence
 from enum import Enum
 from pathlib import Path
-from typing import TypeVar
 
-from cognite.client.data_classes._base import CogniteResource
+from cognite.client.data_classes._base import T_CogniteResource
 from rich import print
 from rich.console import Console
 from rich.table import Table
@@ -28,12 +27,10 @@ from cognite_toolkit._cdf_tk.utils.fileio import Chunk, CSVWriter, NDJsonWriter,
 from cognite_toolkit._cdf_tk.utils.http_client import HTTPClient, HTTPMessage, ItemMessage, SuccessResponseItems
 from cognite_toolkit._cdf_tk.utils.producer_worker import ProducerWorkerExecutor
 from cognite_toolkit._cdf_tk.utils.progress_tracker import AVAILABLE_STATUS, ProgressTracker, Status
+from cognite_toolkit._cdf_tk.utils.useful_types import T_WriteCogniteResource
 
 from .data_model import INSTANCE_SOURCE_VIEW_ID, MODEL_ID, RESOURCE_VIEW_MAPPING_VIEW_ID
 
-T_CogniteResource = TypeVar("T_CogniteResource", bound=CogniteResource)
-T_WriteCogniteResource = TypeVar("T_WriteCogniteResource", bound=CogniteResource)
-
 
 class MigrationCommand(ToolkitCommand):
     class Steps(str, Enum):
@@ -79,7 +76,7 @@ class MigrationCommand(ToolkitCommand):
             ](
                 download_iterable=self._download_iterable(selected, data, tracker),
                 process=self._convert(mapper, data, tracker, log_file),
-                write=self._upload(write_client, data, tracker, log_file, dry_run),
+                write=self._upload(selected, write_client, data, tracker, log_file, dry_run),
                 iteration_count=iteration_count,
                 max_queue_size=10,
                 download_description=f"Downloading {selected.display_name}",
@@ -179,6 +176,7 @@ class MigrationCommand(ToolkitCommand):
 
     def _upload(
         self,
+        selected: T_Selector,
         write_client: HTTPClient,
         target: UploadableStorageIO[T_Selector, T_CogniteResource, T_WriteCogniteResource],
         tracker: ProgressTracker[str],
@@ -192,7 +190,7 @@ class MigrationCommand(ToolkitCommand):
             if dry_run:
                 responses = [SuccessResponseItems(200, "", [item.source_id for item in data_item])]
             else:
-                responses = target.upload_items(data_chunk=data_item, http_client=write_client, selector=None)
+                responses = target.upload_items(data_chunk=data_item, http_client=write_client, selector=selected)
 
             issues: list[Chunk] = []
             for item in responses:
@@ -281,7 +279,7 @@ class MigrationCommand(ToolkitCommand):
             result = deploy_cmd.dry_run_deploy(resources, crud, False, False)
         else:
             result = deploy_cmd.actual_deploy(resources, crud)
-            if result.calculated_total > 0:
+            if result.calculated_total > 0 and creator.HAS_LINEAGE:
                 store_count = creator.store_lineage(resource_list)
                 self.console(f"Stored lineage for {store_count:,} {creator.DISPLAY_NAME}.")
 

cognite_toolkit/_cdf_tk/commands/_migrate/creators.py

@@ -20,6 +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.cruds import NodeCRUD, ResourceCRUD, SpaceCRUD
 from cognite_toolkit._cdf_tk.exceptions import ToolkitRequiredValueError
 from cognite_toolkit._cdf_tk.utils import humanize_collection
@@ -34,10 +35,11 @@ class ResourceConfig:
 
 
 class MigrationCreator(ABC, Generic[T_CogniteResourceList]):
-    """Base class for migration resources configurations that are created from asset-centric resources."""
+    """Base class for migration resources configurations that are created resources."""
 
     CRUD: type[ResourceCRUD]
     DISPLAY_NAME: str
+    HAS_LINEAGE: bool = True
 
     def __init__(self, client: ToolkitClient) -> None:
         self.client = client
@@ -50,7 +52,6 @@ class MigrationCreator(ABC, Generic[T_CogniteResourceList]):
     def resource_configs(self, resources: T_CogniteResourceList) -> list[ResourceConfig]:
         raise NotImplementedError("Subclasses should implement this method")
 
-    @abstractmethod
     def store_lineage(self, resources: T_CogniteResourceList) -> int:
         """Store lineage information for the created resources.
 
@@ -65,6 +66,7 @@ class InstanceSpaceCreator(MigrationCreator[SpaceApplyList]):
 
     CRUD = SpaceCRUD
     DISPLAY_NAME = "Instance Space"
+    HAS_LINEAGE = True
 
     def __init__(
         self, client: ToolkitClient, datasets: DataSetList | None = None, data_set_external_ids: list[str] | None = None
@@ -205,3 +207,25 @@ class SourceSystemCreator(MigrationCreator[NodeApplyList]):
     def store_lineage(self, resources: NodeApplyList) -> int:
         # We already store lineage when creating the resources.
         return len(resources)
+
+
+class InfieldV2ConfigCreator(MigrationCreator[NodeApplyList]):
+    CRUD = NodeCRUD
+    DISPLAY_NAME = "Infield V2 Configuration"
+    HAS_LINEAGE = False
+
+    def create_resources(self) -> NodeApplyList:
+        apm_config_nodes = self.client.data_modeling.instances.list(instance_type="node", sources=APMConfig.view_id)
+        apm_config = APMConfigList.from_nodes(apm_config_nodes)
+
+        new_config_nodes = NodeApplyList([])
+        for config in apm_config:
+            new_config = self._create_infield_v2_config(config)
+            new_config_nodes.append(new_config)
+        return new_config_nodes
+
+    def resource_configs(self, resources: NodeApplyList) -> list[ResourceConfig]:
+        return [ResourceConfig(filestem=node.external_id, data=node.dump()) for node in resources]
+
+    def _create_infield_v2_config(self, config: APMConfig) -> NodeApply:
+        raise NotImplementedError("To be implemented")

cognite_toolkit/_cdf_tk/commands/_migrate/data_classes.py

@@ -1,10 +1,17 @@
+from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Literal
+from typing import Any, Generic, Literal
 
-from cognite.client.data_classes.data_modeling import NodeId, ViewId
+from cognite.client.data_classes._base import (
+    T_WritableCogniteResource,
+    WriteableCogniteResource,
+    WriteableCogniteResourceList,
+)
+from cognite.client.data_classes.data_modeling import InstanceApply, NodeId, ViewId
 from cognite.client.utils._text import to_camel_case
 from pydantic import BaseModel, 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
@@ -12,7 +19,7 @@ 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 AssetCentricType
+from cognite_toolkit._cdf_tk.utils.useful_types import AssetCentricKind, AssetCentricType, JsonVal
 
 
 class MigrationMapping(BaseModel, alias_generator=to_camel_case, extra="ignore", populate_by_name=True):
@@ -122,14 +129,14 @@ class MigrationMappingList(ModelList[MigrationMapping]):
         return {mapping.id: mapping for mapping in self}
 
     @classmethod
-    def read_csv_file(cls, filepath: Path, resource_type: str | None = None) -> "MigrationMappingList":
+    def read_csv_file(cls, filepath: Path, resource_type: AssetCentricKind | 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]] = {
-            "asset": AssetMigrationMappingList,
-            "timeseries": TimeSeriesMigrationMappingList,
-            "file": FileMigrationMappingList,
-            "event": EventMigrationMappingList,
+            "Assets": AssetMigrationMappingList,
+            "TimeSeries": TimeSeriesMigrationMappingList,
+            "FileMetadata": FileMigrationMappingList,
+            "Events": EventMigrationMappingList,
         }
         if resource_type not in cls_by_resource_type:
             raise ToolkitValueError(
@@ -176,3 +183,30 @@ class TimeSeriesMigrationMappingList(MigrationMappingList):
     @classmethod
     def _get_base_model_cls(cls) -> type[TimeSeriesMapping]:
         return TimeSeriesMapping
+
+
+@dataclass
+class AssetCentricMapping(Generic[T_WritableCogniteResource], WriteableCogniteResource[InstanceApply]):
+    mapping: MigrationMapping
+    resource: T_WritableCogniteResource
+
+    def as_write(self) -> InstanceApply:
+        raise NotImplementedError()
+
+    def dump(self, camel_case: bool = True) -> dict[str, JsonVal]:
+        mapping = self.mapping.model_dump(exclude_unset=True, by_alias=camel_case)
+        # Ensure that resource type is always included, even if unset.
+        mapping["resourceType" if camel_case else "resource_type"] = self.mapping.resource_type
+        return {
+            "mapping": mapping,
+            "resource": self.resource.dump(camel_case=camel_case),
+        }
+
+
+class AssetCentricMappingList(
+    WriteableCogniteResourceList[InstanceApply, AssetCentricMapping[T_WritableCogniteResource]]
+):
+    _RESOURCE: type = AssetCentricMapping
+
+    def as_write(self) -> InstanceApplyList:
+        return InstanceApplyList([item.as_write() for item in self])

cognite_toolkit/_cdf_tk/commands/_migrate/data_mapper.py

@@ -8,12 +8,10 @@ from cognite.client.data_classes.data_modeling import DirectRelationReference, I
 
 from cognite_toolkit._cdf_tk.client import ToolkitClient
 from cognite_toolkit._cdf_tk.client.data_classes.migration import ResourceViewMapping
-from cognite_toolkit._cdf_tk.commands._migrate.adapter import (
-    AssetCentricMapping,
-    MigrationSelector,
-)
 from cognite_toolkit._cdf_tk.commands._migrate.conversion import asset_centric_to_dm
+from cognite_toolkit._cdf_tk.commands._migrate.data_classes import AssetCentricMapping
 from cognite_toolkit._cdf_tk.commands._migrate.issues import ConversionIssue, MigrationIssue
+from cognite_toolkit._cdf_tk.commands._migrate.selectors import AssetCentricMigrationSelector
 from cognite_toolkit._cdf_tk.constants import MISSING_INSTANCE_SPACE
 from cognite_toolkit._cdf_tk.exceptions import ToolkitValueError
 from cognite_toolkit._cdf_tk.storageio._base import T_Selector, T_WriteCogniteResource
@@ -45,7 +43,7 @@ class DataMapper(Generic[T_Selector, T_CogniteResource, T_WriteCogniteResource],
         raise NotImplementedError("Subclasses must implement this method.")
 
 
-class AssetCentricMapper(DataMapper[MigrationSelector, AssetCentricMapping, InstanceApply]):
+class AssetCentricMapper(DataMapper[AssetCentricMigrationSelector, AssetCentricMapping, InstanceApply]):
     def __init__(self, client: ToolkitClient) -> None:
         self.client = client
         self._ingestion_view_by_id: dict[ViewId, View] = {}
@@ -55,7 +53,7 @@ class AssetCentricMapper(DataMapper[MigrationSelector, AssetCentricMapping, Inst
         self._asset_mapping_by_id: dict[int, DirectRelationReference] = {}
         self._source_system_mapping_by_id: dict[str, DirectRelationReference] = {}
 
-    def prepare(self, source_selector: MigrationSelector) -> None:
+    def prepare(self, source_selector: AssetCentricMigrationSelector) -> None:
         ingestion_view_ids = source_selector.get_ingestion_mappings()
         ingestion_views = self.client.migration.resource_view_mapping.retrieve(ingestion_view_ids)
         self._view_mapping_by_id = {view.external_id: view for view in ingestion_views}