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

cognite_toolkit/_cdf_tk/resource_classes/workflow_version.py

@@ -1,4 +1,5 @@
-from typing import Literal
+from abc import ABC
+from typing import Annotated, Literal
 
 from cognite.client.data_classes import WorkflowVersionUpsert
 from pydantic import Field, JsonValue
@@ -6,18 +7,150 @@ from pydantic import Field, JsonValue
 from .base import BaseModelResource, ToolkitResource
 
 
+class WorkflowVersionId(BaseModelResource):
+    workflow_external_id: str = Field(
+        max_length=255,
+        description="Identifier for a workflow. Must be unique for the project. No trailing or leading whitespace and no null characters allowed.",
+    )
+    version: str = Field(
+        max_length=255,
+        description="Identifier for a version. Must be unique for the workflow. No trailing or leading whitespace and no null characters allowed.",
+    )
+
+
+class CogniteFunctionRef(BaseModelResource):
+    external_id: str = Field(
+        description="The external id of the Cognite Function in the project. This can be either a function external ID or a reference like ${myTaskExternalId.output.someKey}"
+    )
+    data: str | JsonValue | None = Field(
+        None, description="Input data that will be passed to the Cognite Function. Limited to 100KB in size."
+    )
+
+
+class FunctionTaskParameters(BaseModelResource):
+    function: CogniteFunctionRef
+    is_async_complete: bool = Field(
+        False, description="Defines if the execution of the task should be completed asynchronously."
+    )
+
+
+class TransformationRef(BaseModelResource):
+    external_id: str = Field(
+        description="The external id of the Transformation in the project. This can be either a transformation external ID or a reference like ${myTaskExternalId.output.someKey}"
+    )
+    concurrency_policy: Literal["fail", "waitForCurrent", "restartAfterCurrent"] = Field(
+        "fail",
+        description="""Determines the behavior of the task if the Transformation is already running.
+
+fail: The task fails if another instance of the Transformation is currently running.
+waitForCurrent: The task will pause and wait for the already running Transformation to complete. Once completed, the task is completed. This mode is useful for preventing redundant Transformation runs.
+restartAfterCurrent: The task waits for the ongoing Transformation to finish. After completion, the task restarts the Transformation. This mode ensures that the most recent data can be used by following tasks.""",
+    )
+    use_transformation_credentials: bool = Field(
+        False,
+        description="If set to true, the transformation will run using the client credentials configured on the transformation. If set to false, the transformation will run using the client credentials used to trigger the workflow.",
+    )
+
+
+class TransformationTaskParameters(BaseModelResource):
+    transformation: TransformationRef
+
+
+class CDFRequest(BaseModelResource):
+    resource_path: str = Field(
+        description="The path of the request. The path should be prefixed by {cluster}.cognitedata.com/api/v1/project/{project} based on the relevant cluster and project. It can also contain references like ${myTaskExternalId.output.someKey}"
+    )
+    method: Literal["POST", "GET", "PUT"] | str = Field(
+        description="The HTTP method of the request. It can also be a reference like ${myTaskExternalId.output.someKey}"
+    )
+    query_parameters: dict[str, JsonValue] | str | None = Field(
+        None,
+        description="The query parameters of the request. It can also be a reference like ${myTaskExternalId.output.someKey}",
+    )
+    body: JsonValue | str | None = Field(
+        None, description="The body of the request. It can also be a reference like ${myTaskExternalId.output.someKey}"
+    )
+    request_timeout_in_millis: float | str | None = Field(
+        None,
+        description="The timeout for the request in milliseconds. It can also be a reference like ${myTaskExternalId.output.someKey}",
+    )
+    cdf_version_header: Literal["alpha", "beta"] | str | None = Field(
+        None, description="The Cognite Data Fusion version header to use for the request."
+    )
+
+
+class CDFTaskParameters(BaseModelResource):
+    cdf_request: CDFRequest
+
+
+class DynamicRef(BaseModelResource):
+    tasks: str = Field(
+        description="A Reference is an expression that allows dynamically injecting input to a task during execution. References can be used to reference the input of the Workflow, the output of a previous task in the Workflow, or the input of a previous task in the Workflow. Note that the injected value must be valid in the context of the property it is injected into. Example Task reference: ${myTaskExternalId.output.someKey} Example Workflow input reference: ${workflow.input.myKey}"
+    )
+
+
+class DynamicTaskParameters(BaseModelResource):
+    dynamic: DynamicRef = Field(description="Reference to another task to use as the definition for this task.")
+
+
+class SubworkflowInlineTasks(BaseModelResource):
+    """Inline definition of tasks for a subworkflow."""
+
+    tasks: "list[Task]" = Field(description="Inline definition of tasks for the subworkflow.")
+
+
+# SubworkflowRef can be either a reference to an existing workflow version OR an inline definition of tasks
+SubworkflowRef = Annotated[
+    WorkflowVersionId | SubworkflowInlineTasks,
+    Field(
+        description="Reference to the subworkflow to execute. This can be either a reference to an existing workflow version (with workflowExternalId and version) or an inline definition of tasks."
+    ),
+]
+
+
+class SubworkflowTaskParameters(BaseModelResource):
+    subworkflow: SubworkflowRef = Field(description="Reference to the subworkflow to execute.")
+
+
+class SimulatorInputUnit(BaseModelResource):
+    name: str = Field(description="Name of the unit.")
+
+
+class SimulatorInput(BaseModelResource):
+    reference_id: str = Field(description="Reference id of the value to override.")
+    value: str | int | float | list[str] | list[int] | list[float] = Field(
+        description="Override the value used for a simulation run."
+    )
+    unit: SimulatorInputUnit | None = Field(None, description="Override the unit of the value")
+
+
+class SimulationRef(BaseModelResource):
+    routine_external_id: str = Field(description="The external id of the routine to be executed.")
+    run_time: int | None = Field(
+        None,
+        description="Run time in milliseconds. Reference timestamp used for data pre-processing and data sampling.",
+    )
+    inputs: list[SimulatorInput] | None = Field(
+        None, description="List of inputs to be provided to the simulation.", max_length=200
+    )
+
+
+class SimulationTaskParameters(BaseModelResource):
+    simulation: SimulationRef = Field(description="Reference to the simulation to execute.")
+
+
 class TaskId(BaseModelResource):
     external_id: str = Field(
         max_length=255, description="The external ID provided by the client. Must be unique for the resource type."
     )
 
 
