Enhance validation in NodeTemplate and GraphTemplate models

- Introduced field validators for various attributes in NodeTemplate to ensure non-empty values and uniqueness for node names, identifiers, next nodes, and unites.
- Added validation for name and namespace in GraphTemplate to prevent empty values.
- Implemented model-level validation to ensure node identifiers are unique and that next node identifiers exist within the graph.
- Removed outdated verification functions from verify_graph.py to streamline the validation process.

These changes improve data integrity and validation consistency across the models.

Author: NiveditJain

state-manager/app/models/db/graph_template_model.py

@@ -3,8 +3,8 @@
 import asyncio
 
 from .base import BaseDatabaseModel
-from pydantic import Field, field_validator, PrivateAttr
-from typing import Optional, List
+from pydantic import Field, field_validator, PrivateAttr, model_validator
+from typing import Optional, List, Self
 from ..graph_template_validation_status import GraphTemplateValidationStatus
 from ..node_template_model import NodeTemplate
 from pymongo import IndexModel
@@ -41,6 +41,20 @@ def get_node_by_identifier(self, identifier: str) -> NodeTemplate | None:
         assert self._node_by_identifier is not None
         return self._node_by_identifier.get(identifier)
 
+    @field_validator('name')
+    @classmethod
+    def validate_name(cls, v: str) -> str:
+        if v == "" or v is None:
+            raise ValueError("Name cannot be empty")
+        return v
+    
+    @field_validator('namespace')
+    @classmethod
+    def validate_namespace(cls, v: str) -> str:
+        if v == "" or v is None:
+            raise ValueError("Namespace cannot be empty")
+        return v
+
     @field_validator('secrets')
     @classmethod
     def validate_secrets(cls, v: Dict[str, str]) -> Dict[str, str]:
@@ -55,6 +69,43 @@ def validate_secrets(cls, v: Dict[str, str]) -> Dict[str, str]:
 
         return v
 
+    @model_validator(mode='after')
+    def validate_nodes(self) -> Self:
+        for node in self.nodes:
+            if node.namespace != self.namespace:
+                raise ValueError(f"Node namespace {node.namespace} does not match graph namespace {self.namespace}")
+        return self
+    
+    @field_validator('nodes')
+    @classmethod
+    def validate_unique_identifiers(cls, v: List[NodeTemplate]) -> List[NodeTemplate]:
+        identifiers = set()
+        errors = []
+        for node in v:
+            if node.identifier in identifiers:
+                errors.append(f"Node identifier {node.identifier} is not unique")
+            identifiers.add(node.identifier)
+        if errors:
+            raise ValueError("\n".join(errors))
+        return v
+    
+    @field_validator('nodes')
+    @classmethod
+    def validate_next_nodes_identifiers_exist(cls, v: List[NodeTemplate]) -> List[NodeTemplate]:
+        identifiers = set()
+        for node in v:
+            identifiers.add(node.identifier)
+
+        errors = []    
+        for node in v:
+            if node.next_nodes:
+                for next_node in node.next_nodes:
+                    if next_node not in identifiers:
+                        errors.append(f"Node identifier {next_node} does not exist in the graph")
+        if errors:
+            raise ValueError("\n".join(errors))
+        return v
+    
     @classmethod
     def _validate_secret_value(cls, secret_value: str) -> None:
         # Check minimum length for AES-GCM encrypted string
@@ -70,7 +121,6 @@ def _validate_secret_value(cls, secret_value: str) -> None:
         except Exception:
             raise ValueError("Value is not valid URL-safe base64 encoded")
 
-
     def set_secrets(self, secrets: Dict[str, str]) -> "GraphTemplate":
         self.secrets = {secret_name: get_encrypter().encrypt(secret_value) for secret_name, secret_value in secrets.items()}
         return self

state-manager/app/models/node_template_model.py

@@ -1,4 +1,4 @@
-from pydantic import Field, BaseModel
+from pydantic import Field, BaseModel, field_validator
 from typing import Any, Optional, List
 
 
