update WorkflowGraph

Signed-off-by: Jan Fiedler <jan@union.ai>
flyteorg · Apr 9, 2024 · 5075e5b · 5075e5b
1 parent f970b4e
commit 5075e5b
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 2 deletions.
diff --git a/packages/oss-console/src/components/WorkflowGraph/transformerWorkflowToDag.tsx b/packages/oss-console/src/components/WorkflowGraph/transformerWorkflowToDag.tsx
@@ -6,7 +6,7 @@ import { DISPLAY_NAME_END, DISPLAY_NAME_START } from '../flytegraph/ReactFlow/ut
 import { createDebugLogger } from '../../common/log';
 import { dTypes, dEdge, dNode } from '../../models/Graph/types';
 import { startNodeId, endNodeId } from '../../models/Node/constants';
-import { CompiledNode, ConnectionSet, TaskNode } from '../../models/Node/types';
+import { ArrayNode, CompiledNode, ConnectionSet, TaskNode } from '../../models/Node/types';
 import { CompiledTask } from '../../models/Task/types';
 import { CompiledWorkflow, CompiledWorkflowClosure } from '../../models/Workflow/types';
 import { isStartOrEndNode } from '../../models/Node/utils';
@@ -81,6 +81,7 @@ const createDNode = ({
     gateNode: compiledNode.gateNode,
     level: parentDNode?.level !== undefined ? parentDNode.level + 1 : 0,
     ...nodeMetadata,
+    ...(compiledNode.arrayNode ? { arrayNode: compiledNode.arrayNode } : {}),
     ...(compiledNode.workflowNode ? { workflowNode: compiledNode.workflowNode } : {}),
     ...(compiledNode.gateNode ? { gateNode: compiledNode.gateNode } : {}),
     ...(compiledNode.branchNode ? { taskNode: compiledNode.taskNode } : {}),
@@ -247,6 +248,21 @@ const parseNode = ({
         compiledWorkflowClosure,
       });
     }
+  } else if (node?.arrayNode) {
+    const arrayNode = (node.arrayNode as ArrayNode).node;
+    const taskNode = arrayNode.taskNode as TaskNode;
+    const taskType: CompiledTask = getTaskTypeFromCompiledNode(
+      taskNode,
+      compiledWorkflowClosure.tasks,
+    ) as CompiledTask;
+    dNode = createDNode({
+      compiledNode: node,
+      parentDNode: root,
+      taskTemplate: taskType,
+      nodeMetadataMap,
+      staticExecutionIdsMap,
+      compiledWorkflowClosure,
+    });
   } else if (node.taskNode) {
     const taskNode = node.taskNode as TaskNode;
     const taskType: CompiledTask = getTaskTypeFromCompiledNode(

diff --git a/packages/oss-console/src/components/WorkflowGraph/utils.ts b/packages/oss-console/src/components/WorkflowGraph/utils.ts
@@ -100,6 +100,9 @@ export const getSubWorkflowFromId = (
 };
 
 export const getTaskTypeFromCompiledNode = (taskNode: TaskNode, tasks: CompiledTask[]) => {
+  if (!taskNode.referenceId) {
+    return undefined;
+  }
   for (let i = 0; i < tasks.length; i++) {
     const compiledTask: CompiledTask = tasks[i];
     const taskTemplate: TaskTemplate = compiledTask.template;
@@ -108,7 +111,7 @@ export const getTaskTypeFromCompiledNode = (taskNode: TaskNode, tasks: CompiledT
       return compiledTask;
     }
   }
-  return null;
+  return undefined;
 };
 
 export const getNodeNameFromDag = (dagData: dNode, nodeId: string) => {