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

cognite_toolkit/_cdf_tk/client/resource_classes/workflow_version.py

@@ -0,0 +1,155 @@
+from typing import Annotated, Any, Literal
+
+from pydantic import Field, JsonValue, field_validator
+from pydantic_core.core_schema import ValidationInfo
+
+from cognite_toolkit._cdf_tk.client._resource_base import (
+    BaseModelObject,
+    RequestResource,
+    ResponseResource,
+)
+
+from .identifiers import Identifier, WorkflowVersionId
+
+
+class TaskId(Identifier):
+    external_id: str
+
+    def __str__(self) -> str:
+        return f"externalId='{self.external_id}'"
+
+
+class TaskParameterDefinition(BaseModelObject):
+    type: str
+
+
+class CogniteFunctionRef(BaseModelObject):
+    external_id: str
+    data: JsonValue | str | None = None
+
+
+class FunctionTaskParameters(TaskParameterDefinition):
+    type: Literal["function"] = Field("function", exclude=True)
+    function: CogniteFunctionRef
+    is_async_complete: bool | None = None
+
+
+class TransformationRef(BaseModelObject):
+    external_id: str
+    concurrency_policy: Literal["fail", "waitForCurrent", "restartAfterCurrent"] | None = None
+    use_transformation_credentials: bool | None = None
+
+
+class TransformationTaskParameters(TaskParameterDefinition):
+    type: Literal["transformation"] = Field("transformation", exclude=True)
+    transformation: TransformationRef
+
+
+class CDFRequest(BaseModelObject):
+    resource_path: str
+    method: str
+    query_parameters: JsonValue | str | None = None
+    body: JsonValue | str | None = None
+    request_timeout_in_millis: float | str | None = None
+    cdf_version_header: str | None = None
+
+
+class CDFTaskParameters(TaskParameterDefinition):
+    type: Literal["cdf"] = Field("cdf", exclude=True)
+    cdf_request: CDFRequest
+
+
+class DynamicRef(BaseModelObject):
+    tasks: str
+
+
+class DynamicTaskParameters(TaskParameterDefinition):
+    type: Literal["dynamic"] = Field("dynamic", exclude=True)
+    dynamic: DynamicRef
+
+
+class SubworkflowRef(BaseModelObject):
+    tasks: "list[Task]"
+
+
+class SubworkflowTaskParameters(TaskParameterDefinition):
+    type: Literal["subworkflow"] = Field("subworkflow", exclude=True)
+    subworkflow: WorkflowVersionId | SubworkflowRef
+
+
+class SimulatorInputUnit(BaseModelObject):
+    name: str
+
+
+class SimulatorInput(BaseModelObject):
+    reference_id: str
+    value: str | int | float | list[str] | list[int] | list[float]
+    unit: SimulatorInputUnit | None = None
+
+
+class SimulationRef(BaseModelObject):
+    routine_external_id: str
+    run_time: int | None = None
+    inputs: list[SimulatorInput] | None = None
+
+
+class SimulationTaskParameters(TaskParameterDefinition):
+    type: Literal["simulation"] = Field("simulation", exclude=True)
+    simulation: SimulationRef
+
+
+Parameter = Annotated[
+    FunctionTaskParameters
+    | TransformationTaskParameters
+    | CDFTaskParameters
+    | DynamicTaskParameters
+    | SubworkflowTaskParameters
+    | SimulationTaskParameters,
+    Field(discriminator="type"),
+]
+
+
+class Task(BaseModelObject):
+    external_id: str
+    type: str
+    name: str | None = None
+    description: str | None = None
+    retries: int | None = None
+    timeout: int | None = None
+    on_failure: Literal["abortWorkflow", "skipTask"] = "abortWorkflow"
+    depends_on: list[TaskId] | None = None
+    parameters: Parameter | None = None
+
+    @field_validator("parameters", mode="before")
+    @classmethod
+    def move_type_to_field(cls, value: Any, info: ValidationInfo) -> Any:
+        if not isinstance(value, dict) or "type" not in info.data:
+            return value
+        value = dict(value)
+        value["type"] = info.data["type"]
+        return value
+
+
+class WorkflowDefinition(BaseModelObject):
+    description: str | None = None
+    tasks: list[Task]
+
+
+class WorkflowVersion(BaseModelObject):
+    workflow_external_id: str
+    version: str
+    workflow_definition: WorkflowDefinition
+
+    def as_id(self) -> WorkflowVersionId:
+        return WorkflowVersionId(workflow_external_id=self.workflow_external_id, version=self.version)
+
+
+class WorkflowVersionRequest(WorkflowVersion, RequestResource): ...
+
+
+class WorkflowVersionResponse(WorkflowVersion, ResponseResource[WorkflowVersionRequest]):
+    created_time: int
+    last_updated_time: int
+
+    def as_request_resource(self) -> WorkflowVersionRequest:
+        return WorkflowVersionRequest.model_validate(self.dump(), extra="ignore")