@@ -12,4 +12,49 @@ class NodeTemplate(BaseModel):
     identifier: str = Field(..., description="Identifier of the node")
     inputs: dict[str, Any] = Field(..., description="Inputs of the node")
     next_nodes: Optional[List[str]] = Field(None, description="Next nodes to execute")
-    unites: Optional[Unites] = Field(None, description="Unites of the node")
+    unites: Optional[Unites] = Field(None, description="Unites of the node")
+
+    @field_validator('node_name')
+    @classmethod
+    def validate_node_name(cls, v: str) -> str:
+        if v == "" or v is None:
+            raise ValueError("Node name cannot be empty")
+        return v
+    
+    @field_validator('node_name')
+    @classmethod
+    def validate_node_name_unique(cls, v: str) -> str:
+        if v == "" or v is None:
+            raise ValueError("Node name cannot be empty")
+        return v
+    
+    @field_validator('identifier')
+    @classmethod
+    def validate_identifier_unique(cls, v: str) -> str:
+        if v == "" or v is None:
+            raise ValueError("Node identifier cannot be empty")
+        return v
+    
+    @field_validator('next_nodes')
+    @classmethod
+    def validate_next_nodes(cls, v: Optional[List[str]]) -> Optional[List[str]]:
+        identifiers = set()
+        errors = []
+        if v is not None:
+            for next_node_identifier in v:
+                if next_node_identifier == "" or next_node_identifier is None:
+                    errors.append("Next node identifier cannot be empty")
+                elif next_node_identifier in identifiers:
+                    errors.append(f"Next node identifier {next_node_identifier} is not unique")
+                identifiers.add(next_node_identifier)
+        if errors:
+            raise ValueError("\n".join(errors))
+        return v
+    
+    @field_validator('unites')
+    @classmethod
+    def validate_unites(cls, v: Optional[Unites]) -> Optional[Unites]:
+        if v is not None:
+            if v.identifier == "" or v.identifier is None:
+                raise ValueError("Unites identifier cannot be empty")
+        return v

state-manager/app/tasks/verify_graph.py

@@ -9,20 +9,6 @@
 
 logger = LogsManager().get_logger()
 
-async def verify_nodes_names(nodes: list[NodeTemplate]) -> list[str]:
-    errors = []
-    for node in nodes:
-        if node.node_name is None or node.node_name == "":
-            errors.append(f"Node {node.identifier} has no name")    
-    return errors
-
-async def verify_nodes_namespace(nodes: list[NodeTemplate], graph_namespace: str) -> list[str]:
-    errors = []
-    for node in nodes:
-        if node.namespace != graph_namespace and node.namespace != "exospherehost":
-            errors.append(f"Node {node.identifier} has invalid namespace '{node.namespace}'. Must match graph namespace '{graph_namespace}' or use universal namespace 'exospherehost'")
-    return errors
-
 async def verify_node_exists(nodes: list[NodeTemplate], database_nodes: list[RegisteredNode]) -> list[str]:
     errors = []
     template_nodes_set = set([(node.node_name, node.namespace) for node in nodes])
@@ -33,8 +19,6 @@ async def verify_node_exists(nodes: list[NodeTemplate], database_nodes: list[Reg
     for node in nodes_not_found:
         errors.append(f"Node {node[0]} in namespace {node[1]} does not exist.")
     return errors
-
-async def verify_node_identifiers(nodes: list[NodeTemplate]) -> list[str]:
     errors = []
     identifier_to_nodes = {}
 
@@ -252,10 +236,7 @@ async def verify_graph(graph_template: GraphTemplate):
         database_nodes = await get_database_nodes(graph_template.nodes, graph_template.namespace)
 
         basic_verify_tasks = [
-            verify_nodes_names(graph_template.nodes),
-            verify_nodes_namespace(graph_template.nodes, graph_template.namespace),
             verify_node_exists(graph_template.nodes, database_nodes),
-            verify_node_identifiers(graph_template.nodes),
             verify_secrets(graph_template, database_nodes)
         ]
         errors.extend(await asyncio.gather(*basic_verify_tasks))