This is an automated email from the ASF dual-hosted git repository.
ephraimanierobi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 6547b8182f Improve graph nesting logic (#33421)
6547b8182f is described below
commit 6547b8182f056b6d29445dec50cce3920d36f762
Author: Brent Bovenzi <br...@astronomer.io>
AuthorDate: Wed Aug 16 00:38:46 2023 -0400
Improve graph nesting logic (#33421)
---
airflow/www/static/js/components/Graph/Edge.tsx | 18 +----
airflow/www/static/js/dag/details/graph/Node.tsx | 1 -
airflow/www/static/js/dag/details/graph/index.tsx | 32 ++++-----
airflow/www/static/js/dag/details/graph/utils.ts | 81 ++++++++---------------
airflow/www/static/js/types/index.ts | 41 +++++++++++-
airflow/www/static/js/utils/graph.ts | 62 ++++++++++-------
6 files changed, 124 insertions(+), 111 deletions(-)
diff --git a/airflow/www/static/js/components/Graph/Edge.tsx
b/airflow/www/static/js/components/Graph/Edge.tsx
index 910d2d3e43..3d9e144795 100644
--- a/airflow/www/static/js/components/Graph/Edge.tsx
+++ b/airflow/www/static/js/components/Graph/Edge.tsx
@@ -19,25 +19,13 @@
import React from "react";
import { Text, useTheme } from "@chakra-ui/react";
-import type { ElkEdgeSection, ElkLabel, ElkPoint, LayoutOptions } from "elkjs";
import { Group } from "@visx/group";
import { LinePath } from "@visx/shape";
+import type { EdgeData } from "src/types";
-interface EdgeProps {
- data?: {
- rest: {
- isSelected: boolean;
- sources: string[];
- targets: string[];
- sections: ElkEdgeSection[];
- junctionPoints?: ElkPoint[];
- id: string;
- labels?: ElkLabel[];
- layoutOptions?: LayoutOptions;
- isSetupTeardown?: boolean;
- };
- };
+export interface EdgeProps {
+ data?: EdgeData;
}
const CustomEdge = ({ data }: EdgeProps) => {
diff --git a/airflow/www/static/js/dag/details/graph/Node.tsx
b/airflow/www/static/js/dag/details/graph/Node.tsx
index f2d6f92c09..6ae8489f13 100644
--- a/airflow/www/static/js/dag/details/graph/Node.tsx
+++ b/airflow/www/static/js/dag/details/graph/Node.tsx
@@ -44,7 +44,6 @@ export interface CustomNodeProps {
isOpen?: boolean;
isActive?: boolean;
setupTeardownType?: "setup" | "teardown";
- fullParentNode?: string;
labelStyle?: string;
style?: string;
}
diff --git a/airflow/www/static/js/dag/details/graph/index.tsx
b/airflow/www/static/js/dag/details/graph/index.tsx
index 1bdfe5ece7..29e07d65aa 100644
--- a/airflow/www/static/js/dag/details/graph/index.tsx
+++ b/airflow/www/static/js/dag/details/graph/index.tsx
@@ -24,7 +24,6 @@ import ReactFlow, {
Controls,
Background,
MiniMap,
- Node as ReactFlowNode,
useReactFlow,
Panel,
} from "reactflow";
@@ -35,7 +34,7 @@ import { useOffsetTop } from "src/utils";
import { useGraphLayout } from "src/utils/graph";
import Edge from "src/components/Graph/Edge";
-import Node, { CustomNodeProps } from "./Node";
+import Node from "./Node";
import { buildEdges, nodeStrokeColor, nodeColor, flattenNodes } from "./utils";
const nodeTypes = { custom: Node };
@@ -72,19 +71,17 @@ const Graph = ({ openGroupIds, onToggleGroups,
hoveredTaskState }: Props) => {
const latestDagRunId = dagRuns[dagRuns.length - 1]?.runId;
const offsetTop = useOffsetTop(graphRef);
- const nodes: ReactFlowNode<CustomNodeProps>[] = useMemo(
+ const { nodes, edges: nodeEdges } = useMemo(
() =>
- graphData?.children
- ? flattenNodes({
- children: graphData.children,
- selected,
- openGroupIds,
- onToggleGroups,
- latestDagRunId,
- groups,
- hoveredTaskState,
- })
- : [],
+ flattenNodes({
+ children: graphData?.children,
+ selected,
+ openGroupIds,
+ onToggleGroups,
+ latestDagRunId,
+ groups,
+ hoveredTaskState,
+ }),
[
graphData?.children,
selected,
@@ -110,8 +107,13 @@ const Graph = ({ openGroupIds, onToggleGroups,
hoveredTaskState }: Props) => {
setHasRendered(true);
}, [fitView, hasRendered, selected.taskId, getZoom]);
+ // merge & dedupe edges
+ const flatEdges = [...(graphData?.edges || []), ...(nodeEdges || [])].filter(
+ (value, index, self) => index === self.findIndex((t) => t.id === value.id)
+ );
+
const edges = buildEdges({
- edges: graphData?.edges,
+ edges: flatEdges,
nodes,
selectedTaskId: selected.taskId,
});
diff --git a/airflow/www/static/js/dag/details/graph/utils.ts
b/airflow/www/static/js/dag/details/graph/utils.ts
index aef01e5d53..3d85053619 100644
--- a/airflow/www/static/js/dag/details/graph/utils.ts
+++ b/airflow/www/static/js/dag/details/graph/utils.ts
@@ -28,7 +28,7 @@ import type { Task, TaskInstance, NodeType } from "src/types";
import type { CustomNodeProps } from "./Node";
interface FlattenNodesProps {
- children: NodeType[];
+ children?: NodeType[];
selected: SelectionProps;
groups: Task;
latestDagRunId: string;
@@ -50,15 +50,10 @@ export const flattenNodes = ({
hoveredTaskState,
}: FlattenNodesProps) => {
let nodes: ReactFlowNode<CustomNodeProps>[] = [];
+ let edges: ElkExtendedEdge[] = [];
+ if (!children) return { nodes, edges };
const parentNode = parent ? { parentNode: parent.id } : undefined;
- let fullParentNode: string | undefined;
- if (parent) {
- fullParentNode =
- parent.parentNode && !parent.id.startsWith(parent.parentNode)
- ? `${parent.parentNode}.${parent.id}`
- : parent.id;
- }
children.forEach((node) => {
let instance: TaskInstance | undefined;
const group = getTask({ taskId: node.id, task: groups });
@@ -90,7 +85,6 @@ export const flattenNodes = ({
}
onToggleGroups(newGroupIds);
},
- fullParentNode,
...node.value,
},
type: "custom",
@@ -105,10 +99,14 @@ export const flattenNodes = ({
...parentNode,
};
+ if (node.edges) {
+ edges = [...edges, ...node.edges];
+ }
+
nodes.push(newNode);
if (node.children) {
- const childNodes = flattenNodes({
+ const { nodes: childNodes, edges: childEdges } = flattenNodes({
children: node.children,
selected,
groups,
@@ -119,9 +117,13 @@ export const flattenNodes = ({
hoveredTaskState,
});
nodes = [...nodes, ...childNodes];
+ edges = [...edges, ...childEdges];
}
});
- return nodes;
+ return {
+ nodes,
+ edges,
+ };
};
export const nodeColor = ({
@@ -143,8 +145,12 @@ export const nodeStrokeColor = (
colors: Record<string, string>
) => (isSelected ? colors.blue[500] : "");
+interface Edge extends ElkExtendedEdge {
+ parentNode?: string;
+}
+
interface BuildEdgesProps {
- edges?: ElkExtendedEdge[];
+ edges?: Edge[];
nodes: ReactFlowNode<CustomNodeProps>[];
selectedTaskId?: string | null;
}
@@ -164,57 +170,24 @@ export const buildEdges = ({
type: "custom",
}))
.map((e) => {
- const sourceNode = nodes.find((n) => n.id === e.source);
- const targetNode = nodes.find((n) => n.id === e.target);
-
- // Before finding the depth of the edge, append the parentNode in case
prefix_group_id is false
- const sourceIds = (
- sourceNode?.data.fullParentNode &&
- e.source.startsWith(sourceNode.data.fullParentNode)
- ? e.source
- : `${sourceNode?.data.fullParentNode}.${e.source}`
- ).split(".");
- const targetIds = (
- targetNode?.data.fullParentNode &&
- e.target.startsWith(targetNode.data.fullParentNode)
- ? e.target
- : `${targetNode?.data.fullParentNode}.${e.target}`
- ).split(".");
-
const isSelected =
selectedTaskId &&
(e.source === selectedTaskId || e.target === selectedTaskId);
- if (
- sourceIds.length === targetIds.length &&
- sourceIds[0] === targetIds[0]
- ) {
- let parentIds =
- sourceIds.length > targetIds.length ? sourceIds : targetIds;
-
- if (e.target.endsWith("_join_id") && e.source.endsWith("_join_id")) {
- /** edges between join ids are positioned absolutely,
- * other edges are positioned relative to their parent */
- parentIds = [];
- } else {
- // remove the last node
- parentIds.pop();
- }
- let parentX = 0;
- let parentY = 0;
-
- nodes
- .filter((n) => parentIds.some((p) => p === n.data.label))
- .forEach((p) => {
- parentX += p.position.x;
- parentY += p.position.y;
- });
-
+ if (e.data.rest?.parentNode) {
+ const parentNode = nodes.find((n) => n.id === e.data.rest.parentNode);
+ const parentX =
+ parentNode?.positionAbsolute?.x || parentNode?.position.x || 0;
+ const parentY =
+ parentNode?.positionAbsolute?.y || parentNode?.position.y || 0;
return {
...e,
data: {
rest: {
...e.data.rest,
+ labels: e.data.rest.labels?.map((l) =>
+ l.x && l.y ? { ...l, x: l.x + parentX, y: l.y + parentY } : l
+ ),
isSelected,
sections: e.data.rest.sections.map((s) => ({
...s,
diff --git a/airflow/www/static/js/types/index.ts
b/airflow/www/static/js/types/index.ts
index 478b04d4f6..fc3634d9eb 100644
--- a/airflow/www/static/js/types/index.ts
+++ b/airflow/www/static/js/types/index.ts
@@ -18,7 +18,14 @@
*/
import type { CamelCase } from "type-fest";
-import type { ElkShape } from "elkjs";
+import type {
+ ElkExtendedEdge,
+ ElkEdgeSection,
+ ElkLabel,
+ ElkPoint,
+ ElkShape,
+ LayoutOptions,
+} from "elkjs";
import type * as API from "./api-generated";
type RunState = "success" | "running" | "queued" | "failed";
@@ -109,6 +116,20 @@ type RunOrdering = (
| "dataIntervalEnd"
)[];
+export interface MidEdge {
+ id: string;
+ sources: string[];
+ targets: string[];
+ isSetupTeardown: boolean | undefined;
+ parentNode: string | undefined;
+ labels: {
+ id: string;
+ text: string;
+ height: number;
+ width: number;
+ }[];
+}
+
interface DepNode {
id: string;
value: {
@@ -125,6 +146,7 @@ interface DepNode {
setupTeardownType?: "setup" | "teardown";
};
children?: DepNode[];
+ edges?: MidEdge[];
}
interface DepEdge {
@@ -132,9 +154,25 @@ interface DepEdge {
target: string;
}
+export interface EdgeData {
+ rest: {
+ isSelected: boolean;
+ sources: string[];
+ targets: string[];
+ sections: ElkEdgeSection[];
+ junctionPoints?: ElkPoint[];
+ id: string;
+ labels?: ElkLabel[];
+ layoutOptions?: LayoutOptions;
+ isSetupTeardown?: boolean;
+ parentNode?: string;
+ };
+}
+
export interface NodeType extends ElkShape {
value: DepNode["value"];
children?: NodeType[];
+ edges?: ElkExtendedEdge[];
}
export interface WebserverEdge {
@@ -142,6 +180,7 @@ export interface WebserverEdge {
sourceId: string;
targetId: string;
isSetupTeardown?: boolean;
+ parentNode?: string;
}
interface DatasetListItem extends API.Dataset {
diff --git a/airflow/www/static/js/utils/graph.ts
b/airflow/www/static/js/utils/graph.ts
index cd2af9cf4c..08d5415833 100644
--- a/airflow/www/static/js/utils/graph.ts
+++ b/airflow/www/static/js/utils/graph.ts
@@ -68,6 +68,24 @@ const getDirection = (arrange: string) => {
}
};
+const formatEdge = (e: WebserverEdge, font: string, node?: DepNode) => ({
+ id: `${e.sourceId}-${e.targetId}`,
+ sources: [e.sourceId],
+ targets: [e.targetId],
+ isSetupTeardown: e.isSetupTeardown,
+ parentNode: node?.id,
+ labels: e.label
+ ? [
+ {
+ id: e.label,
+ text: e.label,
+ height: 16,
+ width: getTextWidth(e.label, font),
+ },
+ ]
+ : [],
+});
+
const generateGraph = ({
nodes,
edges: unformattedEdges,
@@ -102,6 +120,7 @@ const generateGraph = ({
const isOpen = openGroupIds?.includes(value.label);
const childCount =
children?.filter((c: DepNode) => !c.id.includes("join_id")).length || 0;
+ const childIds = children?.length ? getNestedChildIds(children) : [];
if (isOpen && children?.length) {
return {
...node,
@@ -116,11 +135,26 @@ const generateGraph = ({
"elk.padding": "[top=60,left=10,bottom=10,right=10]",
},
children: children.map(formatChildNode),
+ edges: filteredEdges
+ .filter((e) => {
+ if (
+ childIds.indexOf(e.sourceId) > -1 &&
+ childIds.indexOf(e.targetId) > -1
+ ) {
+ // Remove edge from array when we add it here
+ filteredEdges = filteredEdges.filter(
+ (fe) =>
+ !(fe.sourceId === e.sourceId && fe.targetId === e.targetId)
+ );
+ return true;
+ }
+ return false;
+ })
+ .map((e) => formatEdge(e, font, node)),
};
}
const isJoinNode = id.includes("join_id");
if (!isOpen && children?.length) {
- const childIds = getNestedChildIds(children);
filteredEdges = filteredEdges
// Filter out internal group edges
.filter(
@@ -153,30 +187,7 @@ const generateGraph = ({
const children = nodes.map(formatChildNode);
- const edges = filteredEdges
- .filter(
- (value, index, self) =>
- index ===
- self.findIndex(
- (t) => t.sourceId === value.sourceId && t.targetId === value.targetId
- )
- )
- .map((e) => ({
- id: `${e.sourceId}-${e.targetId}`,
- sources: [e.sourceId],
- targets: [e.targetId],
- isSetupTeardown: e.isSetupTeardown,
- labels: e.label
- ? [
- {
- id: e.label,
- text: e.label,
- height: 16,
- width: getTextWidth(e.label, font),
- },
- ]
- : [],
- }));
+ const edges = filteredEdges.map((e) => formatEdge(e, font));
return {
id: "root",
@@ -184,6 +195,7 @@ const generateGraph = ({
hierarchyHandling: "INCLUDE_CHILDREN",
"elk.direction": getDirection(arrange),
"spacing.edgeLabel": "10.0",
+ "elk.core.options.EdgeLabelPlacement": "CENTER",
},
children,
edges,
