From 6a0e83f310cc66780b9c986413f658adfa9f087e Mon Sep 17 00:00:00 2001 From: Brent Bovenzi Date: Thu, 4 Apr 2024 15:40:44 -0400 Subject: [PATCH 1/4] Initial pass at graph updates again --- airflow/www/static/js/dag/TaskName.tsx | 1 - .../static/js/dag/details/graph/DagNode.tsx | 35 ++++++++++++------- .../www/static/js/dag/grid/renderTaskRows.tsx | 10 +++--- airflow/www/static/js/types/index.ts | 1 + airflow/www/static/js/utils/graph.ts | 11 +++--- 5 files changed, 35 insertions(+), 23 deletions(-) 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} > diff --git a/airflow/www/static/js/dag/details/graph/DagNode.tsx b/airflow/www/static/js/dag/details/graph/DagNode.tsx index ec04f73c17396..2621e0a2f79be 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,12 @@ const DagNode = ({ }: NodeProps) => { const { onSelect } = useSelection(); const containerRef = useContainerRef(); + const { colors } = useTheme(); if (!task) return null; - const bg = isOpen ? "blackAlpha.50" : "white"; + const isGroup = !!task.children?.length; + const groupBg = isOpen ? `${colors.blue[500]}15` : "blue.50"; const { isMapped } = task; const mappedStates = instance?.mappedStates; @@ -67,9 +69,9 @@ const DagNode = ({ : label; let operatorTextColor = ""; - let operatorBG = ""; + let operatorBg = ""; if (style) { - [, operatorBG] = style.split(":"); + [, operatorBg] = style.split(":"); } if (labelStyle) { @@ -83,6 +85,12 @@ const DagNode = ({ ? stateColors[instance.state] : "gray.400"; + let borderWidth = 2; + if (isZoomedOut) { + if (isSelected) borderWidth = 8; + else borderWidth = 6; + } else if (isSelected) borderWidth = 4; + return ( - {!isZoomedOut && ( @@ -159,7 +170,7 @@ const DagNode = ({ )} )} - + ); }; diff --git a/airflow/www/static/js/dag/grid/renderTaskRows.tsx b/airflow/www/static/js/dag/grid/renderTaskRows.tsx index dd77284fe60b9..a2d91b68dc9f6 100644 --- a/airflow/www/static/js/dag/grid/renderTaskRows.tsx +++ b/airflow/www/static/js/dag/grid/renderTaskRows.tsx @@ -182,11 +182,11 @@ const Row = (props: RowProps) => { pl={level * 4 + 4} setupTeardownType={task.setupTeardownType} pr={4} - fontWeight={ - isGroup || (task.isMapped && !isParentMapped) - ? "bold" - : "normal" - } + // fontWeight={ + // isGroup || (task.isMapped && !isParentMapped) + // ? "bold" + // : "normal" + // } noOfLines={1} /> 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, }; }; From 85675569f1b4b458d774a93aa2ddb86b18eef45c Mon Sep 17 00:00:00 2001 From: Brent Bovenzi Date: Thu, 4 Apr 2024 15:56:18 -0400 Subject: [PATCH 2/4] Initial pass --- airflow/utils/task_group.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/airflow/utils/task_group.py b/airflow/utils/task_group.py index 50ce7b4089427..e3e5ecb52e1ec 100644 --- a/airflow/utils/task_group.py +++ b/airflow/utils/task_group.py @@ -679,9 +679,11 @@ 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 = isinstance(task, MappedOperator) if task.is_setup is True: setup_teardown_type["setupTeardownType"] = "setup" elif task.is_teardown is True: @@ -694,6 +696,7 @@ def task_group_to_dict(task_item_or_group): "style": f"fill:{task.ui_color};", "rx": 5, "ry": 5, + "isMapped": is_mapped, **setup_teardown_type, }, } From e6094c53a9a4b015b064693a91d611a6a1581594 Mon Sep 17 00:00:00 2001 From: Brent Bovenzi Date: Thu, 11 Apr 2024 13:06:09 -0400 Subject: [PATCH 3/4] fixes --- airflow/www/static/js/dag/details/graph/DagNode.tsx | 7 +++---- airflow/www/static/js/dag/grid/renderTaskRows.tsx | 5 ----- 2 files changed, 3 insertions(+), 9 deletions(-) diff --git a/airflow/www/static/js/dag/details/graph/DagNode.tsx b/airflow/www/static/js/dag/details/graph/DagNode.tsx index 2621e0a2f79be..7a448da88058c 100644 --- a/airflow/www/static/js/dag/details/graph/DagNode.tsx +++ b/airflow/www/static/js/dag/details/graph/DagNode.tsx @@ -57,7 +57,6 @@ const DagNode = ({ if (!task) return null; - const isGroup = !!task.children?.length; const groupBg = isOpen ? `${colors.blue[500]}15` : "blue.50"; const { isMapped } = task; const mappedStates = instance?.mappedStates; @@ -87,7 +86,7 @@ const DagNode = ({ let borderWidth = 2; if (isZoomedOut) { - if (isSelected) borderWidth = 8; + if (isSelected) borderWidth = 10; else borderWidth = 6; } else if (isSelected) borderWidth = 4; @@ -130,8 +129,8 @@ const DagNode = ({ }} px={isZoomedOut ? 1 : 2} mt={isZoomedOut ? -2 : 0} - alignItems={isZoomedOut && !isGroup ? "center" : undefined} - justifyContent={isZoomedOut && !isGroup ? "center" : undefined} + alignItems={isZoomedOut && !isOpen ? "center" : undefined} + justifyContent={isZoomedOut && !isOpen ? "center" : undefined} flexDirection="column" overflow="wrap" > diff --git a/airflow/www/static/js/dag/grid/renderTaskRows.tsx b/airflow/www/static/js/dag/grid/renderTaskRows.tsx index a2d91b68dc9f6..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} /> From 327bb0b8580063b22323bb6a652e1b67189e8f66 Mon Sep 17 00:00:00 2001 From: Brent Bovenzi Date: Tue, 16 Apr 2024 13:04:36 -0400 Subject: [PATCH 4/4] Fix task_group tests --- airflow/utils/task_group.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/airflow/utils/task_group.py b/airflow/utils/task_group.py index e3e5ecb52e1ec..37b03e1907f15 100644 --- a/airflow/utils/task_group.py +++ b/airflow/utils/task_group.py @@ -683,11 +683,13 @@ def task_group_to_dict(task_item_or_group): if isinstance(task := task_item_or_group, AbstractOperator): setup_teardown_type = {} - is_mapped = isinstance(task, MappedOperator) + 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": { @@ -696,7 +698,7 @@ def task_group_to_dict(task_item_or_group): "style": f"fill:{task.ui_color};", "rx": 5, "ry": 5, - "isMapped": is_mapped, + **is_mapped, **setup_teardown_type, }, }