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

cognite_toolkit/_cdf_tk/{loaders/_resource_loaders/workflow_loaders.py → cruds/_resource_cruds/workflow.py}

@@ -11,46 +11,47 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from __future__ import annotations
+
 
 import json
 from collections.abc import Hashable, Iterable, Sequence
-from functools import lru_cache
+from graphlib import CycleError, TopologicalSorter
 from pathlib import Path
 from typing import Any, final
 
 from cognite.client.data_classes import (
     ClientCredentials,
     Workflow,
-    WorkflowList,
     WorkflowTrigger,
-    WorkflowTriggerList,
     WorkflowTriggerUpsert,
-    WorkflowTriggerUpsertList,
     WorkflowUpsert,
-    WorkflowUpsertList,
     WorkflowVersion,
     WorkflowVersionId,
-    WorkflowVersionList,
     WorkflowVersionUpsert,
-    WorkflowVersionUpsertList,
 )
 from cognite.client.data_classes.capabilities import (
     Capability,
     WorkflowOrchestrationAcl,
 )
+from cognite.client.data_classes.workflows import (
+    SubworkflowReferenceParameters,
+    WorkflowList,
+    WorkflowTriggerList,
+    WorkflowVersionList,
+)
 from cognite.client.exceptions import CogniteAuthError, CogniteNotFoundError
 from cognite.client.utils.useful_types import SequenceNotStr
 from rich import print
 from rich.console import Console
 
-from cognite_toolkit._cdf_tk._parameters import ANY_INT, ANY_STR, ANYTHING, ParameterSpec, ParameterSpecSet
 from cognite_toolkit._cdf_tk.client import ToolkitClient
+from cognite_toolkit._cdf_tk.cruds._base_cruds import ResourceCRUD
 from cognite_toolkit._cdf_tk.exceptions import (
     ResourceCreationError,
+    ToolkitCycleError,
     ToolkitRequiredValueError,
 )