cognite_toolkit/_cdf_tk/client/testing.py

@@ -7,22 +7,39 @@ from cognite.client._api.datapoints import DatapointsAPI
 from cognite.client._api.datapoints_subscriptions import DatapointsSubscriptionAPI
 from cognite.client._api.functions import FunctionCallsAPI, FunctionSchedulesAPI
 from cognite.client._api.raw import RawDatabasesAPI, RawRowsAPI, RawTablesAPI
+from cognite.client._api.simulators import SimulatorModelsAPI, SimulatorsAPI
 from cognite.client._api.synthetic_time_series import SyntheticDatapointsAPI
 from cognite.client.testing import CogniteClientMock
 from rich.console import Console
 
 from cognite_toolkit._cdf_tk.client._toolkit_client import ToolkitClient
-
-from .api.canvas import CanvasAPI, IndustrialCanvasAPI
-from .api.charts import ChartsAPI
-from .api.dml import DMLAPI
-from .api.extended_data_modeling import ExtendedInstancesAPI
-from .api.extended_files import ExtendedFileMetadataAPI
-from .api.extended_functions import ExtendedFunctionsAPI
-from .api.extended_raw import ExtendedRawAPI
-from .api.extended_timeseries import ExtendedTimeSeriesAPI
+from cognite_toolkit._cdf_tk.client.api.legacy.canvas import CanvasAPI, IndustrialCanvasAPI
+from cognite_toolkit._cdf_tk.client.api.legacy.charts import ChartsAPI
+from cognite_toolkit._cdf_tk.client.api.legacy.dml import DMLAPI
+from cognite_toolkit._cdf_tk.client.api.legacy.extended_data_modeling import ExtendedInstancesAPI
+from cognite_toolkit._cdf_tk.client.api.legacy.extended_files import ExtendedFileMetadataAPI
+from cognite_toolkit._cdf_tk.client.api.legacy.extended_functions import ExtendedFunctionsAPI
+from cognite_toolkit._cdf_tk.client.api.legacy.extended_raw import ExtendedRawAPI
+from cognite_toolkit._cdf_tk.client.api.legacy.extended_timeseries import ExtendedTimeSeriesAPI
+from cognite_toolkit._cdf_tk.client.api.legacy.location_filters import LocationFiltersAPI
+from cognite_toolkit._cdf_tk.client.api.legacy.robotics import (
+    CapabilitiesAPI,
+    DataPostProcessingAPI,
+    FramesAPI,
+    MapsAPI,
+    RoboticsAPI,
+)
+from cognite_toolkit._cdf_tk.client.api.legacy.robotics import LocationsAPI as RoboticsLocationsAPI
+from cognite_toolkit._cdf_tk.client.api.legacy.search_config import SearchConfigurationsAPI
+
+from ._toolkit_client import ToolAPI
+from .api.assets import AssetsAPI
+from .api.datasets import DataSetsAPI
+from .api.events import EventsAPI
+from .api.extraction_pipelines import ExtractionPipelinesAPI
+from .api.filemetadata import FileMetadataAPI
 from .api.infield import InfieldAPI, InFieldCDMConfigAPI, InfieldConfigAPI
