From 7cd476ed506b804d3900c2bd5ffa9dfef0592fb1 Mon Sep 17 00:00:00 2001
From: Brent Bovenzi <brent@astronomer.io>
Date: Wed, 17 Apr 2024 08:59:59 -0400
Subject: [PATCH] Graph view improvements (#38940)

* Initial pass at graph updates again

* Initial pass

* fixes

* Fix task_group tests

(cherry picked from commit 0fbe711dfa73d416a4ddd60caa29f1406cf0774c)
---
 airflow/utils/task_group.py                   |  5 +++
 airflow/www/static/js/dag/TaskName.tsx        |  1 -
 .../static/js/dag/details/graph/DagNode.tsx   | 34 ++++++++++++-------
 .../www/static/js/dag/grid/renderTaskRows.tsx |  5 ---
 airflow/www/static/js/types/index.ts          |  1 +
 airflow/www/static/js/utils/graph.ts          | 11 +++---
 6 files changed, 34 insertions(+), 23 deletions(-)

diff --git a/airflow/utils/task_group.py b/airflow/utils/task_group.py
index 50ce7b4089427..37b03e1907f15 100644
--- a/airflow/utils/task_group.py
+++ b/airflow/utils/task_group.py
@@ -679,13 +679,17 @@ def get_current_task_group(cls, dag: DAG | None) -> TaskGroup | None:
 def task_group_to_dict(task_item_or_group):
     """Create a nested dict representation of this TaskGroup and its children used to construct the Graph."""
     from airflow.models.abstractoperator import AbstractOperator
+    from airflow.models.mappedoperator import MappedOperator
 
     if isinstance(task := task_item_or_group, AbstractOperator):
         setup_teardown_type = {}
+        is_mapped = {}
         if task.is_setup is True:
             setup_teardown_type["setupTeardownType"] = "setup"
         elif task.is_teardown is True:
             setup_teardown_type["setupTeardownType"] = "teardown"
+        if isinstance(task, MappedOperator):
+            is_mapped["isMapped"] = True
         return {
             "id": task.task_id,
             "value": {
@@ -694,6 +698,7 @@ def task_group_to_dict(task_item_or_group):
                 "style": f"fill:{task.ui_color};",
                 "rx": 5,
                 "ry": 5,
+                **is_mapped,
                 **setup_teardown_type,
             },
         }
diff --git a/airflow/www/static/js/dag/TaskName.tsx b/airflow/www/static/js/dag/TaskName.tsx
index 05f851c18db3e..6b5f0a90bde05 100644
--- a/airflow/www/static/js/dag/TaskName.tsx
+++ b/airflow/www/static/js/dag/TaskName.tsx
@@ -54,7 +54,6 @@ const TaskName = ({
       data-testid={id}
       color={colors.gray[800]}
       fontSize={isZoomedOut ? 24 : undefined}
-      textAlign="justify"
       {...rest}
     >
       <chakra.span onClick={onClick}>
diff --git a/airflow/www/static/js/dag/details/graph/DagNode.tsx b/airflow/www/static/js/dag/details/graph/DagNode.tsx
index ec04f73c17396..7a448da88058c 100644
--- a/airflow/www/static/js/dag/details/graph/DagNode.tsx
+++ b/airflow/www/static/js/dag/details/graph/DagNode.tsx
@@ -18,7 +18,7 @@
  */
 
 import React from "react";
-import { Box, Flex, Text } from "@chakra-ui/react";
+import { Flex, Text, useTheme } from "@chakra-ui/react";
 import type { NodeProps } from "reactflow";
 
 import { SimpleStatus } from "src/dag/StatusBox";
@@ -53,10 +53,11 @@ const DagNode = ({
 }: NodeProps<CustomNodeProps>) => {
   const { onSelect } = useSelection();
   const containerRef = useContainerRef();
+  const { colors } = useTheme();
 
   if (!task) return null;
 
-  const bg = isOpen ? "blackAlpha.50" : "white";
+  const groupBg = isOpen ? `${colors.blue[500]}15` : "blue.50";
   const { isMapped } = task;
   const mappedStates = instance?.mappedStates;
 
@@ -67,9 +68,9 @@ const DagNode = ({
     : label;
 
   let operatorTextColor = "";
-  let operatorBG = "";
+  let operatorBg = "";
   if (style) {
-    [, operatorBG] = style.split(":");
+    [, operatorBg] = style.split(":");
   }
 
   if (labelStyle) {
@@ -83,6 +84,12 @@ const DagNode = ({
       ? stateColors[instance.state]
       : "gray.400";
 
+  let borderWidth = 2;
+  if (isZoomedOut) {
+    if (isSelected) borderWidth = 10;
+    else borderWidth = 6;
+  } else if (isSelected) borderWidth = 4;
+
   return (
     <Tooltip
       label={
@@ -95,15 +102,16 @@ const DagNode = ({
       placement="top"
       openDelay={hoverDelay}
     >
-      <Box
+      <Flex
         borderRadius={isZoomedOut ? 10 : 5}
-        borderWidth={(isSelected ? 4 : 2) * (isZoomedOut ? 3 : 1)}
+        borderWidth={borderWidth}
         borderColor={nodeBorderColor}
+        wordBreak="break-word"
         bg={
-          !task.children?.length && operatorBG
+          !task.children?.length && operatorBg
             ? // Fade the operator color to clash less with the task instance status
-              `color-mix(in srgb, ${operatorBG.replace(";", "")} 80%, white)`
-            : bg
+              `color-mix(in srgb, ${operatorBg.replace(";", "")} 80%, white)`
+            : groupBg
         }
         height={`${height}px`}
         width={`${width}px`}
@@ -121,6 +129,10 @@ const DagNode = ({
         }}
         px={isZoomedOut ? 1 : 2}
         mt={isZoomedOut ? -2 : 0}
+        alignItems={isZoomedOut && !isOpen ? "center" : undefined}
+        justifyContent={isZoomedOut && !isOpen ? "center" : undefined}
+        flexDirection="column"
+        overflow="wrap"
       >
         <TaskName
           label={taskName}
@@ -131,9 +143,7 @@ const DagNode = ({
             onToggleCollapse();
           }}
           setupTeardownType={setupTeardownType}
-          fontWeight="bold"
           isZoomedOut={isZoomedOut}
-          mt={isZoomedOut ? -2 : 0}
           noOfLines={2}
         />
         {!isZoomedOut && (
@@ -159,7 +169,7 @@ const DagNode = ({
             )}
           </>
         )}
-      </Box>
+      </Flex>
     </Tooltip>
   );
 };
diff --git a/airflow/www/static/js/dag/grid/renderTaskRows.tsx b/airflow/www/static/js/dag/grid/renderTaskRows.tsx
index dd77284fe60b9..9d688345bb982 100644
--- a/airflow/www/static/js/dag/grid/renderTaskRows.tsx
+++ b/airflow/www/static/js/dag/grid/renderTaskRows.tsx
@@ -182,11 +182,6 @@ const Row = (props: RowProps) => {
               pl={level * 4 + 4}
               setupTeardownType={task.setupTeardownType}
               pr={4}
-              fontWeight={
-                isGroup || (task.isMapped && !isParentMapped)
-                  ? "bold"
-                  : "normal"
-              }
               noOfLines={1}
             />
           </Td>
diff --git a/airflow/www/static/js/types/index.ts b/airflow/www/static/js/types/index.ts
index cf97aa8e0498b..79ed13c9bb6a2 100644
--- a/airflow/www/static/js/types/index.ts
+++ b/airflow/www/static/js/types/index.ts
@@ -143,6 +143,7 @@ interface DepNode {
     labelStyle?: string;
     style?: string;
     setupTeardownType?: "setup" | "teardown";
+    isMapped?: boolean;
   };
   children?: DepNode[];
   edges?: MidEdge[];
diff --git a/airflow/www/static/js/utils/graph.ts b/airflow/www/static/js/utils/graph.ts
index 7db3767994181..dde4fe64c3970 100644
--- a/airflow/www/static/js/utils/graph.ts
+++ b/airflow/www/static/js/utils/graph.ts
@@ -173,8 +173,10 @@ const generateGraph = ({
         }));
       closedGroupIds.push(id);
     }
-    const extraLabelLength =
-      value.label.length > 20 ? value.label.length - 19 : 0;
+
+    const label = value.isMapped ? `${value.label} [100]` : value.label;
+    const labelLength = getTextWidth(label, font);
+    const width = labelLength > 200 ? labelLength : 200;
 
     return {
       id,
@@ -184,9 +186,8 @@ const generateGraph = ({
         isJoinNode,
         childCount,
       },
-      // Make tasks with long names wider
-      width: isJoinNode ? 10 : 200 + extraLabelLength * 5,
-      height: isJoinNode ? 10 : 70,
+      width: isJoinNode ? 10 : width,
+      height: isJoinNode ? 10 : 80,
     };
   };