diff --git a/airflow-core/src/airflow/ui/public/i18n/locales/en/components.json b/airflow-core/src/airflow/ui/public/i18n/locales/en/components.json index af84e92d17401..28ed7c7a15ae1 100644 --- a/airflow-core/src/airflow/ui/public/i18n/locales/en/components.json +++ b/airflow-core/src/airflow/ui/public/i18n/locales/en/components.json @@ -91,7 +91,8 @@ "otherDagRuns": "+Other Dag Runs", "taskCount_one": "{{count}} Task", "taskCount_other": "{{count}} Tasks", - "taskGroup": "Task Group" + "taskGroup": "Task Group", + "zoomToTask": "Zoom to selected task" }, "limitedList": "+{{count}} more", "limitedList.allItems": "All {{count}} items:", diff --git a/airflow-core/src/airflow/ui/src/components/Graph/Edge.tsx b/airflow-core/src/airflow/ui/src/components/Graph/Edge.tsx index 108c3c03496fa..244d440d10bdc 100644 --- a/airflow-core/src/airflow/ui/src/components/Graph/Edge.tsx +++ b/airflow-core/src/airflow/ui/src/components/Graph/Edge.tsx @@ -20,32 +20,32 @@ import { Text, useToken } from "@chakra-ui/react"; import { Group } from "@visx/group"; import { LinePath } from "@visx/shape"; import type { Edge as EdgeType } from "@xyflow/react"; +import { useNodesData } from "@xyflow/react"; import type { ElkPoint } from "elkjs"; import type { EdgeData } from "./reactflowUtils"; type Props = EdgeType; -const CustomEdge = ({ data }: Props) => { +const CustomEdge = ({ data, source, target }: Props) => { const [strokeColor, blueColor, dataEdgeColor] = useToken("colors", [ "border.inverted", "blue.500", "purple.500", ]); + // Read isSelected directly from the node store so that selection changes + // don't require the parent to rebuild and pass down a new edges array. + // useNodesData subscribes to data changes for these specific node IDs only. + const nodesData = useNodesData([source, target]); + const isSelected = nodesData.some((node) => Boolean(node.data.isSelected)); + if (data === undefined) { return undefined; } const { rest } = data; - // Determine edge color based on type - const getEdgeColor = () => { - if (rest.isSelected) { - return rest.edgeType === "data" ? dataEdgeColor : blueColor; - } - - return strokeColor; - }; + const edgeStrokeColor = isSelected ? (rest.edgeType === "data" ? dataEdgeColor : blueColor) : strokeColor; return ( <> @@ -73,9 +73,9 @@ const CustomEdge = ({ data }: Props) => { point.x} y={(point: ElkPoint) => point.y} /> diff --git a/airflow-core/src/airflow/ui/src/components/Graph/elkGraphUtils.ts b/airflow-core/src/airflow/ui/src/components/Graph/elkGraphUtils.ts new file mode 100644 index 0000000000000..0ab685cf37e84 --- /dev/null +++ b/airflow-core/src/airflow/ui/src/components/Graph/elkGraphUtils.ts @@ -0,0 +1,361 @@ +/*! + * 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 type { ElkNode, ElkExtendedEdge, ElkShape } from "elkjs"; + +import type { EdgeResponse, NodeResponse } from "openapi/requests/types.gen"; + +import type { Direction } from "./useGraphLayout"; + +// --------------------------------------------------------------------------- +// Types +// --------------------------------------------------------------------------- + +type EdgeLabel = { + height: number; + id: string; + text: string; + width: number; +}; + +export type FormattedEdge = { + id: string; + isSetupTeardown?: boolean; + labels?: Array; + parentNode?: string; +} & ElkExtendedEdge; + +export type FormattedNode = { + assetCondition?: NodeResponse["asset_condition_type"]; + childCount?: number; + edges?: Array; + isGroup: boolean; + isMapped?: boolean; + isOpen?: boolean; + setupTeardownType?: NodeResponse["setup_teardown_type"]; +} & ElkShape & + NodeResponse; + +export type GenerateElkProps = { + direction: Direction; + edges: Array; + font: string; + nodes: Array; + openGroupIds?: Array; +}; + +// --------------------------------------------------------------------------- +// Canvas singleton for text measurement +// --------------------------------------------------------------------------- + +// Single reusable canvas context — avoids allocating a new HTMLCanvasElement +// on every call (once per node per layout run). +const canvasContext = document.createElement("canvas").getContext("2d"); + +const getTextWidth = (text: string, font: string) => { + if (canvasContext) { + canvasContext.font = font; + + return Math.max(200, canvasContext.measureText(text).width); + } + + return Math.max(200, text.length * 9); +}; + +// --------------------------------------------------------------------------- +// Edge helpers +// --------------------------------------------------------------------------- + +const formatElkEdge = (edge: EdgeResponse, font: string, node?: NodeResponse): FormattedEdge => ({ + id: `${edge.source_id}-${edge.target_id}`, + isSetupTeardown: edge.is_setup_teardown === null ? undefined : edge.is_setup_teardown, + labels: + edge.label === undefined || edge.label === null + ? [] + : [ + { + height: 16, + id: edge.label, + text: edge.label, + width: getTextWidth(edge.label, font), + }, + ], + parentNode: node?.id, + sources: [edge.source_id], + targets: [edge.target_id], +}); + +/** + * Returns true when every child task that has at least one external connection + * shares exactly the same set of external sources AND the same set of external + * targets as every other externally-connected child. + * + * Example — a "cleanup" group where every task fans out from one upstream node + * and funnels into the same downstream node: + * + * upstream → T1 ─┐ + * upstream → T2 ─┼→ downstream + * upstream → T3 ─┘ + * + * Rendering N individual crossing edges adds visual noise without conveying + * any extra information beyond "the group connects upstream → downstream". + * When this returns true, the caller collapses those N edges to a single + * group-level edge while still rendering the children inside the group. + * + * Uses the original, unmodified edge list so that prior sibling group + * transformations do not affect the connectivity check. + */ +export const hasUniformExternalConnectivity = ( + childIdSet: Set, + edges: Array, +): boolean => { + const sourcesPerChild = new Map>(); + const targetsPerChild = new Map>(); + + for (const edge of edges) { + const sourceIsChild = childIdSet.has(edge.source_id); + const targetIsChild = childIdSet.has(edge.target_id); + + if (!sourceIsChild && targetIsChild) { + const existing = sourcesPerChild.get(edge.target_id) ?? new Set(); + + existing.add(edge.source_id); + sourcesPerChild.set(edge.target_id, existing); + } + + if (sourceIsChild && !targetIsChild) { + const existing = targetsPerChild.get(edge.source_id) ?? new Set(); + + existing.add(edge.target_id); + targetsPerChild.set(edge.source_id, existing); + } + } + + // Need at least 2 children with external connections on at least one side + // for the optimisation to be worthwhile. + if (sourcesPerChild.size < 2 && targetsPerChild.size < 2) { + return false; + } + + // Build the union of all external sources / targets across all children. + const allSources = new Set(); + const allTargets = new Set(); + + for (const sources of sourcesPerChild.values()) { + for (const source of sources) { + allSources.add(source); + } + } + for (const targets of targetsPerChild.values()) { + for (const target of targets) { + allTargets.add(target); + } + } + + // Every child's external sources must equal allSources (same size sufficient + // given allSources is already the union — a child with fewer differs in size). + for (const sources of sourcesPerChild.values()) { + if (sources.size !== allSources.size) { + return false; + } + } + for (const targets of targetsPerChild.values()) { + if (targets.size !== allTargets.size) { + return false; + } + } + + return true; +}; + +// --------------------------------------------------------------------------- +// Edge rewriting helper +// --------------------------------------------------------------------------- + +/** + * Given the current working edge list, drops purely-internal edges, rewrites + * crossing edges so both endpoints reference `groupId` instead of a child node, + * then deduplicates the result so N rewritten edges collapse to one per + * (source, target) pair. + */ +const rewriteGroupEdges = ( + edges: Array, + childIdSet: Set, + groupId: string, +): Array => { + const seen = new Set(); + + return edges + .filter((fe) => !(childIdSet.has(fe.source_id) && childIdSet.has(fe.target_id))) + .map((fe) => ({ + ...fe, + source_id: childIdSet.has(fe.source_id) ? groupId : fe.source_id, + target_id: childIdSet.has(fe.target_id) ? groupId : fe.target_id, + })) + .filter((fe) => { + const key = `${fe.source_id}-${fe.target_id}`; + + if (seen.has(key)) { + return false; + } + seen.add(key); + + return true; + }); +}; + +// --------------------------------------------------------------------------- +// Node helpers +// --------------------------------------------------------------------------- + +const getNestedChildIds = (children: Array): Array => { + const childIds: Array = []; + + for (const child of children) { + childIds.push(child.id); + if (child.children) { + childIds.push(...getNestedChildIds(child.children)); + } + } + + return childIds; +}; + +// --------------------------------------------------------------------------- +// Main graph builder +// --------------------------------------------------------------------------- + +export const generateElkGraph = ({ + direction, + edges: unformattedEdges, + font, + nodes, + openGroupIds, +}: GenerateElkProps): ElkNode => { + let filteredEdges = unformattedEdges; + + 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 === null || node.children === undefined ? [] : getNestedChildIds(node.children); + + if (isOpen && node.children !== null && node.children !== undefined) { + const childIdSet = new Set(childIds); + + // Process children first — their formatChildNode calls may modify filteredEdges + // (removing edges that belong to nested open groups). + const formattedChildren = node.children.map(formatChildNode); + + // If every externally-connected task shares the same upstream source(s) + // and downstream target(s), collapse N crossing edges to one group-level + // edge (same as a closed group) while keeping the children visible. + // Checked against unformattedEdges so prior sibling transforms don't interfere. + if (hasUniformExternalConnectivity(childIdSet, unformattedEdges)) { + filteredEdges = rewriteGroupEdges(filteredEdges, childIdSet, node.id); + } + + // Extract any remaining internal edges (both endpoints inside this group). + const internalEdges: Array = []; + + filteredEdges = filteredEdges.filter((edge) => { + if (childIdSet.has(edge.source_id) && childIdSet.has(edge.target_id)) { + internalEdges.push(formatElkEdge(edge, font, node)); + + return false; + } + + return true; + }); + + return { + ...node, + childCount, + children: formattedChildren, + edges: internalEdges, + id: node.id, + isGroup: true, + isOpen, + label: node.label, + layoutOptions: { + "elk.padding": "[top=80,left=15,bottom=15,right=15]", + ...(direction === "RIGHT" ? { "elk.portConstraints": "FIXED_SIDE" } : {}), + }, + }; + } + + if (!isOpen && node.children !== undefined) { + // Use a Set for O(1) membership checks — childIds.includes() would be + // O(n) per edge, turning the filter/map into O(n × E) for large groups. + filteredEdges = rewriteGroupEdges(filteredEdges, new Set(childIds), node.id); + } + + const label = `${node.label}${node.is_mapped ? "[1000]" : ""}${node.children ? ` + ${node.children.length} tasks` : ""}`; + let width = getTextWidth(label, font); + const hasStateBar = Boolean(node.is_mapped) || Boolean(node.children); + let height = hasStateBar ? 90 : 80; + + if (node.type === "join") { + width = 10; + height = 10; + } else if (node.type === "asset-condition") { + width = 30; + height = 30; + } + + return { + assetCondition: node.asset_condition_type, + childCount, + height, + id: node.id, + isGroup: Boolean(node.children), + isMapped: node.is_mapped === null ? undefined : node.is_mapped, + label: node.label, + layoutOptions: direction === "RIGHT" ? { "elk.portConstraints": "FIXED_SIDE" } : undefined, + operator: node.operator, + setupTeardownType: node.setup_teardown_type, + tooltip: node.tooltip, + type: node.type, + width, + }; + }; + + const children = nodes.map(formatChildNode); + const edges = filteredEdges.map((fe) => formatElkEdge(fe, font)); + + return { + children: children as Array, + edges, + id: "root", + layoutOptions: { + "elk.core.options.EdgeLabelPlacement": "CENTER", + "elk.direction": direction, + // SIMPLE placement is a single-pass algorithm — much faster than the + // default BRANDES_KOEPF four-pass approach with acceptable quality for DAGs. + "elk.layered.nodePlacement.strategy": "SIMPLE", + // Crossing minimisation thoroughness (default 7) controls how many random + // sweeps are attempted. Drop to 3 for large graphs where the extra passes + // rarely pay off and add noticeable layout latency. + ...(filteredEdges.length > 100 ? { "elk.layered.thoroughness": "3" } : {}), + // hierarchyHandling is only needed when open groups create cross-hierarchy + // edges. For flat DAGs (no open groups) omitting it simplifies ELK's task. + ...(Boolean(openGroupIds?.length) ? { hierarchyHandling: "INCLUDE_CHILDREN" } : {}), + "spacing.edgeLabel": "10.0", + }, + }; +}; diff --git a/airflow-core/src/airflow/ui/src/components/Graph/reactflowUtils.ts b/airflow-core/src/airflow/ui/src/components/Graph/reactflowUtils.ts index d73eec056ee7b..9c6dcbdf13379 100644 --- a/airflow-core/src/airflow/ui/src/components/Graph/reactflowUtils.ts +++ b/airflow-core/src/airflow/ui/src/components/Graph/reactflowUtils.ts @@ -50,7 +50,9 @@ type FlattenNodesProps = { parent?: NodeType; }; -// Generate a flattened list of nodes for react-flow to render +// Generate a flattened list of nodes for react-flow to render. +// Uses push() throughout to avoid the O(n²) spread-in-loop pattern that the +// previous implementation had (edges = [...edges, ...newEdges] per node). export const flattenGraph = ({ children, level = 0, @@ -59,8 +61,8 @@ export const flattenGraph = ({ edges: Array; nodes: Array; } => { - let nodes: Array = []; - let edges: Array = []; + const nodes: Array = []; + const edges: Array = []; if (!children) { return { edges, nodes }; @@ -74,18 +76,16 @@ export const flattenGraph = ({ data: { ...node, depth: level }, height: node.height, id: node.id, - position: { - x, - y, - }, + position: { x, y }, type: node.type, width: node.width, ...parentNode, } satisfies NodeType; - edges = [ - ...edges, - ...(node.edges ?? []).map((edge) => ({ + nodes.push(newNode); + + for (const edge of node.edges ?? []) { + edges.push({ ...edge, labels: edge.labels?.map((label) => ({ ...label, @@ -94,24 +94,12 @@ export const flattenGraph = ({ })), sections: edge.sections?.map((section) => ({ ...section, - // eslint-disable-next-line max-nested-callbacks - bendPoints: section.bendPoints?.map((bp) => ({ - x: bp.x + x, - y: bp.y + y, - })), - endPoint: { - x: section.endPoint.x + x, - y: section.endPoint.y + y, - }, - startPoint: { - x: section.startPoint.x + x, - y: section.startPoint.y + y, - }, + bendPoints: section.bendPoints?.map((bp) => ({ x: bp.x + x, y: bp.y + y })), + endPoint: { x: section.endPoint.x + x, y: section.endPoint.y + y }, + startPoint: { x: section.startPoint.x + x, y: section.startPoint.y + y }, })), - })), - ]; - - nodes.push(newNode); + }); + } if (node.children) { const { edges: childEdges, nodes: childNodes } = flattenGraph({ @@ -120,15 +108,12 @@ export const flattenGraph = ({ parent: newNode, }); - nodes = [...nodes, ...childNodes]; - edges = [...edges, ...childEdges]; + nodes.push(...childNodes); + edges.push(...childEdges); } }); - return { - edges, - nodes, - }; + return { edges, nodes }; }; type Edge = { diff --git a/airflow-core/src/airflow/ui/src/components/Graph/useGraphLayout.ts b/airflow-core/src/airflow/ui/src/components/Graph/useGraphLayout.ts index a340ddd47a67f..9ffeb0fb673fb 100644 --- a/airflow-core/src/airflow/ui/src/components/Graph/useGraphLayout.ts +++ b/airflow-core/src/airflow/ui/src/components/Graph/useGraphLayout.ts @@ -18,13 +18,31 @@ */ import { createListCollection } from "@chakra-ui/react"; import { useQuery } from "@tanstack/react-query"; -import ELK, { type ElkNode, type ElkExtendedEdge, type ElkShape } from "elkjs"; +import ELK, { type ElkNode } from "elkjs"; +// ?raw imports the file content as a plain string without any transformation. +// We create a blob: URL from it so the Worker is always same-origin to the +// page — avoiding the cross-origin SecurityError that occurs in Airflow's dev +// setup where Vite (5173) and Flask (28080) run on different ports, and all +// URL-based worker approaches (?worker, ?worker&inline, new URL()) resolve to +// the Vite origin which browsers reject for Workers. +import ElkWorkerSource from "elkjs/lib/elk-worker.min.js?raw"; import type { TFunction } from "i18next"; -import type { EdgeResponse, NodeResponse, StructureDataResponse } from "openapi/requests/types.gen"; +import type { NodeResponse, StructureDataResponse } from "openapi/requests/types.gen"; +import { generateElkGraph } from "./elkGraphUtils"; import { flattenGraph, formatFlowEdges } from "./reactflowUtils"; +// Blob URL created once at module load. `type: "classic"` preserves the +// original CJS environment detection in elk-worker: as a classic script +// `typeof module === "undefined"`, so the worker sets self.onmessage rather +// than exporting a FakeWorker. +const elkWorkerBlobUrl = URL.createObjectURL(new Blob([ElkWorkerSource], { type: "application/javascript" })); + +const elk = new ELK({ + workerFactory: () => new Worker(elkWorkerBlobUrl, { type: "classic" }), +}); + export type Direction = "DOWN" | "LEFT" | "RIGHT" | "UP"; export const directionOptions = (translate: TFunction) => createListCollection({ @@ -36,212 +54,8 @@ export const directionOptions = (translate: TFunction) => ], }); -type EdgeLabel = { - height: number; - id: string; - text: string; - width: number; -}; - -type FormattedNode = { - assetCondition?: NodeResponse["asset_condition_type"]; - childCount?: number; - edges?: Array; - isGroup: boolean; - isMapped?: boolean; - isOpen?: boolean; - setupTeardownType?: NodeResponse["setup_teardown_type"]; -} & ElkShape & - NodeResponse; - -type FormattedEdge = { - id: string; - isSetupTeardown?: boolean; - labels?: Array; - parentNode?: string; -} & ElkExtendedEdge; - export type LayoutNode = ElkNode & NodeResponse; -// Take text and font to calculate how long each node should be -const getTextWidth = (text: string, font: string) => { - const context = document.createElement("canvas").getContext("2d"); - - if (context) { - context.font = font; - const metrics = context.measureText(text); - - return metrics.width > 200 ? metrics.width : 200; - } - - const length = text.length * 9; - - return length > 200 ? length : 200; -}; - -const formatElkEdge = (edge: EdgeResponse, font: string, node?: NodeResponse): FormattedEdge => ({ - id: `${edge.source_id}-${edge.target_id}`, - isSetupTeardown: edge.is_setup_teardown === null ? undefined : edge.is_setup_teardown, - // isSourceAsset: e.isSourceAsset, - labels: - edge.label === undefined || edge.label === null - ? [] - : [ - { - height: 16, - id: edge.label, - text: edge.label, - width: getTextWidth(edge.label, font), - }, - ], - parentNode: node?.id, - sources: [edge.source_id], - targets: [edge.target_id], -}); - -const getNestedChildIds = (children: Array) => { - let childIds: Array = []; - - children.forEach((child) => { - childIds.push(child.id); - if (child.children) { - const nestedChildIds = getNestedChildIds(child.children); - - childIds = [...childIds, ...nestedChildIds]; - } - }); - - return childIds; -}; - -type GenerateElkProps = { - direction: Direction; - edges: Array; - font: string; - nodes: Array; - openGroupIds?: Array; -}; - -const generateElkGraph = ({ - direction, - edges: unformattedEdges, - font, - nodes, - openGroupIds, -}: GenerateElkProps): ElkNode => { - const closedGroupIds: Array = []; - let filteredEdges = unformattedEdges; - - 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 === null || node.children === undefined ? [] : getNestedChildIds(node.children); - - if (isOpen && node.children !== null && node.children !== undefined) { - return { - ...node, - childCount, - children: node.children.map(formatChildNode), - edges: filteredEdges - .filter((edge) => { - if (childIds.includes(edge.source_id) && childIds.includes(edge.target_id)) { - // Remove edge from array when we add it here - filteredEdges = filteredEdges.filter( - (fe) => !(fe.source_id === edge.source_id && fe.target_id === edge.target_id), - ); - - return true; - } - - return false; - }) - .map((edge) => formatElkEdge(edge, font, node)), - id: node.id, - isGroup: true, - isOpen, - label: node.label, - layoutOptions: { - "elk.padding": "[top=80,left=15,bottom=15,right=15]", - ...(direction === "RIGHT" ? { "elk.portConstraints": "FIXED_SIDE" } : {}), - }, - }; - } - - if (!Boolean(isOpen) && node.children !== undefined) { - const seenEdges = new Set(); - - filteredEdges = filteredEdges - // Filter out internal group edges - .filter((fe) => !(childIds.includes(fe.source_id) && childIds.includes(fe.target_id))) - // For external group edges, point to the group itself instead of a child node - .map((fe) => ({ - ...fe, - source_id: childIds.includes(fe.source_id) ? node.id : fe.source_id, - target_id: childIds.includes(fe.target_id) ? node.id : fe.target_id, - })) - // Deduplicate edges based on source_id and target_id composite - .filter((fe) => { - const edgeKey = `${fe.source_id}-${fe.target_id}`; - - if (seenEdges.has(edgeKey)) { - return false; - } - seenEdges.add(edgeKey); - - return true; - }); - closedGroupIds.push(node.id); - } - - const label = `${node.label}${node.is_mapped ? "[1000]" : ""}${node.children ? ` + ${node.children.length} tasks` : ""}`; - let width = getTextWidth(label, font); - const hasStateBar = Boolean(node.is_mapped) || Boolean(node.children); - let height = hasStateBar ? 90 : 80; - - if (node.type === "join") { - width = 10; - height = 10; - } else if (node.type === "asset-condition") { - width = 30; - height = 30; - } - - return { - assetCondition: node.asset_condition_type, - childCount, - height, - id: node.id, - isGroup: Boolean(node.children), - isMapped: node.is_mapped === null ? undefined : node.is_mapped, - label: node.label, - layoutOptions: direction === "RIGHT" ? { "elk.portConstraints": "FIXED_SIDE" } : undefined, - operator: node.operator, - setupTeardownType: node.setup_teardown_type, - tooltip: node.tooltip, - type: node.type, - width, - }; - }; - - const children = nodes.map(formatChildNode); - - const edges = filteredEdges.map((fe) => formatElkEdge(fe, font)); - - return { - children: children as Array, - edges, - id: "root", - layoutOptions: { - "elk.core.options.EdgeLabelPlacement": "CENTER", - "elk.direction": direction, - hierarchyHandling: "INCLUDE_CHILDREN", - "spacing.edgeLabel": "10.0", - }, - }; -}; - type LayoutProps = { direction: Direction; openGroupIds: Array; @@ -258,7 +72,6 @@ export const useGraphLayout = ({ useQuery({ queryFn: async () => { const font = `bold 18px ${globalThis.getComputedStyle(document.body).fontFamily}`; - const elk = new ELK(); // 1. Format graph data to pass for elk to process const graph = generateElkGraph({ @@ -277,10 +90,18 @@ export const useGraphLayout = ({ children: (data.children ?? []) as Array, }); - // merge & dedupe edges - const flatEdges = [...(data.edges ?? []), ...flattenedData.edges].filter( - (value, index, self) => index === self.findIndex((edge) => edge.id === value.id), - ); + // merge & dedupe edges — O(n) via Map (first occurrence wins) rather than + // O(n²) findIndex. Root-level edges from ELK come first; child edges from + // flattenedData are skipped when the same id is already present. + const seenEdgeIds = new Set(); + const flatEdges = [...(data.edges ?? []), ...flattenedData.edges].filter((edge) => { + if (seenEdgeIds.has(edge.id)) { + return false; + } + seenEdgeIds.add(edge.id); + + return true; + }); const formattedEdges = formatFlowEdges({ edges: flatEdges }); diff --git a/airflow-core/src/airflow/ui/src/layouts/Details/Graph/Graph.tsx b/airflow-core/src/airflow/ui/src/layouts/Details/Graph/Graph.tsx index 6e3fe7940218a..70fcf23ed9db2 100644 --- a/airflow-core/src/airflow/ui/src/layouts/Details/Graph/Graph.tsx +++ b/airflow-core/src/airflow/ui/src/layouts/Details/Graph/Graph.tsx @@ -16,10 +16,20 @@ * specific language governing permissions and limitations * under the License. */ -import { useToken } from "@chakra-ui/react"; -import { ReactFlow, Controls, Background, MiniMap, type Node as ReactFlowNode } from "@xyflow/react"; +import { Box, Spinner, useToken } from "@chakra-ui/react"; +import { + ReactFlow, + Controls, + ControlButton, + Background, + MiniMap, + useReactFlow, + type Node as ReactFlowNode, +} from "@xyflow/react"; import "@xyflow/react/dist/style.css"; import { useEffect } from "react"; +import { useTranslation } from "react-i18next"; +import { MdCenterFocusStrong } from "react-icons/md"; import { useParams, useSearchParams } from "react-router-dom"; import { useLocalStorage } from "usehooks-ts"; @@ -37,6 +47,47 @@ import { useDependencyGraph } from "src/queries/useDependencyGraph"; import { useGridTiSummariesStream } from "src/queries/useGridTISummaries.ts"; import { getReactFlowThemeStyle } from "src/theme"; +// Hoisted to module scope so ReactFlow receives a stable reference and skips +// its internal shallow-equality check on every render. +const defaultEdgeOptions = { zIndex: 1 }; + +// Fits the viewport whenever a new layout is committed. Must live inside +// to call useReactFlow(). Using layoutData as the dep means it +// only fires when ELK produces a new layout, not on task-instance updates or +// selection changes (unlike the `fitView` prop, which runs on every re-mount). +const FitViewOnLayout = ({ layoutData }: { readonly layoutData: object | undefined }) => { + const { fitView } = useReactFlow(); + + useEffect(() => { + if (layoutData !== undefined) { + void fitView({ padding: 0.1 }); + } + }, [fitView, layoutData]); + + return null; +}; + +const GraphControls = ({ selectedNodeId }: { readonly selectedNodeId?: string }) => { + const { t: translate } = useTranslation("components"); + const { fitView } = useReactFlow(); + + return ( + + {selectedNodeId === undefined ? undefined : ( + { + void fitView({ duration: 500, nodes: [{ id: selectedNodeId }], padding: 0.5 }); + }} + title={translate("graph.zoomToTask")} + > + + + )} + + ); +}; + const nodeColor = ( { data: { depth, height, isOpen, taskInstance, width }, type }: ReactFlowNode, evenColor?: string, @@ -122,22 +173,24 @@ export const Graph = () => { const dagDepEdges = dependencies === "all" ? dagDependencies.edges : []; const dagDepNodes = dependencies === "all" ? dagDependencies.nodes : []; - const { data } = useGraphLayout({ + const layoutEdges = [...graphData.edges, ...dagDepEdges]; + const layoutNodes = dagDepNodes.length + ? dagDepNodes.map((node) => (node.id === `dag:${dagId}` ? { ...node, children: graphData.nodes } : node)) + : graphData.nodes; + const layoutOpenGroupIds = [...openGroupIds, ...(dependencies === "all" ? [`dag:${dagId}`] : [])]; + + const { data, isPending } = useGraphLayout({ direction, - edges: [...graphData.edges, ...dagDepEdges], - nodes: dagDepNodes.length - ? dagDepNodes.map((node) => - node.id === `dag:${dagId}` ? { ...node, children: graphData.nodes } : node, - ) - : graphData.nodes, - openGroupIds: [...openGroupIds, ...(dependencies === "all" ? [`dag:${dagId}`] : [])], + edges: layoutEdges, + nodes: layoutNodes, + openGroupIds: layoutOpenGroupIds, versionNumber: selectedVersion, }); const { summariesByRunId } = useGridTiSummariesStream({ dagId, runIds: runId ? [runId] : [] }); const gridTISummaries = runId ? summariesByRunId.get(runId) : undefined; - // Add task instances to the node data but without having to recalculate how the graph is laid out + // Add task instances and selection state to node data without recalculating layout. const nodes = data?.nodes.map((node) => { const taskInstance = gridTISummaries?.task_instances.find((ti) => ti.task_id === node.id); @@ -151,58 +204,83 @@ export const Graph = () => { }; }); + // isSelected is intentionally absent here — Edge.tsx reads it directly from + // the node store via useNodesData, so the edges array stays stable when only + // the selected task changes, avoiding a full edge reconciliation pass. const edges = (data?.edges ?? []).map((edge) => ({ ...edge, data: { ...edge.data, rest: { ...edge.data?.rest, - isSelected: - taskId === edge.source || - taskId === edge.target || - groupId === edge.source || - groupId === edge.target || - edge.source === `dag:${dagId}` || - edge.target === `dag:${dagId}`, }, }, })); + const selectedNodeId = taskId ?? groupId; + return ( - - - - ) => - nodeColor( - node, - colorMode === "dark" ? evenDark : evenLight, - colorMode === "dark" ? oddDark : oddLight, - ) - } - nodeStrokeColor={(node: ReactFlowNode) => - node.data.isSelected && selectedColor !== undefined ? selectedColor : "" - } - nodeStrokeWidth={15} - pannable - style={{ height: 150, width: 200 }} - zoomable - /> - - + + {isPending ? ( + + + + ) : undefined} + + + {/* Fit the viewport after each new ELK layout instead of using the + fitView prop, which re-fires on every re-mount even when nodes are + served from the React Query cache. */} + + + {/* Hide the MiniMap for large graphs — it processes all nodes even when + onlyRenderVisibleElements is set, adding meaningful paint cost with + little benefit at 500+ nodes where the map is a near-solid blob. */} + {data !== undefined && data.nodes.length <= 500 ? ( + ) => + nodeColor( + node, + colorMode === "dark" ? evenDark : evenLight, + colorMode === "dark" ? oddDark : oddLight, + ) + } + nodeStrokeColor={(node: ReactFlowNode) => + node.data.isSelected && selectedColor !== undefined ? selectedColor : "" + } + nodeStrokeWidth={15} + pannable + style={{ height: 150, width: 200 }} + zoomable + /> + ) : undefined} + + + ); }; diff --git a/airflow-core/src/airflow/ui/testsSetup.ts b/airflow-core/src/airflow/ui/testsSetup.ts index 91d3533b0aa2c..1bb26b7d5ba7a 100644 --- a/airflow-core/src/airflow/ui/testsSetup.ts +++ b/airflow-core/src/airflow/ui/testsSetup.ts @@ -24,6 +24,19 @@ import { beforeEach, beforeAll, afterAll, afterEach, vi } from "vitest"; import { handlers } from "src/mocks/handlers"; +// ELK creates a Worker at module-load time, which is not available in +// happy-dom. Mock useGraphLayout so the Worker is never constructed. Any +// test that specifically exercises graph layout should override this mock. +vi.mock("src/components/Graph/useGraphLayout", () => ({ + directionOptions: () => ({ items: [] }), + useGraphLayout: vi.fn().mockReturnValue({ data: undefined, isPending: false }), +})); + +// Render nothing for the two graph-rendering components so tests that render +// full pages don't need to stub layout data. +vi.mock("src/layouts/Details/Graph/Graph", () => ({ Graph: () => null })); +vi.mock("src/pages/Asset/AssetGraph", () => ({ AssetGraph: () => null })); + // Mock Chart.js to prevent DOM access errors during test cleanup vi.mock("react-chartjs-2", () => ({ Bar: vi.fn(() => {