-from .api.location_filters import LocationFiltersAPI
+from .api.labels import LabelsAPI
 from .api.lookup import (
     AssetLookUpAPI,
     DataSetLookUpAPI,
@@ -44,17 +61,18 @@ from .api.migration import (
     ResourceViewMappingAPI,
 )
 from .api.project import ProjectAPI
-from .api.robotics import RoboticsAPI
-from .api.robotics.capabilities import CapabilitiesAPI
-from .api.robotics.data_postprocessing import DataPostProcessingAPI
-from .api.robotics.frames import FramesAPI
-from .api.robotics.locations import LocationsAPI as RoboticsLocationsAPI
-from .api.robotics.maps import MapsAPI
 from .api.search import SearchAPI
-from .api.search_config import SearchConfigurationsAPI
+from .api.security_categories import SecurityCategoriesAPI
+from .api.sequences import SequencesAPI
 from .api.streams import StreamsAPI
+from .api.three_d import ThreeDAPI, ThreeDClassicModelsAPI
+from .api.timeseries import TimeSeriesAPI
 from .api.token import TokenAPI
+from .api.transformations import TransformationsAPI
 from .api.verify import VerifyAPI
+from .api.workflow_triggers import WorkflowTriggersAPI
+from .api.workflow_versions import WorkflowVersionsAPI
+from .api.workflows import WorkflowsAPI
 
 
 class ToolkitClientMock(CogniteClientMock):
@@ -131,6 +149,25 @@ class ToolkitClientMock(CogniteClientMock):
         self.time_series.data.synthetic = MagicMock(spec_set=SyntheticDatapointsAPI)
         self.time_series.subscriptions = MagicMock(spec_set=DatapointsSubscriptionAPI)
 
+        self.tool = MagicMock(spec=ToolAPI)
+        self.tool.three_d = MagicMock(spec=ThreeDAPI)
+        self.tool.three_d.models_classic = MagicMock(spec_set=ThreeDClassicModelsAPI)
+        self.tool.assets = MagicMock(spec_set=AssetsAPI)
+        self.tool.timeseries = MagicMock(spec_set=TimeSeriesAPI)
+        self.tool.filemetadata = MagicMock(spec_set=FileMetadataAPI)
+        self.tool.events = MagicMock(spec_set=EventsAPI)
+        self.tool.simulators = MagicMock(spec=SimulatorsAPI)
+        self.tool.simulators.models = MagicMock(spec_set=SimulatorModelsAPI)
+        self.tool.datasets = MagicMock(spec_set=DataSetsAPI)
+        self.tool.extraction_pipelines = MagicMock(spec_set=ExtractionPipelinesAPI)
+        self.tool.labels = MagicMock(spec_set=LabelsAPI)
+        self.tool.security_categories = MagicMock(spec_set=SecurityCategoriesAPI)
+        self.tool.sequences = MagicMock(spec_set=SequencesAPI)
+        self.tool.transformations = MagicMock(spec_set=TransformationsAPI)
+        self.tool.workflows = MagicMock(spec=WorkflowsAPI)
+        self.tool.workflows.triggers = MagicMock(spec_set=WorkflowTriggersAPI)
+        self.tool.workflows.versions = MagicMock(spec_set=WorkflowVersionsAPI)
+
         self.streams = MagicMock(spec=StreamsAPI)
 
         # This is a helper API, not a real API.

cognite_toolkit/_cdf_tk/commands/__init__.py

@@ -24,6 +24,7 @@ __all__ = [
     "AboutCommand",
     "AuthCommand",
     "BuildCommand",
+    "BuildCommandV2",
     "CleanCommand",
     "CollectCommand",
     "DeployCommand",

cognite_toolkit/_cdf_tk/commands/_migrate/command.py

@@ -1,12 +1,19 @@
-from collections.abc import Callable, Iterable, Sequence
-from enum import Enum
+from collections.abc import Callable, Sequence
+from dataclasses import dataclass
 from pathlib import Path
+from typing import get_args
 
 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.client.http_client import (
+    FailedRequestItems,
+    FailedResponseItems,
+    HTTPClient,
+    SuccessResponseItems,
+)
 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
@@ -21,26 +28,30 @@ from cognite_toolkit._cdf_tk.exceptions import (
 )
 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.storageio.logger import FileDataLogger, OperationStatus
 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.fileio import NDJsonWriter, Uncompressed
 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
+from .issues import WriteIssue
 
 
-class MigrationCommand(ToolkitCommand):
-    class Steps(str, Enum):
-        DOWNLOAD = "download"
-        CONVERT = "convert"
-        UPLOAD = "upload"
+@dataclass
+class OperationIssue:
+    message: str
+    count: int
+
 
-        @classmethod
-        def list(cls) -> list[str]:
-            return [step.value for step in cls.__members__.values()]
+@dataclass
+class MigrationStatusResult:
+    status: OperationStatus
+    issues: list[OperationIssue]
+    count: int
 
+
+class MigrationCommand(ToolkitCommand):
     def migrate(
         self,
         selected: T_Selector,
@@ -49,7 +60,7 @@ class MigrationCommand(ToolkitCommand):
         log_dir: Path,
         dry_run: bool = False,
         verbose: bool = False,
-    ) -> ProgressTracker[str]:
+    ) -> list[MigrationStatusResult]:
         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."
@@ -65,15 +76,18 @@ class MigrationCommand(ToolkitCommand):
             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,
         ):
