Add UnitesStrategyEnum and integrate into NodeTemplate validation

state-manager/app/models/node_template_model.py

@@ -1,10 +1,17 @@
 from pydantic import Field, BaseModel, field_validator
 from typing import Any, Optional, List
 from .dependent_string import DependentString
+from enum import Enum
+
+
+class UnitesStrategyEnum(str, Enum):
+    ALL_SUCCESS = "ALL_SUCCESS"
+    ALL_DONE = "ALL_DONE"
 
 
 class Unites(BaseModel):
     identifier: str = Field(..., description="Identifier of the node")
+    strategy: UnitesStrategyEnum = Field(..., description="Strategy of the unites")
 
 
 class NodeTemplate(BaseModel):

state-manager/app/tasks/create_next_states.py

@@ -1,13 +1,14 @@
 from beanie import PydanticObjectId
 from pymongo.errors import DuplicateKeyError, BulkWriteError
-from beanie.operators import In, NE
+from beanie.operators import In, NE, NotIn
 from app.singletons.logs_manager import LogsManager
 from app.models.db.graph_template_model import GraphTemplate
 from app.models.db.state import State
 from app.models.state_status_enum import StateStatusEnum
 from app.models.node_template_model import NodeTemplate
 from app.models.db.registered_node import RegisteredNode
 from app.models.dependent_string import DependentString
+from app.models.node_template_model import UnitesStrategyEnum
 from json_schema_to_pydantic import create_model
 from pydantic import BaseModel
 from typing import Type
@@ -30,7 +31,7 @@ async def check_unites_satisfied(namespace: str, graph_name: str, node_template:
     if not unites_id:
         raise ValueError(f"Unit identifier not found in parents: {node_template.unites.identifier}")
     else:
-        if await State.find(
+        if node_template.unites.strategy == UnitesStrategyEnum.ALL_SUCCESS and await State.find(
                 State.namespace_name == namespace,
                 State.graph_name == graph_name,
                 NE(State.status, StateStatusEnum.SUCCESS),
@@ -39,6 +40,18 @@ async def check_unites_satisfied(namespace: str, graph_name: str, node_template:
                 }
             ).count() > 0:
                 return False
+
+        if node_template.unites.strategy == UnitesStrategyEnum.ALL_DONE and await State.find(
+                State.namespace_name == namespace,
+                State.graph_name == graph_name,
+                In(State.status, [StateStatusEnum.CREATED, StateStatusEnum.QUEUED, StateStatusEnum.EXECUTED]),
+                {
+                    f"parents.{node_template.unites.identifier}": unites_id
+                }
+            ).count() > 0:
+                return False
+
+
     return True