-class TaskDefinition(BaseModelResource):
+class TaskDefinition(BaseModelResource, ABC):
     external_id: str = Field(
         max_length=255,
         description="Identifier for the task. Must be unique within the version. No trailing or leading whitespace and no null characters allowed.",
     )
-    type: Literal["function", "transformation", "cdf", "dynamic", "subworkflow", "simulation"]
+    type: str
     name: str | None = Field(
         default=None,
         max_length=255,
@@ -28,7 +161,6 @@ class TaskDefinition(BaseModelResource):
         max_length=500,
         description="Description of the intention of the task",
     )
-    parameters: JsonValue = Field()
     retries: int = Field(
         3,
         ge=0,
@@ -52,13 +184,49 @@ class TaskDefinition(BaseModelResource):
     )
 
 
+class FunctionTask(TaskDefinition):
+    type: Literal["function"] = "function"
+    parameters: FunctionTaskParameters
+
+
+class TransformationTask(TaskDefinition):
+    type: Literal["transformation"] = "transformation"
+    parameters: TransformationTaskParameters
+
+
+class CDFTask(TaskDefinition):
+    type: Literal["cdfRequest"] = "cdfRequest"
+    parameters: CDFTaskParameters
+
+
+class DynamicTask(TaskDefinition):
+    type: Literal["dynamic"] = "dynamic"
+    parameters: DynamicTaskParameters
+
+
+class SubworkflowTask(TaskDefinition):
+    type: Literal["subworkflow"] = "subworkflow"
+    parameters: SubworkflowTaskParameters
+
+
+class SimulationTask(TaskDefinition):
+    type: Literal["simulation"] = "simulation"
+    parameters: SimulationTaskParameters
+
+
+Task = Annotated[
+    FunctionTask | TransformationTask | CDFTask | DynamicTask | SubworkflowTask | SimulationTask,
+    Field(discriminator="type"),
+]
+
+
 class WorkflowDefinition(BaseModelResource):
     description: str | None = Field(
         default=None,
         max_length=500,
         description="The description of the workflow version.",
     )
