[CDF-27582] 🚨Check instance space existence.

cognite_toolkit/_cdf_tk/commands/_migrate/migration_io.py

@@ -14,17 +14,18 @@
     ItemsResultList,
     ItemsSuccessResponse,
 )
-from cognite_toolkit._cdf_tk.client.identifiers import InternalId
+from cognite_toolkit._cdf_tk.client.identifiers import InternalId, SpaceId
 from cognite_toolkit._cdf_tk.client.resource_classes.annotation import AnnotationResponse
 from cognite_toolkit._cdf_tk.client.resource_classes.data_modeling import EdgeId, InstanceRequest, NodeId
+from cognite_toolkit._cdf_tk.client.resource_classes.migration import SpaceSource
 from cognite_toolkit._cdf_tk.client.resource_classes.pending_instance_id import PendingInstanceId
 from cognite_toolkit._cdf_tk.client.resource_classes.three_d import (
     AssetMappingClassicResponse,
     AssetMappingDMRequestId,
     ThreeDModelClassicResponse,
 )
 from cognite_toolkit._cdf_tk.commands._migrate.data_classes import ThreeDMigrationRequest
-from cognite_toolkit._cdf_tk.constants import MISSING_EXTERNAL_ID, MISSING_INSTANCE_SPACE
+from cognite_toolkit._cdf_tk.constants import MISSING_EXTERNAL_ID
 from cognite_toolkit._cdf_tk.exceptions import ToolkitNotImplementedError, ToolkitValueError
 from cognite_toolkit._cdf_tk.storageio import (
     AnnotationIO,
@@ -41,7 +42,7 @@
     ThreeDSelector,
 )
 from cognite_toolkit._cdf_tk.tk_warnings import MediumSeverityWarning