+            logger = FileDataLogger(log_file)
+            data.logger = logger
+            mapper.logger = logger
+
             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),
+                download_iterable=(page.items for page in data.stream_data(selected)),
+                process=self._convert(mapper, data),
+                write=self._upload(selected, write_client, data, dry_run),
                 iteration_count=iteration_count,
                 max_queue_size=10,
                 download_description=f"Downloading {selected.display_name}",
@@ -86,91 +100,71 @@ class MigrationCommand(ToolkitCommand):
             executor.run()
             total = executor.total_items
 
-        self._print_table(tracker.aggregate(), console)
-        self._print_csv(tracker, log_dir, f"{selected.kind}Items", console)
+        results = self._create_status_summary(logger)
+
+        self._print_rich_tables(results, console)
+        self._print_txt(results, 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
+        return results
+
+    # Todo: Move to the logger module
+    @classmethod
+    def _create_status_summary(cls, logger: FileDataLogger) -> list[MigrationStatusResult]:
+        results: list[MigrationStatusResult] = []
+        status_counts = logger.tracker.get_status_counts()
+        for status in get_args(OperationStatus):
+            issue_counts = logger.tracker.get_issue_counts(status)
+            issues = [OperationIssue(message=issue, count=count) for issue, count in issue_counts.items()]
+            result = MigrationStatusResult(
+                status=status,
+                issues=issues,
+                count=status_counts.get(status, 0),
             )
+            results.append(result)
+        return results
 
+    def _print_rich_tables(self, results: list[MigrationStatusResult], console: Console) -> None:
         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)
-
+        table.add_column("Status", style="bold")
+        table.add_column("Count", justify="right", style="bold")
+        table.add_column("Issues", style="bold")
+        for result in results:
+            issues_str = "\n".join(f"{issue.message}: {issue.count}" for issue in result.issues) or ""
+            table.add_row(result.status, str(result.count), issues_str)
         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 _print_txt(self, results: list[MigrationStatusResult], log_dir: Path, kind: str, console: Console) -> None:
+        summary_file = log_dir / f"{kind}_migration_summary.txt"
+        with summary_file.open("w", encoding="utf-8") as f:
+            f.write("Migration Summary\n")
+            f.write("=================\n\n")
+            for result in results:
+                f.write(f"Status: {result.status}\n")
+                f.write(f"Count: {result.count}\n")
+                f.write("Issues:\n")
+                if result.issues:
+                    for issue in result.issues:
+                        f.write(f"  - {issue.message}: {issue.count}\n")
+                else:
+                    f.write("  None\n")
+                f.write("\n")
+        console.print(f"Summary written to {log_dir}")
 
+    @staticmethod
     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 [
+                UploadItem(source_id=data.as_id(item), item=target)
+                for target, item in zip(mapped, source)
+                if target is not None
+            ]
 
         return track_mapping
 
@@ -179,36 +173,37 @@ class MigrationCommand(ToolkitCommand):
         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] = []
+                target.logger.tracker.finalize_item([item.source_id for item in data_item], "pending")
+                return None
+
+            responses = target.upload_items(data_chunk=data_item, http_client=write_client, selector=selected)
+
+            # Todo: Move logging into the UploadableStorageIO class
+            issues: list[WriteIssue] = []
             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]
+                    target.logger.tracker.finalize_item(item.ids, "success")
+                    continue
+                if isinstance(item, FailedResponseItems):
+                    error = item.error
+                    for id_ in item.ids:
+                        issue = WriteIssue(id=str(id_), status_code=error.code, message=error.message)
+                        issues.append(issue)
+                elif isinstance(item, FailedRequestItems):
+                    for id_ in item.ids:
+                        issue = WriteIssue(id=str(id_), status_code=0, message=item.error)
+                        issues.append(issue)
+
+                if isinstance(item, FailedResponseItems | FailedRequestItems):
+                    target.logger.tracker.finalize_item(item.ids, "failure")
             if issues:
-                log_file.write_chunks(issues)
+                target.logger.log(issues)
             return None
 
         return upload_items