From 64440c01e83f854632545830ef6b0ace6ab7d1b3 Mon Sep 17 00:00:00 2001
From: olga-union <101579322+olga-union@users.noreply.github.com>
Date: Fri, 15 Jul 2022 14:58:14 -0500
Subject: [PATCH] fix: graph has realtime updates as execution progresses
 (#543)

Signed-off-by: Olga Nad <olga@union.ai>
---
 .../ExecutionWorkflowGraph.tsx                | 116 +++++++++++-------
 .../Executions/nodeExecutionQueries.ts        |  10 +-
 2 files changed, 77 insertions(+), 49 deletions(-)

diff --git a/packages/zapp/console/src/components/Executions/ExecutionDetails/ExecutionWorkflowGraph.tsx b/packages/zapp/console/src/components/Executions/ExecutionDetails/ExecutionWorkflowGraph.tsx
index a536ecd46..b0181cdec 100644
--- a/packages/zapp/console/src/components/Executions/ExecutionDetails/ExecutionWorkflowGraph.tsx
+++ b/packages/zapp/console/src/components/Executions/ExecutionDetails/ExecutionWorkflowGraph.tsx
@@ -13,8 +13,8 @@ import * as React from 'react';
 import { useEffect, useMemo, useState } from 'react';
 import { useQuery, useQueryClient } from 'react-query';
 import { NodeExecutionsContext } from '../contexts';
+import { fetchTaskExecutionList } from '../taskExecutionQueries';
 import { getGroupedLogs } from '../TaskExecutionsList/utils';
-import { useTaskExecutions, useTaskExecutionsRefresher } from '../useTaskExecutions';
 import { NodeExecutionDetailsPanelContent } from './NodeExecutionDetailsPanelContent';
 
 export interface ExecutionWorkflowGraphProps {
@@ -22,56 +22,27 @@ export interface ExecutionWorkflowGraphProps {
   workflowId: WorkflowId;
 }
 
+interface WorkflowNodeExecution extends NodeExecution {
+  logsByPhase?: LogsByPhase;
+}
+
 /** Wraps a WorkflowGraph, customizing it to also show execution statuses */
 export const ExecutionWorkflowGraph: React.FC<ExecutionWorkflowGraphProps> = ({
   nodeExecutions,
   workflowId,
 }) => {
-  const workflowQuery = useQuery<Workflow, Error>(makeWorkflowQuery(useQueryClient(), workflowId));
-
-  const nodeExecutionsWithResources = nodeExecutions.map((nodeExecution) => {
-    const taskExecutions = useTaskExecutions(nodeExecution.id);
-    useTaskExecutionsRefresher(nodeExecution, taskExecutions);
-
-    const useNewMapTaskView = taskExecutions.value.every((taskExecution) => {
-      const {
-        closure: { taskType, metadata, eventVersion = 0 },
-      } = taskExecution;
-      return isMapTaskV1(
-        eventVersion,
-        metadata?.externalResources?.length ?? 0,
-        taskType ?? undefined,
-      );
-    });
-    const externalResources: ExternalResource[] = taskExecutions.value
-      .map((taskExecution) => taskExecution.closure.metadata?.externalResources)
-      .flat()
-      .filter((resource): resource is ExternalResource => !!resource);
+  const queryClient = useQueryClient();
+  const workflowQuery = useQuery<Workflow, Error>(makeWorkflowQuery(queryClient, workflowId));
 
-    const logsByPhase: LogsByPhase = getGroupedLogs(externalResources);
-
-    return {
-      ...nodeExecution,
-      ...(useNewMapTaskView && logsByPhase.size > 0 && { logsByPhase }),
-    };
-  });
+  const [nodeExecutionsWithResources, setNodeExecutionsWithResources] = useState<
+    WorkflowNodeExecution[]
+  >([]);
+  const [selectedNodes, setSelectedNodes] = useState<string[]>([]);
 
   const nodeExecutionsById = useMemo(
     () => keyBy(nodeExecutionsWithResources, 'scopedId'),
     [nodeExecutionsWithResources],
   );
-
-  const [selectedNodes, setSelectedNodes] = useState<string[]>([]);
-  const onNodeSelectionChanged = (newSelection: string[]) => {
-    const validSelection = newSelection.filter((nodeId) => {
-      if (nodeId === startNodeId || nodeId === endNodeId) {
-        return false;
-      }
-      return true;
-    });
-    setSelectedNodes(validSelection);
-  };
-
   // Note: flytegraph allows multiple selection, but we only support showing
   // a single item in the details panel
   const selectedExecution = selectedNodes.length
@@ -83,19 +54,72 @@ export const ExecutionWorkflowGraph: React.FC<ExecutionWorkflowGraphProps> = ({
         }
     : null;
 
-  const onCloseDetailsPanel = () => {
-    setSelectedPhase(undefined);
-    setIsDetailsTabClosed(true);
-    setSelectedNodes([]);
-  };
-
   const [selectedPhase, setSelectedPhase] = useState<TaskExecutionPhase | undefined>(undefined);
   const [isDetailsTabClosed, setIsDetailsTabClosed] = useState<boolean>(!selectedExecution);
 
+  useEffect(() => {
+    let isCurrent = true;
+    async function fetchData(nodeExecutions, queryClient) {
+      const newValue = await Promise.all(
+        nodeExecutions.map(async (nodeExecution) => {
+          const taskExecutions = await fetchTaskExecutionList(queryClient, nodeExecution.id);
+
+          const useNewMapTaskView = taskExecutions.every((taskExecution) => {
+            const {
+              closure: { taskType, metadata, eventVersion = 0 },
+            } = taskExecution;
+            return isMapTaskV1(
+              eventVersion,
+              metadata?.externalResources?.length ?? 0,
+              taskType ?? undefined,
+            );
+          });
+          const externalResources: ExternalResource[] = taskExecutions
+            .map((taskExecution) => taskExecution.closure.metadata?.externalResources)
+            .flat()
+            .filter((resource): resource is ExternalResource => !!resource);
+
+          const logsByPhase: LogsByPhase = getGroupedLogs(externalResources);
+
+          return {
+            ...nodeExecution,
+            ...(useNewMapTaskView && logsByPhase.size > 0 && { logsByPhase }),
+          };
+        }),
+      );
+
+      if (isCurrent) {
+        setNodeExecutionsWithResources(newValue);
+      }
+    }
+
+    fetchData(nodeExecutions, queryClient);
+
+    return () => {
+      isCurrent = false;
+    };
+  }, [nodeExecutions]);
+
   useEffect(() => {
     setIsDetailsTabClosed(!selectedExecution);
   }, [selectedExecution]);
 
+  const onNodeSelectionChanged = (newSelection: string[]) => {
+    const validSelection = newSelection.filter((nodeId) => {
+      if (nodeId === startNodeId || nodeId === endNodeId) {
+        return false;
+      }
+      return true;
+    });
+    setSelectedNodes(validSelection);
+  };
+
+  const onCloseDetailsPanel = () => {
+    setSelectedPhase(undefined);
+    setIsDetailsTabClosed(true);
+    setSelectedNodes([]);
+  };
+
   const renderGraph = (workflow: Workflow) => (
     <WorkflowGraph
       onNodeSelectionChanged={onNodeSelectionChanged}
diff --git a/packages/zapp/console/src/components/Executions/nodeExecutionQueries.ts b/packages/zapp/console/src/components/Executions/nodeExecutionQueries.ts
index 7d5e3b18b..1db75713e 100644
--- a/packages/zapp/console/src/components/Executions/nodeExecutionQueries.ts
+++ b/packages/zapp/console/src/components/Executions/nodeExecutionQueries.ts
@@ -333,9 +333,11 @@ export function useAllChildNodeExecutionGroupsQuery(
   const shouldEnableFn = (groups) => {
     if (groups.length > 0) {
       return groups.some((group) => {
-        if (group.nodeExecutions?.length > 0) {
+        // non-empty groups are wrapped in array
+        const unwrappedGroup = Array.isArray(group) ? group[0] : group;
+        if (unwrappedGroup.nodeExecutions?.length > 0) {
           /* Return true is any executions are not yet terminal (ie, they can change) */
-          return group.nodeExecutions.some((ne) => {
+          return unwrappedGroup.nodeExecutions.some((ne) => {
             return !nodeExecutionIsTerminal(ne);
           });
         } else {
@@ -347,9 +349,11 @@ export function useAllChildNodeExecutionGroupsQuery(
     }
   };
 
+  const key = `${nodeExecutions?.[0]?.scopedId}-${nodeExecutions?.[0]?.closure?.phase}`;
+
   return useConditionalQuery<Array<NodeExecutionGroup[]>>(
     {
-      queryKey: [QueryType.NodeExecutionChildList, nodeExecutions[0]?.id, config],
+      queryKey: [QueryType.NodeExecutionChildList, key, config],
       queryFn: () => fetchAllChildNodeExecutions(queryClient, nodeExecutions, config),
     },
     shouldEnableFn,