From 2beaff72356af28388c898369bdcac30120193c0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Andr=C3=A9s=20G=C3=B3mez?= <andresgomezfrr@gmail.com>
Date: Mon, 1 Jul 2024 14:40:38 +0200
Subject: [PATCH] Handle conditional in dynamic (#301)

Signed-off-by: Andres Gomez Ferrer <andresgomezfrr@users.noreply.github.com>
Co-authored-by: Andres Gomez Ferrer <andresgomezfrr@users.noreply.github.com>
---
 .../flyte/jflyte/ExecuteDynamicWorkflow.java  | 24 +++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/jflyte/src/main/java/org/flyte/jflyte/ExecuteDynamicWorkflow.java b/jflyte/src/main/java/org/flyte/jflyte/ExecuteDynamicWorkflow.java
index b46c4489..9c4c8729 100644
--- a/jflyte/src/main/java/org/flyte/jflyte/ExecuteDynamicWorkflow.java
+++ b/jflyte/src/main/java/org/flyte/jflyte/ExecuteDynamicWorkflow.java
@@ -24,6 +24,7 @@
 
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Maps;
+import java.util.ArrayList;
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.List;
@@ -41,6 +42,7 @@
 import org.flyte.api.v1.DynamicJobSpec;
 import org.flyte.api.v1.DynamicWorkflowTask;
 import org.flyte.api.v1.DynamicWorkflowTaskRegistrar;
+import org.flyte.api.v1.IfBlock;
 import org.flyte.api.v1.Literal;
 import org.flyte.api.v1.NamedEntityIdentifier;
 import org.flyte.api.v1.Node;
@@ -280,6 +282,28 @@ private static List<Node> collectAllUsedTaskTemplates(
             flyteAdminClient,
             cache);
 
+    // collect task templates used by conditionals
+    spec.nodes().stream()
+        .filter(node -> node.branchNode() != null)
+        .forEach(
+            node -> {
+              List<Node> nodes = new ArrayList<>();
+              nodes.add(node.branchNode().ifElse().case_().thenNode());
+              nodes.add(node.branchNode().ifElse().elseNode());
+              nodes.addAll(
+                  node.branchNode().ifElse().other().stream()
+                      .map(IfBlock::thenNode)
+                      .collect(toList()));
+
+              collectTaskTemplates(
+                  nodes,
+                  nodesRewriter,
+                  allUsedTaskTemplates,
+                  allTaskTemplates,
+                  flyteAdminClient,
+                  cache);
+            });
+
     // collect task templates used by subworkflows
     allUsedSubWorkflows
         .values()