[DOP-23920] Highlight related nodes on double click

dolfinus · dolfinus · commit db24a3e12b24 · 2025-02-27T16:39:29.000+03:00
diff --git a/src/components/lineage/LineageGraph.tsx b/src/components/lineage/LineageGraph.tsx
@@ -11,6 +11,7 @@ import {
 import { DatasetNode, JobNode, RunNode, OperationNode } from "./nodes";
 import { useEffect } from "react";
 import { BaseEdge, IOEdge } from "./edges";
+import { useLineageSelection } from "./selection";
 
 export const MIN_ZOOM_VALUE = 0.1;
 export const MAX_ZOOM_VALUE = 2.5;
@@ -29,6 +30,7 @@ const nodeTypes = {
 
 const LineageGraph = (props: ReactFlowProps) => {
     const { fitView } = useReactFlow();
+    const lineageSelection = useLineageSelection();
 
     const nodesInitialized = useNodesInitialized();
 
@@ -38,6 +40,9 @@ const LineageGraph = (props: ReactFlowProps) => {
 
     return (
         <ReactFlow
+            className={
+                lineageSelection.hideNonSelected ? "hideNonSelected" : undefined
+            }
             nodeTypes={nodeTypes}
             edgeTypes={edgeTypes}
             nodesFocusable={true}
@@ -53,6 +58,11 @@ const LineageGraph = (props: ReactFlowProps) => {
             zoomOnDoubleClick={false}
             fitView
             onDoubleClick={() => fitView()}
+            onPaneClick={lineageSelection.onPaneClick}
+            onEdgeClick={lineageSelection.onEdgeClick}
+            onNodeClick={lineageSelection.onNodeClick}
+            onEdgeDoubleClick={lineageSelection.onEdgeDoubleClick}
+            onNodeDoubleClick={lineageSelection.onNodeDoubleClick}
             {...props}
         >
             <Background variant={BackgroundVariant.Dots} />
diff --git a/src/components/lineage/LineageView.tsx b/src/components/lineage/LineageView.tsx
@@ -1,13 +1,12 @@
 import { useDataProvider, useNotify } from "react-admin";
 
-import buildLineageLayout from "./utils/buildLayout";
+import { buildLineageLayout } from "./layout";
 import { Edge, Node, useNodesState, useEdgesState } from "@xyflow/react";
 import LineageFilters from "./LineageFilters";
 import LineageGraph from "./LineageGraph";
 
 import "@xyflow/react/dist/style.css";
-import getGraphNodes from "./utils/getGraphNodes";
-import getGraphEdges from "./utils/getGraphEdges";
+import { getGraphNodes, getGraphEdges } from "./converters";
 import { LineageResponseV1 } from "@/dataProvider/types";
 
 type LineageViewProps = {
@@ -22,6 +21,9 @@ const LineageView = (props: LineageViewProps) => {
     const dataProvider = useDataProvider();
     const notify = useNotify();
 
+    // datasets + 123 -> DATASET-123
+    const currentNodeId = `${props.resource.slice(0, -1).toUpperCase()}-${props.recordId}`;
+
     const [nodes, setNodes, onNodesChange] = useNodesState<Node>([]);
     const [edges, setEdges, onEdgesChange] = useEdgesState<Edge>([]);
 
@@ -44,6 +46,13 @@ const LineageView = (props: LineageViewProps) => {
                 .then((data: LineageResponseV1) => {
                     const initialNodes = getGraphNodes(data);
                     const initialEdges = getGraphEdges(data);
+
+                    initialNodes
+                        .filter((node) => node.id == currentNodeId)
+                        .forEach((node) => {
+                            node.selected = true;
+                        });
+
                     const { nodes: layoutedNodes, edges: layoutedEdges } =
                         buildLineageLayout({
                             nodes: initialNodes,
diff --git a/src/components/lineage/converters/getGraphEdges.tsx b/src/components/lineage/converters/getGraphEdges.tsx
diff --git a/src/components/lineage/converters/getGraphNodes.ts b/src/components/lineage/converters/getGraphNodes.ts
diff --git a/src/components/lineage/converters/index.ts b/src/components/lineage/converters/index.ts
@@ -0,0 +1,4 @@
+import getGraphEdges from "./getGraphEdges";
+import getGraphNodes from "./getGraphNodes";
+
+export { getGraphEdges, getGraphNodes };
diff --git a/src/components/lineage/edges/BaseEdge.css b/src/components/lineage/edges/BaseEdge.css
@@ -0,0 +1,7 @@
+.react-flow.hideNonSelected .react-flow__node:not(.selected),
+.react-flow.hideNonSelected .react-flow__edge:not(.selected),
+.react-flow.hideNonSelected
+    .react-flow__edgelabel-renderer
+    > div:not(.selected) {
+    opacity: 0.3 !important;
+}
diff --git a/src/components/lineage/edges/BaseEdge.tsx b/src/components/lineage/edges/BaseEdge.tsx
@@ -7,7 +7,14 @@ import {
 } from "@xyflow/react";
 import { Chip } from "@mui/material";
 
-const BaseEdge = ({ id, label, ...props }: EdgeProps): ReactElement => {
+import "./BaseEdge.css";
+
+const BaseEdge = ({
+    id,
+    label,
+    selected,
+    ...props
+}: EdgeProps): ReactElement => {
     const [edgePath, labelX, labelY] = getBezierPath(props);
 
     return (
@@ -21,7 +28,7 @@ const BaseEdge = ({ id, label, ...props }: EdgeProps): ReactElement => {
                             position: "absolute",
                             transform: `translate(-50%, -50%) translate(${labelX}px,${labelY}px)`,
                         }}
-                        className="nodrag nopan"
+                        className={`nodrag nopan ${selected ? "selected" : ""}`}
                     >
                         <Chip
                             size="small"
diff --git a/src/components/lineage/edges/IOEdge.tsx b/src/components/lineage/edges/IOEdge.tsx
@@ -2,7 +2,7 @@ import { ReactElement } from "react";
 import {
     getBezierPath,
     EdgeLabelRenderer,
-    BaseEdge,
+    BaseEdge as ReactFlowBaseEdge,
     EdgeProps,
 } from "@xyflow/react";
 import {
@@ -13,11 +13,14 @@ import formatNumberApprox from "@/utils/numbers";
 import formatBytes from "@/utils/bytes";
 import { Card, Chip, Typography } from "@mui/material";
 
+import "./BaseEdge.css";
+
 const IOEdge = ({
     id,
     /* eslint-disable @typescript-eslint/no-unused-vars */
     label,
     data,
+    selected,
     ...props
 }: EdgeProps & {
     data: InputRelationLineageResponseV1 | OutputRelationLineageResponseV1;
@@ -28,12 +31,12 @@ const IOEdge = ({
         data.type || data.num_rows || data.num_bytes || data.num_files;
 
     if (!hasContent) {
-        return <BaseEdge id={id} path={edgePath} {...props} />;
+        return <ReactFlowBaseEdge id={id} path={edgePath} {...props} />;
     }
 
     return (
         <>
-            <BaseEdge id={id} path={edgePath} {...props} />
+            <ReactFlowBaseEdge id={id} path={edgePath} {...props} />
             <EdgeLabelRenderer>
                 <Card
                     style={{
@@ -43,7 +46,7 @@ const IOEdge = ({
                         position: "absolute",
                         transform: `translate(-50%, -50%) translate(${labelX}px,${labelY}px)`,
                     }}
-                    className="nodrag nopan"
+                    className={`nodrag nopan ${selected ? "selected" : ""}`}
                 >
                     {data.type ? (
                         <Chip
diff --git a/src/components/lineage/layout/buildLineageLayout.ts b/src/components/lineage/layout/buildLineageLayout.ts
diff --git a/src/components/lineage/layout/index.ts b/src/components/lineage/layout/index.ts
@@ -0,0 +1,3 @@
+import buildLineageLayout from "./buildLineageLayout";
+
+export { buildLineageLayout };
diff --git a/src/components/lineage/selection/index.ts b/src/components/lineage/selection/index.ts
@@ -0,0 +1,3 @@
+import useLineageSelection from "./useLineageSelection";
+
+export { useLineageSelection };
diff --git a/src/components/lineage/selection/useLineageSelection.ts b/src/components/lineage/selection/useLineageSelection.ts
@@ -0,0 +1,133 @@
+import { Edge, Node, useReactFlow } from "@xyflow/react";
+import { useCallback, useState, MouseEvent } from "react";
+
+const useLineageSelection = () => {
+    const { getNodes, getEdges, setEdges, setNodes } = useReactFlow();
+    const [hideNonSelected, setHideNonSelected] = useState(false);
+
+    const setSelection = (
+        nodesToSelect: Set<string>,
+        edgesToSelect: Set<string>,
+        hideOthers: boolean = true,
+    ) => {
+        const nodes = getNodes().map((node) => {
+            node.selected = nodesToSelect.has(node.id);
+            return node;
+        });
+
+        const edges = getEdges().map((edge) => {
+            edge.selected = edgesToSelect.has(edge.id);
+            return edge;
+        });
+
+        setNodes(nodes);
+        setEdges(edges);
+        setHideNonSelected(
+            hideOthers && nodesToSelect.size + edgesToSelect.size > 0,
+        );
+    };
+
+    const trackUpstreams = (
+        nodeId: string,
+    ): { nodes: Set<string>; edges: Set<string> } => {
+        const result = {
+            nodes: new Set<string>([nodeId]),
+            edges: new Set<string>(),
+        };
+
+        getEdges()
+            .filter((edge) => edge.target == nodeId)
+            .forEach((edge) => {
+                result.nodes.add(edge.source);
+                result.edges.add(edge.id);
+
+                const { nodes, edges } = trackUpstreams(edge.source);
+                nodes.forEach((node) => result.nodes.add(node));
+                edges.forEach((edge) => result.edges.add(edge));
+            });
+
+        return result;
+    };
+
+    const trackDownstreams = (
+        nodeId: string,
+    ): { nodes: Set<string>; edges: Set<string> } => {
+        const result = {
+            nodes: new Set<string>([nodeId]),
+            edges: new Set<string>(),
+        };
+
+        getEdges()
+            .filter((edge) => edge.source == nodeId)
+            .forEach((edge) => {
+                result.nodes.add(edge.target);
+                result.edges.add(edge.id);
+
+                const { nodes, edges } = trackDownstreams(edge.target);
+                nodes.forEach((node) => result.nodes.add(node));
+                edges.forEach((edge) => result.edges.add(edge));
+            });
+
+        return result;
+    };
+
+    const onEdgeClick = useCallback((e: MouseEvent, currentEdge: Edge) => {
+        setSelection(
+            new Set([currentEdge.source, currentEdge.target]),
+            new Set([currentEdge.id]),
+        );
+        e.stopPropagation();
+    }, []);
+
+    const onEdgeDoubleClick = useCallback(
+        (e: MouseEvent, currentEdge: Edge) => {
+            const upstreams = trackUpstreams(currentEdge.source);
+            const downstreams = trackDownstreams(currentEdge.target);
+
+            const relatedNodes = upstreams.nodes.union(downstreams.nodes);
+            const relatedEdges = upstreams.edges.union(downstreams.edges);
+            relatedEdges.add(currentEdge.id);
+
+            setSelection(relatedNodes, relatedEdges);
+            e.stopPropagation();
+        },
+        [],
+    );
+
+    const onNodeClick = useCallback((e: MouseEvent, currentNode: Node) => {
+        // If only one node is selected, don't hide others, as it not very convenient
+        setSelection(new Set([currentNode.id]), new Set(), false);
+        e.stopPropagation();
+    }, []);
+
+    const onNodeDoubleClick = useCallback(
+        (e: MouseEvent, currentNode: Node) => {
+            const upstreams = trackUpstreams(currentNode.id);
+            const downstreams = trackDownstreams(currentNode.id);
+
+            const relatedNodes = upstreams.nodes.union(downstreams.nodes);
+            const relatedEdges = upstreams.edges.union(downstreams.edges);
+            relatedNodes.add(currentNode.id);
+
+            setSelection(relatedNodes, relatedEdges);
+            e.stopPropagation();
+        },
+        [],
+    );
+
+    const onPaneClick = useCallback((e: MouseEvent) => {
+        setSelection(new Set(), new Set());
+        e.stopPropagation();
+    }, []);
+
+    return {
+        hideNonSelected,
+        onEdgeClick,
+        onEdgeDoubleClick,
+        onNodeClick,
+        onNodeDoubleClick,
+        onPaneClick,
+    };
+};
+
+export default useLineageSelection;

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+import buildLineageLayout from "./buildLineageLayout";`
	`2`	`+`
	`3`	`+export { buildLineageLayout };`