-    tasks: list[TaskDefinition]
+    tasks: list[Task]
 
 
 class WorkflowVersionYAML(ToolkitResource):

cognite_toolkit/_cdf_tk/storageio/_applications.py

@@ -1,16 +1,25 @@
 from collections.abc import Iterable, Sequence
+from itertools import zip_longest
 from typing import Any
 
+from pydantic import JsonValue
+
 from cognite_toolkit._cdf_tk.client import ToolkitClient
-from cognite_toolkit._cdf_tk.client.data_classes.canvas import (
+from cognite_toolkit._cdf_tk.client.http_client import (
+    HTTPClient,
+    HTTPMessage,
+    HTTPResult2,
+    RequestMessage2,
+    SuccessResponse2,
+)
+from cognite_toolkit._cdf_tk.client.resource_classes.legacy.canvas import (
     IndustrialCanvas,
     IndustrialCanvasApply,
 )
-from cognite_toolkit._cdf_tk.client.data_classes.charts import Chart, ChartList, ChartWrite
+from cognite_toolkit._cdf_tk.client.resource_classes.legacy.charts import Chart, ChartList, ChartWrite
 from cognite_toolkit._cdf_tk.exceptions import ToolkitNotImplementedError
 from cognite_toolkit._cdf_tk.tk_warnings import HighSeverityWarning, MediumSeverityWarning
 from cognite_toolkit._cdf_tk.utils.collection import chunker_sequence
-from cognite_toolkit._cdf_tk.utils.http_client import HTTPClient, HTTPMessage, SimpleBodyRequest
 from cognite_toolkit._cdf_tk.utils.useful_types import JsonVal
 
 from ._base import Page, UploadableStorageIO, UploadItem
@@ -185,23 +194,52 @@ class CanvasIO(UploadableStorageIO[CanvasSelector, IndustrialCanvas, IndustrialC
         results: list[HTTPMessage] = []
         for item in data_chunk:
             instances = item.item.as_instances()
-            items: list[dict[str, JsonVal]] = []
+            upsert_items: list[dict[str, JsonValue]] = []
             for instance in instances:
                 dumped = instance.dump()
                 if self.exclude_existing_version:
                     dumped.pop("existingVersion", None)
-                items.append(dumped)
-
-            responses = http_client.request_with_retries(
-                message=SimpleBodyRequest(
-                    endpoint_url=config.create_api_url("/models/instances"),
-                    method="POST",
-                    # MyPy does not understand that .dump is valid json
-                    body_content={"items": items},  # type: ignore[dict-item]
-                )
-            )
-            results.extend(responses.as_item_responses(item.source_id))
+                upsert_items.append(dumped)
+
+            canvas = item.item.canvas
+            existing = self.client.canvas.industrial.retrieve(canvas.external_id)
+            if existing is not None:
+                existing_instance_ids = existing.as_write().as_instance_ids(include_solution_tags=False)
+                delete_set = set(existing_instance_ids) - set(item.item.as_instance_ids())
+                to_delete: list[Any] = [
+                    {
+                        "space": instance_id.space,
+                        "externalId": instance_id.external_id,
+                        "instanceType": instance_id.instance_type,
+                    }
+                    for instance_id in delete_set
+                ]
+            else:
+                to_delete = []
 
+            last_response: HTTPResult2 | None = None
+            for upsert_chunk, delete_chunk in zip_longest(
+                chunker_sequence(upsert_items, 1000), chunker_sequence(to_delete, 1000), fillvalue=None
+            ):
+                body_content: dict[str, JsonValue] = {}
+                if upsert_chunk:
+                    # MyPy fails do understand that list[dict[str, JsonValue]] is a subtype of JsonValue
+                    body_content["items"] = upsert_chunk  # type: ignore[assignment]
+                if delete_chunk:
+                    body_content["delete"] = delete_chunk
+
+                response = http_client.request_single_retries(
+                    message=RequestMessage2(
+                        endpoint_url=config.create_api_url("/models/instances"),
+                        method="POST",
+                        body_content=body_content,
+                    )
+                )
+                if not isinstance(response, SuccessResponse2):
+                    results.append(response.as_item_response(item.source_id))
+                last_response = response
+            if last_response is not None:
+                results.append(last_response.as_item_response(item.source_id))
         return results
 
     def data_to_json_chunk(