diff --git a/airflow/api_fastapi/core_api/datamodels/ui/structure.py b/airflow/api_fastapi/core_api/datamodels/ui/structure.py index e3df958a1a81f..cc234500eb7a9 100644 --- a/airflow/api_fastapi/core_api/datamodels/ui/structure.py +++ b/airflow/api_fastapi/core_api/datamodels/ui/structure.py @@ -34,12 +34,12 @@ class NodeResponse(BaseModel): """Node serializer for responses.""" children: list[NodeResponse] | None = None - id: str | None + id: str is_mapped: bool | None = None - label: str | None = None + label: str tooltip: str | None = None setup_teardown_type: Literal["setup", "teardown"] | None = None - type: Literal["join", "sensor", "task", "task_group"] + type: Literal["join", "sensor", "task", "asset_condition"] class StructureDataResponse(BaseModel): diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml index d6ea8dcec50e6..010a4141e06b6 100644 --- a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml +++ b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml @@ -7782,9 +7782,7 @@ components: - type: 'null' title: Children id: - anyOf: - - type: string - - type: 'null' + type: string title: Id is_mapped: anyOf: @@ -7792,9 +7790,7 @@ components: - type: 'null' title: Is Mapped label: - anyOf: - - type: string - - type: 'null' + type: string title: Label tooltip: anyOf: @@ -7815,11 +7811,12 @@ components: - join - sensor - task - - task_group + - asset_condition title: Type type: object required: - id + - label - type title: NodeResponse description: Node serializer for responses. diff --git a/airflow/ui/openapi-gen/requests/schemas.gen.ts b/airflow/ui/openapi-gen/requests/schemas.gen.ts index 0f704159c8dc5..cb8e2e7396e3f 100644 --- a/airflow/ui/openapi-gen/requests/schemas.gen.ts +++ b/airflow/ui/openapi-gen/requests/schemas.gen.ts @@ -3128,14 +3128,7 @@ export const $NodeResponse = { title: "Children", }, id: { - anyOf: [ - { - type: "string", - }, - { - type: "null", - }, - ], + type: "string", title: "Id", }, is_mapped: { @@ -3150,14 +3143,7 @@ export const $NodeResponse = { title: "Is Mapped", }, label: { - anyOf: [ - { - type: "string", - }, - { - type: "null", - }, - ], + type: "string", title: "Label", }, tooltip: { @@ -3185,12 +3171,12 @@ export const $NodeResponse = { }, type: { type: "string", - enum: ["join", "sensor", "task", "task_group"], + enum: ["join", "sensor", "task", "asset_condition"], title: "Type", }, }, type: "object", - required: ["id", "type"], + required: ["id", "label", "type"], title: "NodeResponse", description: "Node serializer for responses.", } as const; diff --git a/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow/ui/openapi-gen/requests/types.gen.ts index e5aed24a98cd5..13321ad442483 100644 --- a/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow/ui/openapi-gen/requests/types.gen.ts @@ -761,15 +761,15 @@ export type JobResponse = { */ export type NodeResponse = { children?: Array | null; - id: string | null; + id: string; is_mapped?: boolean | null; - label?: string | null; + label: string; tooltip?: string | null; setup_teardown_type?: "setup" | "teardown" | null; - type: "join" | "sensor" | "task" | "task_group"; + type: "join" | "sensor" | "task" | "asset_condition"; }; -export type type = "join" | "sensor" | "task" | "task_group"; +export type type = "join" | "sensor" | "task" | "asset_condition"; /** * Request body for Clear Task Instances endpoint. diff --git a/airflow/ui/src/pages/DagsList/Dag/DagVizModal.tsx b/airflow/ui/src/pages/DagsList/Dag/DagVizModal.tsx index 726e97527cca3..9ab868b8f8d28 100644 --- a/airflow/ui/src/pages/DagsList/Dag/DagVizModal.tsx +++ b/airflow/ui/src/pages/DagsList/Dag/DagVizModal.tsx @@ -24,24 +24,28 @@ import { Dialog } from "src/components/ui"; import { Graph } from "./Graph"; type TriggerDAGModalProps = { - dagDisplayName: DAGResponse["dag_display_name"]; + dagDisplayName?: DAGResponse["dag_display_name"]; + dagId?: DAGResponse["dag_id"]; onClose: () => void; open: boolean; }; export const DagVizModal: React.FC = ({ dagDisplayName, + dagId, onClose, open, }) => ( - {dagDisplayName} + + {Boolean(dagDisplayName) ? dagDisplayName : "Dag Undefined"} + - + {dagId === undefined ? undefined : } diff --git a/airflow/ui/src/pages/DagsList/Dag/Graph/Graph.tsx b/airflow/ui/src/pages/DagsList/Dag/Graph/Graph.tsx index d80f0ba2f1a0f..84fbc2be114f0 100644 --- a/airflow/ui/src/pages/DagsList/Dag/Graph/Graph.tsx +++ b/airflow/ui/src/pages/DagsList/Dag/Graph/Graph.tsx @@ -20,13 +20,14 @@ import { Flex } from "@chakra-ui/react"; import { ReactFlow, Controls, Background, MiniMap } from "@xyflow/react"; import "@xyflow/react/dist/style.css"; +import { useStructureServiceStructureData } from "openapi/queries"; +import type { DAGResponse } from "openapi/requests/types.gen"; import { useColorMode } from "src/context/colorMode"; import { useOpenGroups } from "src/context/openGroups"; import Edge from "./Edge"; import { JoinNode } from "./JoinNode"; import { TaskNode } from "./TaskNode"; -import { graphData } from "./data"; import { useGraphLayout } from "./useGraphLayout"; const nodeTypes = { @@ -35,10 +36,16 @@ const nodeTypes = { }; const edgeTypes = { custom: Edge }; -export const Graph = () => { +export const Graph = ({ dagId }: { readonly dagId: DAGResponse["dag_id"] }) => { const { colorMode } = useColorMode(); const { openGroupIds } = useOpenGroups(); + + const { data: graphData = { arrange: "LR", edges: [], nodes: [] } } = + useStructureServiceStructureData({ + dagId, + }); + const { data } = useGraphLayout({ ...graphData, openGroupIds, diff --git a/airflow/ui/src/pages/DagsList/Dag/Graph/TaskName.tsx b/airflow/ui/src/pages/DagsList/Dag/Graph/TaskName.tsx index 2747d8d1e82a0..6807225360cef 100644 --- a/airflow/ui/src/pages/DagsList/Dag/Graph/TaskName.tsx +++ b/airflow/ui/src/pages/DagsList/Dag/Graph/TaskName.tsx @@ -20,7 +20,7 @@ import { Text, type TextProps } from "@chakra-ui/react"; import type { CSSProperties } from "react"; import { FiArrowUpRight, FiArrowDownRight } from "react-icons/fi"; -import type { Node } from "./data"; +import type { NodeResponse } from "openapi/requests/types.gen"; type Props = { readonly id: string; @@ -29,7 +29,7 @@ type Props = { readonly isOpen?: boolean; readonly isZoomedOut?: boolean; readonly label: string; - readonly setupTeardownType?: Node["setup_teardown_type"]; + readonly setupTeardownType?: NodeResponse["setup_teardown_type"]; } & TextProps; export const TaskName = ({ diff --git a/airflow/ui/src/pages/DagsList/Dag/Graph/data.ts b/airflow/ui/src/pages/DagsList/Dag/Graph/data.ts deleted file mode 100644 index 68759fc4ebf73..0000000000000 --- a/airflow/ui/src/pages/DagsList/Dag/Graph/data.ts +++ /dev/null @@ -1,216 +0,0 @@ -/*! - * 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. - */ - -export type Edge = { - is_setup_teardown?: boolean; - label?: string; - source_id: string; - target_id: string; -}; - -export type Node = { - children?: Array; - id: string; - is_mapped?: boolean; - label: string; - setup_teardown_type?: "setup" | "teardown"; - tooltip?: string; - type: - | "asset_alias" - | "asset_condition" - | "asset" - | "dag" - | "join" - | "sensor" - | "task" - | "trigger"; -}; - -export type GraphData = { - arrange: "BT" | "LR" | "RL" | "TB"; - edges: Array; - nodes: Array; -}; - -export const graphData: GraphData = { - arrange: "LR", - edges: [ - { - source_id: "section_1.upstream_join_id", - target_id: "section_1.taskgroup_setup", - }, - { - source_id: "section_1.downstream_join_id", - target_id: "section_2.upstream_join_id", - }, - { - source_id: "section_1.normal", - target_id: "section_1.taskgroup_teardown", - }, - { - is_setup_teardown: true, - label: "setup and teardown", - source_id: "section_1.taskgroup_setup", - target_id: "section_1.taskgroup_teardown", - }, - { - label: "test", - source_id: "section_1.taskgroup_teardown", - target_id: "section_1.downstream_join_id", - }, - { - source_id: "section_1.taskgroup_setup", - target_id: "section_1.normal", - }, - { - source_id: "section_2.downstream_join_id", - target_id: "end", - }, - { - source_id: "section_2.inner_section_2.task_2", - target_id: "section_2.inner_section_2.task_4", - }, - { - source_id: "section_2.inner_section_2.task_3", - target_id: "section_2.inner_section_2.task_4", - }, - { - source_id: "section_2.inner_section_2.task_4", - target_id: "section_2.downstream_join_id", - }, - { - source_id: "section_2.task_1", - target_id: "section_2.downstream_join_id", - }, - { - source_id: "section_2.upstream_join_id", - target_id: "section_2.inner_section_2.task_2", - }, - { - source_id: "section_2.upstream_join_id", - target_id: "section_2.inner_section_2.task_3", - }, - { - source_id: "section_2.upstream_join_id", - target_id: "section_2.task_1", - }, - { - label: "I am a realllllllllllllllllly long label", - source_id: "start", - target_id: "section_1.upstream_join_id", - }, - ], - nodes: [ - { - id: "end", - label: "end", - type: "task", - }, - { - children: [ - { - id: "section_1.normal", - label: "normal", - type: "task", - }, - { - id: "section_1.taskgroup_setup", - label: "taskgroup_setup", - setup_teardown_type: "setup", - type: "task", - }, - { - id: "section_1.taskgroup_teardown", - label: "taskgroup_teardown", - setup_teardown_type: "teardown", - type: "task", - }, - { - id: "section_1.upstream_join_id", - label: "", - type: "join", - }, - { - id: "section_1.downstream_join_id", - label: "", - type: "join", - }, - ], - id: "section_1", - is_mapped: false, - label: "section_1", - tooltip: "Tasks for section_1", - type: "task", - }, - { - children: [ - { - children: [ - { - id: "section_2.inner_section_2.task_2", - label: "task_2", - type: "task", - }, - { - id: "section_2.inner_section_2.task_3", - is_mapped: true, - label: "task_3", - type: "task", - }, - { - id: "section_2.inner_section_2.task_4", - label: "task_4", - type: "task", - }, - ], - id: "section_2.inner_section_2", - label: "inner_section_2", - tooltip: "Tasks for inner_section2", - type: "task", - }, - { - id: "section_2.task_1", - is_mapped: true, - label: "task_1", - type: "task", - }, - { - id: "section_2.upstream_join_id", - label: "", - type: "join", - }, - { - id: "section_2.downstream_join_id", - label: "", - type: "join", - }, - ], - id: "section_2", - is_mapped: false, - label: "section_2", - tooltip: "Tasks for section_2", - type: "task", - }, - { - id: "start", - label: "start", - type: "task", - }, - ], -}; diff --git a/airflow/ui/src/pages/DagsList/Dag/Graph/reactflowUtils.ts b/airflow/ui/src/pages/DagsList/Dag/Graph/reactflowUtils.ts index 76eb4c7489f36..793eb840c5d77 100644 --- a/airflow/ui/src/pages/DagsList/Dag/Graph/reactflowUtils.ts +++ b/airflow/ui/src/pages/DagsList/Dag/Graph/reactflowUtils.ts @@ -19,7 +19,8 @@ import type { Node as FlowNodeType, Edge as FlowEdgeType } from "@xyflow/react"; import type { ElkExtendedEdge } from "elkjs"; -import type { Node } from "./data"; +import type { NodeResponse } from "openapi/requests/types.gen"; + import type { LayoutNode } from "./useGraphLayout"; export type CustomNodeProps = { @@ -30,7 +31,7 @@ export type CustomNodeProps = { isMapped?: boolean; isOpen?: boolean; label: string; - setupTeardownType?: Node["setup_teardown_type"]; + setupTeardownType?: NodeResponse["setup_teardown_type"]; width?: number; }; @@ -103,7 +104,7 @@ export const flattenGraph = ({ if (node.children) { const { edges: childEdges, nodes: childNodes } = flattenGraph({ - children: node.children, + children: node.children as Array, parent: newNode, }); diff --git a/airflow/ui/src/pages/DagsList/Dag/Graph/useGraphLayout.ts b/airflow/ui/src/pages/DagsList/Dag/Graph/useGraphLayout.ts index cd0d1089078ca..352c9ada00891 100644 --- a/airflow/ui/src/pages/DagsList/Dag/Graph/useGraphLayout.ts +++ b/airflow/ui/src/pages/DagsList/Dag/Graph/useGraphLayout.ts @@ -19,7 +19,12 @@ import { useQuery } from "@tanstack/react-query"; import ELK, { type ElkNode, type ElkExtendedEdge, type ElkShape } from "elkjs"; -import type { Edge, Node } from "./data"; +import type { + EdgeResponse, + NodeResponse, + StructureDataResponse, +} from "openapi/requests/types.gen"; + import { flattenGraph, formatFlowEdges } from "./reactflowUtils"; type EdgeLabel = { @@ -35,9 +40,9 @@ type FormattedNode = { isGroup: boolean; isMapped?: boolean; isOpen?: boolean; - setupTeardownType?: Node["setup_teardown_type"]; + setupTeardownType?: NodeResponse["setup_teardown_type"]; } & ElkShape & - Node; + NodeResponse; type FormattedEdge = { id: string; @@ -46,7 +51,7 @@ type FormattedEdge = { parentNode?: string; } & ElkExtendedEdge; -export type LayoutNode = ElkNode & Node; +export type LayoutNode = ElkNode & NodeResponse; // Take text and font to calculate how long each node should be const getTextWidth = (text: string, font: string) => { @@ -76,15 +81,16 @@ const getDirection = (arrange: string) => { }; const formatElkEdge = ( - edge: Edge, + edge: EdgeResponse, font: string, - node?: Node, + node?: NodeResponse, ): FormattedEdge => ({ id: `${edge.source_id}-${edge.target_id}`, - isSetupTeardown: edge.is_setup_teardown, + isSetupTeardown: + edge.is_setup_teardown === null ? undefined : edge.is_setup_teardown, // isSourceAsset: e.isSourceAsset, labels: - edge.label === undefined + edge.label === undefined || edge.label === null ? [] : [ { @@ -99,7 +105,7 @@ const formatElkEdge = ( targets: [edge.target_id], }); -const getNestedChildIds = (children: Array) => { +const getNestedChildIds = (children: Array) => { let childIds: Array = []; children.forEach((child) => { @@ -116,9 +122,9 @@ const getNestedChildIds = (children: Array) => { type GenerateElkProps = { arrange: string; - edges: Array; + edges: Array; font: string; - nodes: Array; + nodes: Array; openGroupIds?: Array; }; @@ -132,15 +138,17 @@ const generateElkGraph = ({ const closedGroupIds: Array = []; let filteredEdges = unformattedEdges; - const formatChildNode = (node: Node): FormattedNode => { + const formatChildNode = (node: NodeResponse): FormattedNode => { const isOpen = openGroupIds?.includes(node.id); const childCount = node.children?.filter((child) => child.type !== "join").length ?? 0; const childIds = - node.children === undefined ? [] : getNestedChildIds(node.children); + node.children === null || node.children === undefined + ? [] + : getNestedChildIds(node.children); - if (isOpen && node.children !== undefined) { + if (isOpen && node.children !== null && node.children !== undefined) { return { ...node, childCount, @@ -212,7 +220,7 @@ const generateElkGraph = ({ height, id: node.id, isGroup: Boolean(node.children), - isMapped: node.is_mapped, + isMapped: node.is_mapped === null ? undefined : node.is_mapped, label: node.label, setupTeardownType: node.setup_teardown_type, type: node.type, @@ -225,7 +233,7 @@ const generateElkGraph = ({ const edges = filteredEdges.map((fe) => formatElkEdge(fe, font)); return { - children, + children: children as Array, edges, id: "root", layoutOptions: { @@ -238,11 +246,8 @@ const generateElkGraph = ({ }; type LayoutProps = { - arrange?: string; - edges: Array; - nodes: Array; openGroupIds: Array; -}; +} & StructureDataResponse; export const useGraphLayout = ({ arrange = "LR", @@ -271,7 +276,7 @@ export const useGraphLayout = ({ // 3. Flatten the nodes and edges for xyflow to actually render the graph const flattenedData = flattenGraph({ - children: data.children, + children: (data.children ?? []) as Array, }); // merge & dedupe edges diff --git a/airflow/ui/src/pages/DagsList/Dag/Tabs.tsx b/airflow/ui/src/pages/DagsList/Dag/Tabs.tsx index e3447051dbf25..5c2d44109ad3b 100644 --- a/airflow/ui/src/pages/DagsList/Dag/Tabs.tsx +++ b/airflow/ui/src/pages/DagsList/Dag/Tabs.tsx @@ -81,7 +81,8 @@ export const DagTabs = ({ dag }: { readonly dag?: DAGResponse }) => { diff --git a/airflow/utils/task_group.py b/airflow/utils/task_group.py index 3d3738f5a8d7d..2f850826ca990 100644 --- a/airflow/utils/task_group.py +++ b/airflow/utils/task_group.py @@ -123,7 +123,7 @@ def task_group_to_dict(task_item_or_group): "tooltip": task_group.tooltip, "is_mapped": is_mapped, "children": children, - "type": "task_group", + "type": "task", } diff --git a/tests/utils/test_task_group.py b/tests/utils/test_task_group.py index 7b5c960f9c298..bd336fa4c944a 100644 --- a/tests/utils/test_task_group.py +++ b/tests/utils/test_task_group.py @@ -190,17 +190,17 @@ def my_task(): {"id": "group234.group34.task4", "label": "task4", "type": "task"}, {"id": "group234.group34.downstream_join_id", "label": "", "type": "join"}, ], - "type": "task_group", + "type": "task", }, {"id": "group234.task2", "label": "task2", "type": "task"}, {"id": "group234.upstream_join_id", "label": "", "type": "join"}, ], - "type": "task_group", + "type": "task", }, {"id": "task1", "label": "task1", "type": "task"}, {"id": "task5", "label": "task5", "type": "task"}, ], - "type": "task_group", + "type": "task", }