feat(core): change error to branchType on Topology Graph

kestra-io · Jan 13, 2025 · b58107e · b58107e
1 parent 089e4d4
commit b58107e
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 12 deletions.
diff --git a/core/src/main/java/io/kestra/core/models/hierarchies/AbstractGraph.java b/core/src/main/java/io/kestra/core/models/hierarchies/AbstractGraph.java
@@ -4,7 +4,6 @@
 import com.fasterxml.jackson.annotation.JsonInclude;
 import com.fasterxml.jackson.annotation.JsonTypeInfo;
 import io.micronaut.core.annotation.Introspected;
-import lombok.Builder;
 import lombok.Getter;
 import lombok.Setter;
 import lombok.ToString;
@@ -19,7 +18,7 @@ public abstract class AbstractGraph {
     @JsonInclude
     protected String type;
     @Setter
-    protected boolean error;
+    protected BranchType branchType;
 
     public AbstractGraph() {
         this.type = this.getClass().getName();
@@ -39,8 +38,8 @@ public void updateUidWithChildren(String uid) {
         this.uid = uid;
     }
 
-    public void updateErrorWithChildren(boolean error) {
-        this.error = error;
+    public void updateWithChildren(BranchType branchType) {
+        this.branchType = branchType;
     }
 
     public AbstractGraph forExecution() {
@@ -53,4 +52,9 @@ public boolean equals(Object o) {
         if (!(o instanceof AbstractGraph)) return false;
         return o.hashCode() == this.hashCode();
     }
+
+    public enum BranchType {
+        ERROR,
+        FINALLY
+    }
 }
diff --git a/core/src/main/java/io/kestra/core/models/hierarchies/GraphCluster.java b/core/src/main/java/io/kestra/core/models/hierarchies/GraphCluster.java
@@ -139,12 +139,12 @@ public void updateUidWithChildren(String uid) {
     }
 
     @Override
-    public void updateErrorWithChildren(boolean error) {
-        this.error = error;
+    public void updateWithChildren(BranchType branchType) {
+        this.branchType = branchType;
 
-        this.taskNode.error = error;
-        this.root.error = error;
-        this.end.error = error;
+        this.taskNode.branchType = branchType;
+        this.root.branchType = branchType;
+        this.end.branchType = branchType;
     }
 
     @Override

diff --git a/core/src/main/java/io/kestra/core/services/GraphService.java b/core/src/main/java/io/kestra/core/services/GraphService.java
@@ -153,8 +153,8 @@ public void replace() {
                 });
             parentCluster.getGraph().removeNode(taskToReplace);
 
-            if (taskToReplace.isError()) {
-                clusterForReplacement.updateErrorWithChildren(true);
+            if (taskToReplace.getBranchType() != null) {
+                clusterForReplacement.updateWithChildren(taskToReplace.getBranchType());
             }
         }
     }

diff --git a/core/src/main/java/io/kestra/core/utils/GraphUtils.java b/core/src/main/java/io/kestra/core/utils/GraphUtils.java
@@ -337,8 +337,11 @@ private static void fillGraph(
                 // add the node
                 graph.addNode(currentGraph);
 
+                if (relationType == RelationType.ERROR || relationType == RelationType.FINALLY) {
+                    currentGraph.updateWithChildren(AbstractGraph.BranchType.valueOf(relationType.name()));
+                }
+
                 if (relationType == RelationType.ERROR) {
-                    currentGraph.updateErrorWithChildren(true);
                     if (isFirst) {
                         previous = graph.getRoot();
                     }