-from cognite_toolkit._cdf_tk.utils.collection import chunker_sequence
+from cognite_toolkit._cdf_tk.utils.collection import chunker_sequence, humanize_collection
 from cognite_toolkit._cdf_tk.utils.useful_types import (
     AssetCentricKindExtended,
     AssetCentricType,
@@ -89,12 +90,28 @@ def stream_data(
         bookmark: Bookmark | None = None,
     ) -> Iterator[Page]:
         file_location = bookmark if isinstance(bookmark, FileBookmark) else None
+
         if isinstance(selector, MigrationCSVFileSelector):
+            instance_spaces = list({SpaceId(space=item.instance_id.space) for item in selector.items})
             iterator = self._stream_from_csv(selector, limit, file_location)
         elif isinstance(selector, MigrateDataSetSelector):
-            iterator = self._stream_given_dataset(selector, limit)
+            space_source = self.client.migration.space_source.retrieve(
+                data_set_external_id=selector.data_set_external_id
+            )
+            if space_source is None:
+                raise ToolkitValueError(
+                    f"Missing instance space that maps to {selector.data_set_external_id!r}. Have you run `cdf migrate data-sets`?"
+                )
+            instance_spaces = [SpaceId(space=space_source.space)]
+            iterator = self._stream_given_dataset(selector, space_source, limit)
         else:
             raise ToolkitNotImplementedError(f"Selector {type(selector)} is not supported for stream_data")
+        existing = self.client.tool.spaces.retrieve(instance_spaces)
+        if missing := set(instance_spaces).difference({item.as_id() for item in existing}):
+            raise ToolkitValueError(
+                f"The following instance spaces do not exist in CDF: {humanize_collection(missing)}. Please create these spaces before running the migration."
+            )
+
         yield from (
             Page(
                 worker_id="main",
@@ -132,11 +149,10 @@ def count(self, selector: AssetCentricMigrationSelector) -> int | None:
             raise ToolkitNotImplementedError(f"Selector {type(selector)} is not supported for count")
 
     def _stream_given_dataset(
-        self, selector: MigrateDataSetSelector, limit: int | None = None
+        self, selector: MigrateDataSetSelector, space_source: SpaceSource, limit: int | None = None
     ) -> Iterator[Sequence[AssetCentricMapping[T_AssetCentricResource]]]:
         asset_centric_selector = selector.as_asset_centric_selector()
-        space_source = self.client.migration.space_source.retrieve(data_set_external_id=selector.data_set_external_id)
-        instance_space = space_source.instance_space if space_source else MISSING_INSTANCE_SPACE
+        instance_space = space_source.instance_space
         for data_chunk in self.hierarchy.stream_data(asset_centric_selector, limit):
             mapping_list: list[AssetCentricMapping[T_AssetCentricResource]] = []
             for data_item in data_chunk.items:

tests/test_unit/test_cdf_tk/test_commands/test_migration_cmd/test_command.py

@@ -31,7 +31,12 @@
 )
 from cognite_toolkit._cdf_tk.client.resource_classes.chart import ChartResponse
 from cognite_toolkit._cdf_tk.client.resource_classes.charts_data import ChartData, ChartSource, ChartTimeseries
-from cognite_toolkit._cdf_tk.client.resource_classes.data_modeling import InstanceSource, NodeRequest, ViewId
+from cognite_toolkit._cdf_tk.client.resource_classes.data_modeling import (
+    InstanceSource,
+    NodeRequest,
+    SpaceResponse,
+    ViewId,
+)
 from cognite_toolkit._cdf_tk.client.resource_classes.migration import InstanceSource as LegacyInstanceSource
 from cognite_toolkit._cdf_tk.client.testing import monkeypatch_toolkit_client
 from cognite_toolkit._cdf_tk.commands._migrate.command import MigrationCommand
@@ -203,6 +208,18 @@ def test_migrate_assets(
             + "\n".join(f"{1000 + i},{space},asset_{i},{ASSET_ID},cdf_cdm,CogniteAsset,v1" for i in range(len(assets)))
         )
 
+        # Retrieve space
+        respx.post(
+            config.create_api_url("/models/spaces/byids"),
+        ).mock(
+            return_value=httpx.Response(
+                status_code=200,
+                json={
+                    "items": [SpaceResponse(space=space, created_time=1, last_updated_time=1, is_global=False).dump()]
+                },
+            )
+        )
+
         # Asset retrieve ids
         respx.post(
             config.create_api_url("/assets/byids"),
@@ -318,6 +335,17 @@ def test_migrate_resume(
             "id,space,externalId,ingestionView,consumerViewSpace,consumerViewExternalId,consumerViewVersion\n"
             f"{assets[0].id},{space},{assets[0].external_id},{ASSET_ID},cdf_cdm,CogniteAsset,v1"
         )
+        # Retrieve space
+        respx.post(
+            config.create_api_url("/models/spaces/byids"),
+        ).mock(
+            return_value=httpx.Response(
+                status_code=200,
+                json={
+                    "items": [SpaceResponse(space=space, created_time=1, last_updated_time=1, is_global=False).dump()]
+                },
+            )
+        )
 
         # Asset retrieve ids
         respx.post(

tests/test_unit/test_cdf_tk/test_commands/test_migration_cmd/test_migration_io.py

@@ -1,13 +1,15 @@
 from pathlib import Path
 from unittest.mock import MagicMock
 
+import httpx
 import pytest
 import respx
 from httpx import Response
 
 from cognite_toolkit._cdf_tk.client import ToolkitClient, ToolkitClientConfig
 from cognite_toolkit._cdf_tk.client.http_client import HTTPClient
 from cognite_toolkit._cdf_tk.client.resource_classes.annotation import AnnotationResponse
+from cognite_toolkit._cdf_tk.client.resource_classes.data_modeling import SpaceResponse
 from cognite_toolkit._cdf_tk.commands._migrate.migration_io import (
     AnnotationMigrationIO,
     AssetCentricMigrationIO,
@@ -39,6 +41,19 @@ def test_download(self, toolkit_client: ToolkitClient, respx_mock: respx.MockRou
             }
             for i in range(N)
         ]
+        respx.post(
+            config.create_api_url("/models/spaces/byids"),
+        ).mock(
+            return_value=httpx.Response(
+                status_code=200,
+                json={
+                    "items": [
+                        SpaceResponse(space="mySpace", created_time=1, last_updated_time=1, is_global=False).dump()
+                    ]
+                },
+            )
+        )
+
         respx_mock.post(config.create_api_url("/assets/byids")).mock(
             side_effect=[
                 Response(status_code=200, json={"items": items[: AssetIO.CHUNK_SIZE]}),