-from cognite_toolkit._cdf_tk.loaders._base_loaders import ResourceLoader
+from cognite_toolkit._cdf_tk.resource_classes import WorkflowTriggerYAML, WorkflowVersionYAML, WorkflowYAML
 from cognite_toolkit._cdf_tk.tk_warnings import (
     LowSeverityWarning,
     MissingReferencedWarning,
@@ -60,36 +61,33 @@ from cognite_toolkit._cdf_tk.utils import (
     calculate_secure_hash,
     humanize_collection,
     load_yaml_inject_variables,
-    to_directory_compatible,
+    sanitize_filename,
 )
 from cognite_toolkit._cdf_tk.utils.cdf import read_auth, try_find_error
-from cognite_toolkit._cdf_tk.utils.collection import chunker
 from cognite_toolkit._cdf_tk.utils.diff_list import diff_list_hashable, diff_list_identifiable
 
-from .auth_loaders import GroupAllScopedLoader
-from .data_organization_loaders import DataSetsLoader
-from .function_loaders import FunctionLoader
-from .group_scoped_loader import GroupResourceScopedLoader
-from .transformation_loaders import TransformationLoader
+from .auth import GroupAllScopedCRUD
+from .data_organization import DataSetsCRUD
+from .function import FunctionCRUD
+from .group_scoped import GroupResourceScopedCRUD
+from .transformation import TransformationCRUD
 
 
 @final
-class WorkflowLoader(ResourceLoader[str, WorkflowUpsert, Workflow, WorkflowUpsertList, WorkflowList]):
+class WorkflowCRUD(ResourceCRUD[str, WorkflowUpsert, Workflow]):
     folder_name = "workflows"
-    filename_pattern = r"^.*Workflow$"
     resource_cls = Workflow
     resource_write_cls = WorkflowUpsert
-    list_cls = WorkflowList
-    list_write_cls = WorkflowUpsertList
     kind = "Workflow"
     dependencies = frozenset(
         {
-            GroupAllScopedLoader,
-            TransformationLoader,
-            FunctionLoader,
-            DataSetsLoader,
+            GroupAllScopedCRUD,
+            TransformationCRUD,
+            FunctionCRUD,
+            DataSetsCRUD,
         }
     )
+    yaml_cls = WorkflowYAML
     _doc_base_url = "https://api-docs.cognite.com/20230101-beta/tag/"
     _doc_url = "Workflows/operation/CreateOrUpdateWorkflow"
 
@@ -138,23 +136,14 @@ class WorkflowLoader(ResourceLoader[str, WorkflowUpsert, Workflow, WorkflowUpser
             dumped["dataSetExternalId"] = self.client.lookup.data_sets.external_id(data_set_id)
         return dumped
 
-    def retrieve(self, ids: SequenceNotStr[str]) -> WorkflowList:
-        workflows = []
-        for ext_id in ids:
-            workflow = self.client.workflows.retrieve(external_id=ext_id)
-            if workflow:
-                workflows.append(workflow)
-        return WorkflowList(workflows)
-
-    def _upsert(self, items: WorkflowUpsert | WorkflowUpsertList) -> WorkflowList:
-        upserts = [items] if isinstance(items, WorkflowUpsert) else items
-        return WorkflowList([self.client.workflows.upsert(upsert) for upsert in upserts])
+    def retrieve(self, ids: SequenceNotStr[str]) -> Sequence[Workflow]:
+        return self.client.workflows.retrieve(external_id=ids, ignore_unknown_ids=True)
 
-    def create(self, items: WorkflowUpsert | WorkflowUpsertList) -> WorkflowList:
-        return self._upsert(items)
+    def create(self, items: Sequence[WorkflowUpsert]) -> WorkflowList:
+        return self.client.workflows.upsert(items)
 
-    def update(self, items: WorkflowUpsertList) -> WorkflowList:
-        return self._upsert(items)
+    def update(self, items: Sequence[WorkflowUpsert]) -> WorkflowList:
+        return self.client.workflows.upsert(items)
 
     def delete(self, ids: SequenceNotStr[str]) -> int:
         successes = 0
@@ -184,53 +173,25 @@ class WorkflowLoader(ResourceLoader[str, WorkflowUpsert, Workflow, WorkflowUpser
                 yield workflow
 
     @classmethod
-    @lru_cache(maxsize=1)
-    def get_write_cls_parameter_spec(cls) -> ParameterSpecSet:
-        spec = super().get_write_cls_parameter_spec()
-        spec.add(
-            ParameterSpec(
-                ("dataSetExternalId",),
-                frozenset({"str"}),
-                is_required=False,
-                _is_nullable=True,
-            )
-        )
-        spec.discard(
-            ParameterSpec(
-                ("dataSetId",),
-                frozenset({"str"}),
-                is_required=False,
-                _is_nullable=True,
-            )
-        )
-        return spec
-
-    @classmethod
-    def get_dependent_items(cls, item: dict) -> Iterable[tuple[type[ResourceLoader], Hashable]]:
+    def get_dependent_items(cls, item: dict) -> Iterable[tuple[type[ResourceCRUD], Hashable]]:
         """Returns all items that this item requires.
 
         For example, a TimeSeries requires a DataSet, so this method would return the
         DatasetLoader and identifier of that dataset.
         """
         if "dataSetExternalId" in item:
-            yield DataSetsLoader, item["dataSetExternalId"]
+            yield DataSetsCRUD, item["dataSetExternalId"]
 
 
 @final
-class WorkflowVersionLoader(
-    ResourceLoader[
-        WorkflowVersionId, WorkflowVersionUpsert, WorkflowVersion, WorkflowVersionUpsertList, WorkflowVersionList
-    ]
-):
+class WorkflowVersionCRUD(ResourceCRUD[WorkflowVersionId, WorkflowVersionUpsert, WorkflowVersion]):
     folder_name = "workflows"
-    filename_pattern = r"^.*WorkflowVersion$"
     resource_cls = WorkflowVersion
     resource_write_cls = WorkflowVersionUpsert
-    list_cls = WorkflowVersionList
-    list_write_cls = WorkflowVersionUpsertList
     kind = "WorkflowVersion"
-    dependencies = frozenset({WorkflowLoader})
-    parent_resource = frozenset({WorkflowLoader})
+    dependencies = frozenset({WorkflowCRUD})
+    parent_resource = frozenset({WorkflowCRUD})
+    yaml_cls = WorkflowVersionYAML
 
     _doc_base_url = "https://api-docs.cognite.com/20230101-beta/tag/"
     _doc_url = "Workflow-versions/operation/CreateOrUpdateWorkflowVersion"
@@ -341,9 +302,15 @@ class WorkflowVersionLoader(
             if task_id not in local_task_by_id:
                 continue
             local_task = local_task_by_id[task_id]
-            for default_key, default_value in [("retries", 3), ("timeout", 3600), ("onFailure", "abortWorkflow")]:
-                if default_key not in local_task and cdf_task.get(default_key) == default_value:
-                    del cdf_task[default_key]
+            for key, default_value in [("retries", 3), ("timeout", 3600), ("onFailure", "abortWorkflow")]:
+                if key not in local_task and cdf_task.get(key) == default_value:
+                    del cdf_task[key]
+                elif (
+                    key in local_task
+                    and local_task[key] is None
+                    and (cdf_task.get(key) == default_value or key not in cdf_task)
+                ):
+                    cdf_task[key] = None
 
             if local_task["type"] == "function" and cdf_task["type"] == "function":
                 cdf_parameters = cdf_task["parameters"]
@@ -353,6 +320,19 @@ class WorkflowVersionLoader(
                     local_function = local_parameters["function"]
                     if local_function.get("data") == {} and "data" not in cdf_function:
                         cdf_parameters["function"] = local_function
+            elif local_task["type"] == "transformation" and cdf_task["type"] == "transformation":
+                cdf_parameters = cdf_task["parameters"]
+                local_parameters = local_task["parameters"]
+                if "transformation" in cdf_parameters and "transformation" in local_parameters:
+                    for key, default_transformation_value in [
+                        ("concurrencyPolicy", "fail"),
+                        ("useTransformationCredentials", False),
+                    ]:
+                        if (
+                            key not in local_parameters["transformation"]
+                            and cdf_parameters["transformation"].get(key) == default_transformation_value
+                        ):
+                            del cdf_parameters["transformation"][key]
         return dumped
 
     def diff_list(
@@ -368,9 +348,9 @@ class WorkflowVersionLoader(
         return super().diff_list(local, cdf, json_path)
 
     @classmethod
-    def get_dependent_items(cls, item: dict) -> Iterable[tuple[type[ResourceLoader], Hashable]]:
+    def get_dependent_items(cls, item: dict) -> Iterable[tuple[type[ResourceCRUD], Hashable]]:
         if "workflowExternalId" in item:
-            yield WorkflowLoader, item["workflowExternalId"]
+            yield WorkflowCRUD, item["workflowExternalId"]
 
     @classmethod
     def check_item(cls, item: dict, filepath: Path, element_no: int | None) -> list[ToolkitWarning]:
@@ -399,28 +379,21 @@ class WorkflowVersionLoader(
         return warnings
 
     def retrieve(self, ids: SequenceNotStr[WorkflowVersionId]) -> WorkflowVersionList:
-        retrieved = WorkflowVersionList([])
         if not ids:
-            return retrieved
-        for chunk in chunker(ids, self._list_workflow_filter_limit):
-            retrieved_chunk = self.client.workflows.versions.list(chunk)
-            retrieved.extend(retrieved_chunk)
-        return retrieved
+            return WorkflowVersionList([])
+        return self.client.workflows.versions.retrieve(workflow_external_id=list(ids), ignore_unknown_ids=True)
 
-    def _upsert(self, items: WorkflowVersionUpsertList) -> WorkflowVersionList:
-        return WorkflowVersionList([self.client.workflows.versions.upsert(item) for item in items])
+    def _upsert(self, items: Sequence[WorkflowVersionUpsert]) -> WorkflowVersionList:
+        return self.client.workflows.versions.upsert(items)
 
-    def create(self, items: WorkflowVersionUpsertList) -> WorkflowVersionList:
-        upserted = []
-        for item in items:
-            upserted.append(self.client.workflows.versions.upsert(item))
+    def create(self, items: Sequence[WorkflowVersionUpsert]) -> WorkflowVersionList:
+        upserted: list[WorkflowVersion] = []
+        for item in self.topological_sort(items):
+            upserted.extend(self.client.workflows.versions.upsert([item]))
         return WorkflowVersionList(upserted)
 
-    def update(self, items: WorkflowVersionUpsertList) -> WorkflowVersionList:
-        updated = []
-        for item in items:
-            updated.append(self.client.workflows.versions.upsert(item))
-        return WorkflowVersionList(updated)
+    def update(self, items: Sequence[WorkflowVersionUpsert]) -> WorkflowVersionList:
+        return self._upsert(items)
 
     def delete(self, ids: SequenceNotStr[WorkflowVersionId]) -> int:
         successes = 0
@@ -442,48 +415,6 @@ class WorkflowVersionLoader(
         workflow_ids = [parent_id for parent_id in parent_ids if isinstance(parent_id, str)] if parent_ids else None
         return self.client.workflows.versions.list(limit=-1, workflow_version_ids=workflow_ids)  # type: ignore[arg-type]
 
-    @classmethod
-    @lru_cache(maxsize=1)
-    def get_write_cls_parameter_spec(cls) -> ParameterSpecSet:
-        spec = super().get_write_cls_parameter_spec()
-        # The Parameter class in the SDK class WorkflowVersion implementation is deviating from the API.
-        # So we need to modify the spec to match the API.
-        parameter_path = ("workflowDefinition", "tasks", ANY_INT, "parameters")
-        length = len(parameter_path)
-        for item in spec:
-            if len(item.path) >= length + 1 and item.path[:length] == parameter_path[:length]:
-                # Add extra ANY_STR layer
-                # The spec class is immutable, so we use this trick to modify it.
-                object.__setattr__(item, "path", item.path[:length] + (ANY_STR,) + item.path[length:])
-        spec.add(ParameterSpec((*parameter_path, ANY_STR), frozenset({"dict"}), is_required=True, _is_nullable=False))
-        # The depends on is implemented as a list of string in the SDK, but in the API spec it
-        # is a list of objects with one 'externalId' field.
-        spec.add(
-            ParameterSpec(
-                ("workflowDefinition", "tasks", ANY_INT, "dependsOn", ANY_INT, "externalId"),
-                frozenset({"str"}),
-                is_required=False,
-                _is_nullable=False,
-            )
-        )
-        spec.add(
-            ParameterSpec(
-                ("workflowDefinition", "tasks", ANY_INT, "type"),
-                frozenset({"str"}),
-                is_required=True,
-                _is_nullable=False,
-            )
-        )
-        spec.add(
-            ParameterSpec(
-                ("workflowDefinition", "tasks", ANY_INT, "parameters", "subworkflow", ANYTHING),
-                frozenset({"dict"}),
-                is_required=False,
-                _is_nullable=False,
-            )
-        )
-        return spec
-
     @classmethod
     def as_str(cls, id: WorkflowVersionId) -> str:
         if id.version is None:
@@ -493,22 +424,42 @@ class WorkflowVersionLoader(
         else:
             version = f"_v{id.version}"
 
-        return to_directory_compatible(f"{id.workflow_external_id}{version}")
+        return sanitize_filename(f"{id.workflow_external_id}{version}")
+
+    @classmethod
+    def topological_sort(cls, items: Sequence[WorkflowVersionUpsert]) -> list[WorkflowVersionUpsert]:
+        workflow_by_id: dict[WorkflowVersionId, WorkflowVersionUpsert] = {item.as_id(): item for item in items}
+        dependencies: dict[WorkflowVersionId, set[WorkflowVersionId]] = {}
+        for item_id, item in workflow_by_id.items():
+            dependencies[item_id] = set()
+            for task in item.workflow_definition.tasks:
+                if isinstance(task.parameters, SubworkflowReferenceParameters):
+                    dependencies[item_id].add(
+                        WorkflowVersionId(task.parameters.workflow_external_id, task.parameters.version)
+                    )
+
+        try:
+            return [
+                workflow_by_id[item_id]
+                for item_id in TopologicalSorter(dependencies).static_order()
+                if item_id in workflow_by_id
+            ]
+        except CycleError as e:
+            raise ToolkitCycleError(
+                f"Cannot deploy workflows. Cycle detected {e.args} in the 'subworkflow' dependencies of the workflows.",
+                *e.args[1:],
+            ) from None
 
 
 @final
-class WorkflowTriggerLoader(
-    ResourceLoader[str, WorkflowTriggerUpsert, WorkflowTrigger, WorkflowTriggerUpsertList, WorkflowTriggerList]
-):
+class WorkflowTriggerCRUD(ResourceCRUD[str, WorkflowTriggerUpsert, WorkflowTrigger]):
     folder_name = "workflows"
-    filename_pattern = r"^.*WorkflowTrigger$"
     resource_cls = WorkflowTrigger
     resource_write_cls = WorkflowTriggerUpsert
-    list_cls = WorkflowTriggerList
-    list_write_cls = WorkflowTriggerUpsertList
     kind = "WorkflowTrigger"
-    dependencies = frozenset({WorkflowLoader, WorkflowVersionLoader, GroupResourceScopedLoader, GroupAllScopedLoader})
-    parent_resource = frozenset({WorkflowLoader})
+    dependencies = frozenset({WorkflowCRUD, WorkflowVersionCRUD, GroupResourceScopedCRUD, GroupAllScopedCRUD})
+    parent_resource = frozenset({WorkflowCRUD})
+    yaml_cls = WorkflowTriggerYAML
 
     _doc_url = "Workflow-triggers/operation/CreateOrUpdateTriggers"
 
@@ -551,13 +502,16 @@ class WorkflowTriggerLoader(
             WorkflowOrchestrationAcl.Scope.All(),
         )
 
-    def create(self, items: WorkflowTriggerUpsertList) -> WorkflowTriggerList:
+    def create(self, items: Sequence[WorkflowTriggerUpsert]) -> WorkflowTriggerList:
+        return self._upsert(items)
+
+    def _upsert(self, items: Sequence[WorkflowTriggerUpsert]) -> WorkflowTriggerList:
         created = WorkflowTriggerList([])
         for item in items:
-            created.append(self._create_item(item))
+            created.append(self._upsert_item(item))
         return created
 
-    def _create_item(self, item: WorkflowTriggerUpsert) -> WorkflowTrigger:
+    def _upsert_item(self, item: WorkflowTriggerUpsert) -> WorkflowTrigger:
         credentials = self._authentication_by_id.get(item.external_id)
         try:
             return self.client.workflows.triggers.upsert(item, credentials)
@@ -571,17 +525,8 @@ class WorkflowTriggerLoader(
         lookup = set(ids)
         return WorkflowTriggerList([trigger for trigger in all_triggers if trigger.external_id in lookup])
 
-    def update(self, items: WorkflowTriggerUpsertList) -> WorkflowTriggerList:
-        exising = self.client.workflows.triggers.list(limit=-1)
-        existing_lookup = {trigger.external_id: trigger for trigger in exising}
-        updated = WorkflowTriggerList([])
-        for item in items:
-            if item.external_id in existing_lookup:
-                self.client.workflows.triggers.delete(external_id=item.external_id)
-
-            created = self._create_item(item)
-            updated.append(created)
-        return updated
+    def update(self, items: Sequence[WorkflowTriggerUpsert]) -> WorkflowTriggerList:
+        return self._upsert(items)
 
     def delete(self, ids: SequenceNotStr[str]) -> int:
         successes = 0
@@ -608,52 +553,13 @@ class WorkflowTriggerLoader(
         return triggers
 
     @classmethod
-    @lru_cache(maxsize=1)
-    def get_write_cls_parameter_spec(cls) -> ParameterSpecSet:
-        spec = super().get_write_cls_parameter_spec()
-        # Removed by the SDK
-        spec.add(
-            ParameterSpec(
-                ("triggerRule", "triggerType"),
-                frozenset({"str"}),
-                is_required=True,
-                _is_nullable=False,
-            )
-        )
-        spec.add(
-            ParameterSpec(
-                ("authentication",),
-                frozenset({"dict"}),
-                is_required=True,
-                _is_nullable=False,
-            )
-        )
-        spec.add(
-            ParameterSpec(
-                ("authentication", "clientId"),
-                frozenset({"str"}),
-                is_required=False,
-                _is_nullable=False,
-            )
-        )
-        spec.add(
-            ParameterSpec(
-                ("authentication", "clientSecret"),
-                frozenset({"str"}),
-                is_required=False,
-                _is_nullable=False,
-            )
-        )
-        return spec
-
-    @classmethod
-    def get_dependent_items(cls, item: dict) -> Iterable[tuple[type[ResourceLoader], Hashable]]:
+    def get_dependent_items(cls, item: dict) -> Iterable[tuple[type[ResourceCRUD], Hashable]]:
         """Returns all items that this item requires."""
         if "workflowExternalId" in item:
-            yield WorkflowLoader, item["workflowExternalId"]
+            yield WorkflowCRUD, item["workflowExternalId"]
 
             if "workflowVersion" in item:
-                yield WorkflowVersionLoader, WorkflowVersionId(item["workflowExternalId"], item["workflowVersion"])
+                yield WorkflowVersionCRUD, WorkflowVersionId(item["workflowExternalId"], item["workflowVersion"])
 
     def load_resource_file(
         self, filepath: Path, environment_variables: dict[str, str | None] | None = None
@@ -686,6 +592,13 @@ class WorkflowTriggerLoader(
         if isinstance(dumped.get("data"), str) and isinstance(local.get("data"), dict):
             dumped["data"] = json.loads(dumped["data"])
 
+        cdf_rule = dumped.get("triggerRule", {})
+        local_rule = local.get("triggerRule", {})
+        if cdf_rule.get("triggerType") == "schedule" and local_rule.get("triggerType") == "schedule":
+            if cdf_rule.get("timezone") == "UTC" and "timezone" not in local_rule:
+                # The server defaults to UTC if not set, so we remove it if the local does not have it set.
+                del cdf_rule["timezone"]
+
         if "authentication" in local:
             # Changes in auth will be detected by the hash. We need to do this to ensure
             # that the pull command works.