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

cognite_toolkit/_cdf_tk/commands/_migrate/__init__.py

@@ -1,13 +1,7 @@
-from .assets import MigrateAssetsCommand
-from .canvas import MigrationCanvasCommand
-from .files import MigrateFilesCommand
+from .command import MigrationCommand
 from .prepare import MigrationPrepareCommand
-from .timeseries import MigrateTimeseriesCommand
 
 __all__ = [
-    "MigrateAssetsCommand",
-    "MigrateFilesCommand",
-    "MigrateTimeseriesCommand",
-    "MigrationCanvasCommand",
+    "MigrationCommand",
     "MigrationPrepareCommand",
 ]

cognite_toolkit/_cdf_tk/commands/_migrate/command.py

@@ -0,0 +1,309 @@
+from collections.abc import Callable, Iterable, Sequence
+from enum import Enum
+from pathlib import Path
+
+from rich import print
+from rich.console import Console
+from rich.table import Table
+
+from cognite_toolkit._cdf_tk.client import ToolkitClient
+from cognite_toolkit._cdf_tk.commands._base import ToolkitCommand
+from cognite_toolkit._cdf_tk.commands._migrate.creators import MigrationCreator
+from cognite_toolkit._cdf_tk.commands._migrate.data_mapper import DataMapper
+from cognite_toolkit._cdf_tk.commands.deploy import DeployCommand
+from cognite_toolkit._cdf_tk.constants import DMS_INSTANCE_LIMIT_MARGIN
+from cognite_toolkit._cdf_tk.cruds import ResourceWorker
+from cognite_toolkit._cdf_tk.data_classes import DeployResults
+from cognite_toolkit._cdf_tk.exceptions import (
+    ToolkitFileExistsError,
+    ToolkitMigrationError,
+    ToolkitValueError,
+)
+from cognite_toolkit._cdf_tk.protocols import T_ResourceRequest, T_ResourceResponse
+from cognite_toolkit._cdf_tk.storageio import T_Selector, UploadableStorageIO, UploadItem
+from cognite_toolkit._cdf_tk.utils import humanize_collection, safe_write, sanitize_filename
+from cognite_toolkit._cdf_tk.utils.file import yaml_safe_dump
+from cognite_toolkit._cdf_tk.utils.fileio import Chunk, CSVWriter, NDJsonWriter, SchemaColumn, Uncompressed
+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 .data_model import INSTANCE_SOURCE_VIEW_ID, MODEL_ID, RESOURCE_VIEW_MAPPING_VIEW_ID
+
+
+class MigrationCommand(ToolkitCommand):
+    class Steps(str, Enum):
+        DOWNLOAD = "download"
+        CONVERT = "convert"
+        UPLOAD = "upload"
+
+        @classmethod
+        def list(cls) -> list[str]:
+            return [step.value for step in cls.__members__.values()]
+
+    def migrate(
+        self,
+        selected: T_Selector,
+        data: UploadableStorageIO[T_Selector, T_ResourceResponse, T_ResourceRequest],
+        mapper: DataMapper[T_Selector, T_ResourceResponse, T_ResourceRequest],
+        log_dir: Path,
+        dry_run: bool = False,
+        verbose: bool = False,
+    ) -> ProgressTracker[str]:
+        if log_dir.exists() and any(log_dir.iterdir()):
+            raise ToolkitFileExistsError(
+                f"Log directory {log_dir} already exists. Please remove it or choose another directory."
+            )
+        self.validate_migration_model_available(data.client)
+        log_dir.mkdir(parents=True, exist_ok=True)
+        mapper.prepare(selected)
+
+        iteration_count: int | None = None
+        total_items = data.count(selected)
+        if total_items is not None:
+            iteration_count = (total_items // data.CHUNK_SIZE) + (1 if total_items % data.CHUNK_SIZE > 0 else 0)
+            self.validate_available_capacity(data.client, total_items)
+
+        console = Console()
+        tracker = ProgressTracker[str](self.Steps.list())
+        with (
+            NDJsonWriter(log_dir, kind=f"{selected.kind}MigrationIssues", compression=Uncompressed) as log_file,
+            HTTPClient(config=data.client.config) as write_client,
+        ):
+            executor = ProducerWorkerExecutor[Sequence[T_ResourceResponse], Sequence[UploadItem[T_ResourceRequest]]](
+                download_iterable=self._download_iterable(selected, data, tracker),
+                process=self._convert(mapper, data, tracker, log_file),
+                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}",
+                process_description="Converting",
+                write_description="Uploading",
+                console=console,
+                verbose=verbose,
+            )
+
+            executor.run()
+            total = executor.total_items
+
+        self._print_table(tracker.aggregate(), console)
+        self._print_csv(tracker, log_dir, f"{selected.kind}Items", console)
+        executor.raise_on_error()
+        action = "Would migrate" if dry_run else "Migrating"
+        console.print(f"{action} {total:,} {selected.display_name} to instances.")
+        return tracker
+
+    def _print_table(self, results: dict[tuple[str, Status], int], console: Console) -> None:
+        for step in self.Steps:
+            # We treat pending as failed for summary purposes
+            results[(step.value, "failed")] = results.get((step.value, "failed"), 0) + results.get(
+                (step.value, "pending"), 0
+            )
+
+        table = Table(title="Migration Summary", show_lines=True)
+        table.add_column("Status", style="cyan", no_wrap=True)
+        for step in self.Steps:
+            table.add_column(step.value.capitalize(), style="magenta")
+        for status in AVAILABLE_STATUS:
+            if status == "pending":
+                # Skip pending as we treat it as failed
+                continue
+            row = [status]
+            for step in self.Steps:
+                row.append(str(results.get((step.value, status), 0)))
+            table.add_row(*row)
+
+        console.print(table)
+
+    def _print_csv(self, tracker: ProgressTracker[str], log_dir: Path, kind: str, console: Console) -> None:
+        with CSVWriter(log_dir, kind=kind, compression=Uncompressed, columns=self._csv_columns()) as csv_file:
+            batch: list[Chunk] = []
+            steps = self.Steps.list()
+            for item_id, progress in tracker.result().items():
+                batch.append({"ID": str(item_id), **{step: progress[step] for step in steps}})
+                if len(batch) >= 1000:
+                    csv_file.write_chunks(batch)
+                    batch = []
+            if batch:
+                csv_file.write_chunks(batch)
+        console.print(f"Migration items written to {log_dir}")
+
+    @classmethod
+    def _csv_columns(cls) -> list[SchemaColumn]:
+        return [
+            SchemaColumn(name="ID", type="string"),
+            *(SchemaColumn(name=step, type="string") for step in cls.Steps.list()),
+        ]
+
+    def _download_iterable(
+        self,
+        selected: T_Selector,
+        data: UploadableStorageIO[T_Selector, T_ResourceResponse, T_ResourceRequest],
+        tracker: ProgressTracker[str],
+    ) -> Iterable[Sequence[T_ResourceResponse]]:
+        for page in data.stream_data(selected):
+            for item in page.items:
+                tracker.set_progress(data.as_id(item), self.Steps.DOWNLOAD, "success")
+            yield page.items
+
+    def _convert(
+        self,
+        mapper: DataMapper[T_Selector, T_ResourceResponse, T_ResourceRequest],
+        data: UploadableStorageIO[T_Selector, T_ResourceResponse, T_ResourceRequest],
+        tracker: ProgressTracker[str],
+        log_file: NDJsonWriter,
+    ) -> Callable[[Sequence[T_ResourceResponse]], Sequence[UploadItem[T_ResourceRequest]]]:
+        def track_mapping(source: Sequence[T_ResourceResponse]) -> list[UploadItem[T_ResourceRequest]]:
+            mapped = mapper.map(source)
+            issues: list[Chunk] = []
+            targets: list[UploadItem[T_ResourceRequest]] = []
+
+            for (target, issue), item in zip(mapped, source):
+                id_ = data.as_id(item)
+                result: Status = "failed" if target is None else "success"
+                tracker.set_progress(id_, step=self.Steps.CONVERT, status=result)
+
+                if issue.has_issues:
+                    # MyPy fails to understand that dict[str, JsonVal] is a Chunk
+                    issues.append(issue.dump())  # type: ignore[arg-type]
+                if target is not None:
+                    targets.append(UploadItem(source_id=id_, item=target))
+            if issues:
+                log_file.write_chunks(issues)
+            return targets
+
+        return track_mapping
+
+    def _upload(
+        self,
+        selected: T_Selector,
+        write_client: HTTPClient,
+        target: UploadableStorageIO[T_Selector, T_ResourceResponse, T_ResourceRequest],
+        tracker: ProgressTracker[str],
+        log_file: NDJsonWriter,
+        dry_run: bool,
+    ) -> Callable[[Sequence[UploadItem[T_ResourceRequest]]], None]:
+        def upload_items(data_item: Sequence[UploadItem[T_ResourceRequest]]) -> None:
+            if not data_item:
+                return None
+            responses: Sequence[HTTPMessage]
+            if dry_run:
+                responses = [
+                    SuccessResponseItems(
+                        status_code=200, body="", content=b"", ids=[item.source_id for item in data_item]
+                    )
+                ]
+            else:
+                responses = target.upload_items(data_chunk=data_item, http_client=write_client, selector=selected)
+
+            issues: list[Chunk] = []
+            for item in responses:
+                if isinstance(item, SuccessResponseItems):
+                    for success_id in item.ids:
+                        tracker.set_progress(success_id, step=self.Steps.UPLOAD, status="success")
+                elif isinstance(item, ItemMessage):
+                    for failed_id in item.ids:
+                        tracker.set_progress(failed_id, step=self.Steps.UPLOAD, status="failed")
+
+                if not isinstance(item, SuccessResponseItems):
+                    issues.append(item.dump())  # type: ignore[arg-type]
+            if issues:
+                log_file.write_chunks(issues)
+            return None
+
+        return upload_items
+
+    @staticmethod
+    def validate_migration_model_available(client: ToolkitClient) -> None:
+        models = client.data_modeling.data_models.retrieve([MODEL_ID], inline_views=False)
+        if not models:
+            raise ToolkitMigrationError(
+                f"The migration data model {MODEL_ID!r} does not exist. "
+                "Please run the `cdf migrate prepare` command to deploy the migration data model."
+            )
+        elif len(models) > 1:
+            raise ToolkitMigrationError(
+                f"Multiple migration models {MODEL_ID!r}. "
+                "Please delete the duplicate models before proceeding with the migration."
+            )
+        model = models[0]
+        missing_views = {INSTANCE_SOURCE_VIEW_ID, RESOURCE_VIEW_MAPPING_VIEW_ID} - set(model.views or [])
+        if missing_views:
+            raise ToolkitMigrationError(
+                f"Invalid migration model. Missing views {humanize_collection(missing_views)}. "
+                f"Please run the `cdf migrate prepare` command to deploy the migration data model."
+            )
+
+    def validate_available_capacity(self, client: ToolkitClient, instance_count: int) -> None:
+        """Validate that the project has enough capacity to accommodate the migration."""
+
+        stats = client.data_modeling.statistics.project()
+
+        available_capacity = stats.instances.instances_limit - stats.instances.instances
+        available_capacity_after = available_capacity - instance_count
+
+        if available_capacity_after < DMS_INSTANCE_LIMIT_MARGIN:
+            raise ToolkitValueError(
+                "Cannot proceed with migration, not enough instance capacity available. Total capacity after migration"
+                f" would be {available_capacity_after:,} instances, which is less than the required margin of"
+                f" {DMS_INSTANCE_LIMIT_MARGIN:,} instances. Please increase the instance capacity in your CDF project"
+                f" or delete some existing instances before proceeding with the migration of {instance_count:,} assets."
+            )
+        total_instances = stats.instances.instances + instance_count
+        self.console(
+            f"Project has enough capacity for migration. Total instances after migration: {total_instances:,}."
+        )
+
+    def create(
+        self,
+        client: ToolkitClient,
+        creator: MigrationCreator,
+        dry_run: bool,
+        output_dir: Path,
+        verbose: bool = False,
+    ) -> DeployResults:
+        """This method is used to create migration resource in CDF."""
+        self.validate_migration_model_available(client)
+
+        deploy_cmd = DeployCommand(self.print_warning, silent=self.silent)
+        deploy_cmd.tracker = self.tracker
+
+        crud_cls = creator.CRUD
+        resource_list = creator.create_resources()
+
+        results = DeployResults([], "deploy", dry_run=dry_run)
+        crud = crud_cls.create_loader(client)
+        worker = ResourceWorker(crud, "deploy")
+        local_by_id = {crud.get_id(item): (item.dump(), item) for item in resource_list}
+        worker.validate_access(local_by_id, is_dry_run=dry_run)
+        cdf_resources = crud.retrieve(list(local_by_id.keys()))
+        resources = worker.categorize_resources(local_by_id, cdf_resources, False, verbose)
+
+        if dry_run:
+            result = deploy_cmd.dry_run_deploy(resources, crud, False, False)
+        else:
+            result = deploy_cmd.actual_deploy(resources, crud)
+            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}.")
+
+        verb = "Would deploy" if dry_run else "Deploying"
+        self.console(f"{verb} {creator.DISPLAY_NAME} to CDF.")
+
+        resource_configs = creator.resource_configs(resource_list)
+        for config in resource_configs:
+            filepath = output_dir / crud_cls.folder_name / f"{sanitize_filename(config.filestem)}.{crud_cls.kind}.yaml"
+            filepath.parent.mkdir(parents=True, exist_ok=True)
+            safe_write(filepath, yaml_safe_dump(config.data))
+        self.console(
+            f"{len(resource_configs)} {crud_cls.kind} resource configurations written to {(output_dir / crud_cls.folder_name).as_posix()!r}"
+        )
+        self.console("It is recommended to add these files to a Toolkit governed module.")
+
+        if result:
+            results[result.name] = result
+
+        if results.has_counts:
+            print(results.counts_table())
+
+        return results