Improve task names in graph view

airflow/www/static/js/dag/grid/TaskName.test.tsx → airflow/www/static/js/dag/TaskName.test.tsx

@@ -38,26 +38,19 @@ describe("Test TaskName", () => {
 
   test("Displays a mapped task name", () => {
     const { getByText } = render(
-      <TaskName
-        level={0}
-        label="test"
-        id="test"
-        isMapped
-        onToggle={() => {}}
-      />,
+      <TaskName label="test" id="test" isMapped onToggle={() => {}} />,
       { wrapper: ChakraWrapper }
     );
 
     expect(getByText("test [ ]")).toBeDefined();
   });
 
   test("Displays a group task name", () => {
-    const { getByText, getByTestId } = render(
-      <TaskName level={0} label="test" id="test" isGroup onToggle={() => {}} />,
+    const { getByText } = render(
+      <TaskName label="test" id="test" isGroup onToggle={() => {}} />,
       { wrapper: ChakraWrapper }
     );
 
     expect(getByText("test")).toBeDefined();
-    expect(getByTestId("open-group")).toBeDefined();
   });
 });

airflow/www/static/js/dag/TaskName.tsx

@@ -0,0 +1,89 @@
+/*!
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+import React, { MouseEventHandler, CSSProperties } from "react";
+import { Text, TextProps, useTheme } from "@chakra-ui/react";
+import { FiChevronUp, FiArrowUpRight, FiArrowDownRight } from "react-icons/fi";
+
+interface Props extends TextProps {
+  isGroup?: boolean;
+  isMapped?: boolean;
+  onToggle: MouseEventHandler<HTMLDivElement>;
+  isOpen?: boolean;
+  label: string;
+  id?: string;
+  setupTeardownType?: string;
+  isZoomedOut?: boolean;
+}
+
+const TaskName = ({
+  isGroup = false,
+  isMapped = false,
+  onToggle,
+  isOpen = false,
+  label,
+  id,
+  setupTeardownType,
+  isZoomedOut,
+  ...rest
+}: Props) => {
+  const { colors } = useTheme();
+  const iconStyle: CSSProperties = {
+    display: "inline",
+    position: "relative",
+    verticalAlign: "middle",
+  };
+  return (
+    <Text
+      cursor={isGroup ? "pointer" : undefined}
+      onClick={onToggle}
+      data-testid={id}
+      width="100%"
+      color={colors.gray[800]}
+      fontSize={isZoomedOut ? 24 : undefined}
+      textAlign="justify"
+      {...rest}
+    >
+      {label}
+      {isMapped && " [ ]"}
+      {isGroup && (
+        <FiChevronUp
+          size={isZoomedOut ? 24 : 15}
+          strokeWidth={3}
+          style={{
+            transition: "transform 0.5s",
+            transform: `rotate(${isOpen ? 0 : 180}deg)`,
+            ...iconStyle,
+          }}
+        />
+      )}
+      {setupTeardownType === "setup" && (
+        <FiArrowUpRight size={isZoomedOut ? 24 : 15} style={iconStyle} />
+      )}
+      {setupTeardownType === "teardown" && (
+        <FiArrowDownRight size={isZoomedOut ? 24 : 15} style={iconStyle} />
+      )}
+    </Text>
+  );
+};
+
+// Only rerender the component if props change
+const MemoizedTaskName = React.memo(TaskName);
+
+export default MemoizedTaskName;

airflow/www/static/js/dag/details/graph/Node.test.tsx

@@ -110,7 +110,6 @@ describe("Test Graph Node", () => {
 
     expect(getByText("success")).toBeInTheDocument();
     expect(getByText("task_id")).toBeInTheDocument();
-    expect(getByText("+ 5 tasks")).toBeInTheDocument();
   });
 
   test("Renders normal task correctly", async () => {

airflow/www/static/js/dag/details/graph/Node.tsx

@@ -18,7 +18,7 @@
  */
 
 import React from "react";
-import { Box, Text, Flex, useTheme } from "@chakra-ui/react";
+import { Box, Text, Flex } from "@chakra-ui/react";
 import { Handle, NodeProps, Position } from "reactflow";
 
 import { SimpleStatus } from "src/dag/StatusBox";
@@ -28,7 +28,7 @@ import { getGroupAndMapSummary, hoverDelay } from "src/utils";
 import Tooltip from "src/components/Tooltip";
 import InstanceTooltip from "src/dag/InstanceTooltip";
 import { useContainerRef } from "src/context/containerRef";
-import { ImArrowUpRight2, ImArrowDownRight2 } from "react-icons/im";
+import TaskName from "src/dag/TaskName";
 
 export interface CustomNodeProps {
   label: string;
@@ -69,7 +69,6 @@ export const BaseNode = ({
     isZoomedOut,
   },
 }: NodeProps<CustomNodeProps>) => {
-  const { colors } = useTheme();
   const { onSelect } = useSelection();
   const containerRef = useContainerRef();
 
@@ -138,84 +137,48 @@ export const BaseNode = ({
             });
           }
         }}
+        px={isZoomedOut ? 1 : 2}
+        mt={isZoomedOut ? -2 : 0}
       >
-        <Flex
-          justifyContent="space-between"
-          width={width}
-          p={2}
-          flexWrap={isZoomedOut ? "nowrap" : "wrap"}
-          flexDirection={isZoomedOut ? "row" : "column"}
-        >
-          <Flex flexDirection="column" width="100%">
-            <Flex
-              justifyContent="space-between"
-              alignItems="center"
-              width="100%"
-              fontSize={isZoomedOut ? 25 : undefined}
-              fontWeight="bold"
-            >
-              <Text noOfLines={1} maxWidth={`calc(${width}px - 8px)`}>
-                {taskName}
+        <TaskName
+          label={taskName}
+          isOpen={isOpen}
+          isGroup={!!childCount}
+          onToggle={(e) => {
+            e.stopPropagation();
+            onToggleCollapse();
+          }}
+          setupTeardownType={setupTeardownType}
+          fontWeight="bold"
+          isZoomedOut={isZoomedOut}
+          mt={isZoomedOut ? -2 : 0}
+          noOfLines={2}
+        />
+        {!isZoomedOut && (
+          <>
+            {!!instance && instance.state && (
+              <Flex alignItems="center">
+                <SimpleStatus state={instance.state} />
+                <Text ml={2} color="gray.500" fontWeight={400} fontSize="md">
+                  {instance.state}
+                </Text>
+              </Flex>
+            )}
+            {task?.operator && (
+              <Text
+                noOfLines={1}
+                maxWidth={`calc(${width}px - 12px)`}
+                fontWeight={400}
+                fontSize="md"
+                width="fit-content"
+                color={operatorTextColor}
+                px={1}
+              >
+                {task.operator}
               </Text>
-              {setupTeardownType === "setup" && (
-                <ImArrowUpRight2 size={15} color={colors.gray[800]} />
-              )}
-              {setupTeardownType === "teardown" && (
-                <ImArrowDownRight2 size={15} color={colors.gray[800]} />
-              )}
-            </Flex>
-            {!isZoomedOut && (
-              <>
-                {!!instance && instance.state && (
-                  <Flex alignItems="center">
-                    <SimpleStatus state={instance.state} />
-                    <Text
-                      ml={2}
-                      color="gray.500"
-                      fontWeight={400}
-                      fontSize="md"
-                    >
-                      {instance.state}
-                    </Text>
-                  </Flex>
-                )}
-                {task?.operator && (
-                  <Text
-                    noOfLines={1}
-                    maxWidth={`calc(${width}px - 12px)`}
-                    fontWeight={400}
-                    fontSize="md"
-                    width="fit-content"
-                    color={operatorTextColor}
-                    px={1}
-                  >
-                    {task.operator}
-                  </Text>
-                )}
-              </>
             )}
-          </Flex>
-          {!!childCount && (
-            <Text
-              noOfLines={1}
-              fontSize={isZoomedOut ? 25 : undefined}
-              color="blue.600"
-              cursor="pointer"
-              width="150px"
-              fontWeight="bold"
-              // Increase the target area to expand/collapse a group
-              p={2}
-              m={-2}
-              onClick={(e) => {
-                e.stopPropagation();
-                onToggleCollapse();
-              }}
-            >
-              {isOpen ? "- " : "+ "}
-              {childCount} tasks
-            </Text>
-          )}
-        </Flex>
+          </>
+        )}
       </Box>
     </Tooltip>
   );