[DOP-24183] Highlight column lineage on click

MobileTeleSystems · Mar 5, 2025 · 5e6f79b · 5e6f79b
1 parent cdd4e9a
commit 5e6f79b
Show file tree

Hide file tree

Showing 16 changed files with 526 additions and 222 deletions.
diff --git a/src/components/lineage/LineageGraph.tsx b/src/components/lineage/LineageGraph.tsx
@@ -8,11 +8,19 @@ import {
     useNodesInitialized,
     BackgroundVariant,
     Edge,
+    Node,
 } from "@xyflow/react";
 import { DatasetNode, JobNode, RunNode, OperationNode } from "./nodes";
-import { useEffect } from "react";
+import { MouseEvent, useEffect } from "react";
 import { BaseEdge, IOEdge, ColumnLineageEdge } from "./edges";
-import { useLineageSelection } from "./selection";
+import useLineageSelectionProvider from "./selection/useLineageSelectionProvider";
+import LineageSelectionContext from "./selection/LineageSelectionContext";
+import { getAllNodeRelations } from "./selection/utils/nodeSelection";
+import {
+    getNearestEdgeRelations,
+    getAllEdgeRelations,
+} from "./selection/utils/edgeSelection";
+import { isSubgraphSelected } from "./selection/utils/common";
 
 export const MIN_ZOOM_VALUE = 0.1;
 export const MAX_ZOOM_VALUE = 2.5;
@@ -30,54 +38,83 @@ const nodeTypes = {
     operationNode: OperationNode,
 };
 
-const subgraphSelected = (edges?: Edge[]) => {
-    if (!edges) {
-        return false;
-    }
-    for (const edge of edges) {
-        if (edge.selected) {
-            return true;
-        }
-    }
-    return false;
-};
-
 const LineageGraph = (props: ReactFlowProps) => {
-    const { fitView } = useReactFlow();
-    const selectionHandlers = useLineageSelection();
+    const { fitView, getEdges } = useReactFlow();
     const nodesInitialized = useNodesInitialized();
 
+    const lineageSelection = useLineageSelectionProvider();
+    const { selection, setSelection, resetSelection } = lineageSelection;
+
+    const onEdgeClick = (e: MouseEvent, edge: Edge) => {
+        const selection = getNearestEdgeRelations(edge);
+        setSelection(selection);
+        e.stopPropagation();
+    };
+
+    const onEdgeDoubleClick = (e: MouseEvent, edge: Edge) => {
+        const selection = getAllEdgeRelations(getEdges(), edge);
+        setSelection(selection);
+        e.stopPropagation();
+    };
+
+    const onNodeClick = (e: MouseEvent, node: Node) => {
+        setSelection({
+            nodeWithColumns: new Map([[node.id, new Set<string>()]]),
+            edges: new Set(),
+        });
+        e.stopPropagation();
+    };
+
+    const onNodeDoubleClick = (e: MouseEvent, node: Node) => {
+        const selection = getAllNodeRelations(getEdges(), node.id);
+        setSelection(selection);
+        e.stopPropagation();
+    };
+
+    const onPaneClick = (e: MouseEvent) => {
+        resetSelection();
+        e.stopPropagation();
+    };
+
     useEffect(() => {
         fitView();
     }, [nodesInitialized]);
 
     return (
-        <ReactFlow
-            className={
-                subgraphSelected(props.edges) ? "subgraphSelected" : undefined
-            }
-            nodeTypes={nodeTypes}
-            edgeTypes={edgeTypes}
-            nodesFocusable={true}
-            elementsSelectable={true}
-            nodesConnectable={false}
-            nodesDraggable={true}
-            panOnScroll={false}
-            panOnDrag={true}
-            minZoom={MIN_ZOOM_VALUE}
-            maxZoom={MAX_ZOOM_VALUE}
-            zoomOnScroll={true}
-            zoomOnPinch={true}
-            zoomOnDoubleClick={false}
-            fitView
-            onDoubleClick={() => fitView()}
-            {...selectionHandlers}
-            {...props}
-        >
-            <Background variant={BackgroundVariant.Dots} />
-            <Controls />
-            <MiniMap pannable zoomable />
-        </ReactFlow>
+        <LineageSelectionContext.Provider value={lineageSelection}>
+            <ReactFlow
+                className={
+                    isSubgraphSelected(selection)
+                        ? "subgraphSelected"
+                        : undefined
+                }
+                nodeTypes={nodeTypes}
+                edgeTypes={edgeTypes}
+                nodesFocusable={true}
+                elementsSelectable={true}
+                nodesConnectable={false}
+                nodesDraggable={true}
+                panOnScroll={false}
+                panOnDrag={true}
+                minZoom={MIN_ZOOM_VALUE}
+                maxZoom={MAX_ZOOM_VALUE}
+                zoomOnScroll={true}
+                zoomOnPinch={true}
+                zoomOnDoubleClick={false}
+                fitView
+                onDoubleClick={() => fitView()}
+                onEdgeClick={onEdgeClick}
+                onEdgeDoubleClick={onEdgeDoubleClick}
+                onNodeClick={onNodeClick}
+                onNodeDoubleClick={onNodeDoubleClick}
+                onPaneClick={onPaneClick}
+                {...props}
+            >
+                <Background variant={BackgroundVariant.Dots} />
+                <Controls />
+                <MiniMap pannable zoomable />
+            </ReactFlow>
+        </LineageSelectionContext.Provider>
     );
 };
 

diff --git a/src/components/lineage/LineageView.tsx b/src/components/lineage/LineageView.tsx
@@ -78,7 +78,7 @@ const LineageView = (props: LineageViewProps) => {
                         edges={edges}
                         onNodesChange={onNodesChange}
                         onEdgesChange={onEdgesChange}
-                    ></LineageGraph>
+                    />
                 </div>
             )}
         </>

diff --git a/src/components/lineage/converters/getGraphEdges.tsx b/src/components/lineage/converters/getGraphEdges.tsx
@@ -250,7 +250,7 @@ const getSymlinkEdge = (
 const getDirectColumnLineageEdges = (
     relation: DirectColumnLineageRelationLineageResponseV1,
 ): Edge[] => {
-    const color = "#b1b1b7";
+    const color = "gray";
     return Object.keys(relation.fields).flatMap((target_field_name) => {
         return relation.fields[target_field_name].map((source_field) => {
             return {
@@ -265,17 +265,17 @@ const getDirectColumnLineageEdges = (
                     types: source_field.types,
                     kind: "DIRECT_COLUMN_LINEAGE",
                 },
-                animated: true,
                 markerEnd: {
                     type: MarkerType.ArrowClosed,
                     color: color,
                 },
                 style: {
-                    strokeWidth: STOKE_MEDIUM,
+                    strokeWidth: STOKE_THIN,
                     stroke: color,
                 },
                 labelStyle: {
-                    backgroundColor: color,
+                    // label is shown only then edge is selected, so color is different from stroke
+                    backgroundColor: "#89b2f3",
                 },
             };
         });
@@ -285,7 +285,7 @@ const getDirectColumnLineageEdges = (
 const getIndirectColumnLineageEdges = (
     relation: IndirectColumnLineageRelationLineageResponseV1,
 ): Edge[] => {
-    const color = "#b1b1b7";
+    const color = "gray";
     return relation.fields.map((source_field) => {
         return {
             ...getMinimalEdge(relation),
@@ -298,17 +298,17 @@ const getIndirectColumnLineageEdges = (
                 types: source_field.types,
                 kind: "INDIRECT_COLUMN_LINEAGE",
             },
-            animated: true,
             markerEnd: {
                 type: MarkerType.ArrowClosed,
                 color: color,
             },
             style: {
-                strokeWidth: STOKE_MEDIUM,
+                strokeWidth: STOKE_THIN,
                 stroke: color,
             },
             labelStyle: {
-                backgroundColor: color,
+                // label is shown only then edge is selected, so color is different from stroke
+                backgroundColor: "#89b2f3",
             },
         };
     });

diff --git a/src/components/lineage/edges/ColumnLineageEdge.css b/src/components/lineage/edges/ColumnLineageEdge.css
@@ -0,0 +1,5 @@
+.react-flow__edge.react-flow__edge-columnLineageEdge.selected
+    .react-flow__edge-path {
+    stroke: #89b2f3 !important;
+    stroke-width: 2px !important;
+}
diff --git a/src/components/lineage/edges/ColumnLineageEdge.tsx b/src/components/lineage/edges/ColumnLineageEdge.tsx
@@ -6,9 +6,9 @@ import {
     EdgeProps,
 } from "@xyflow/react";
 import { ColumnLineageFieldResponseV1 } from "@/dataProvider/types";
-import { Card, Chip } from "@mui/material";
+import { Chip } from "@mui/material";
 
-import "./BaseEdge.css";
+import "./ColumnLineageEdge.css";
 
 const ColumnLineageEdge = ({
     id,

diff --git a/src/components/lineage/nodes/dataset_node/DatasetNode.css b/src/components/lineage/nodes/dataset_node/DatasetNode.css
@@ -10,3 +10,11 @@
     transform: none;
     border: none;
 }
+
+.react-flow__node .columnLineageField.selected {
+    background-color: #b4d1ff;
+    border-color: var(
+        --xy-selection-border,
+        var(--xy-selection-border-default)
+    );
+}
diff --git a/src/components/lineage/nodes/dataset_node/DatasetNode.tsx b/src/components/lineage/nodes/dataset_node/DatasetNode.tsx
@@ -52,7 +52,10 @@ const DatasetNode = (props: NodeProps<DatasetNode>): ReactElement => {
                                 { smart_count: props.data.schemaCount },
                             )}
                         </Typography>
-                        <DatasetSchemaTable fields={props.data.schema.fields} />
+                        <DatasetSchemaTable
+                            nodeId={props.id}
+                            fields={props.data.schema.fields}
+                        />
                     </>
                 ) : null
             }

diff --git a/src/components/lineage/nodes/dataset_node/DatasetSchemaTable.tsx b/src/components/lineage/nodes/dataset_node/DatasetSchemaTable.tsx
@@ -10,20 +10,28 @@ import {
     TextField,
 } from "@mui/material";
 import { useTranslate } from "react-admin";
-import { Handle, Position } from "@xyflow/react";
-import { useMemo, useState } from "react";
+import { Handle, Position, useReactFlow } from "@xyflow/react";
+import { MouseEvent, useContext, useMemo, useState } from "react";
 import { Search } from "@mui/icons-material";
 import { IORelationSchemaFieldV1 } from "@/dataProvider/types";
 import { paginateArray } from "../../utils/pagination";
 import { fieldMatchesText, flattenFields } from "./utils";
+import {
+    getNearestColumnRelations,
+    getAllColumnRelations,
+} from "../../selection/utils/columnSelection";
+import LineageSelectionContext from "../../selection/LineageSelectionContext";
 
 const DatasetSchemaTable = ({
+    nodeId,
     fields,
     defaultRowsPerPage = 10,
 }: {
+    nodeId: string;
     fields: IORelationSchemaFieldV1[];
     defaultRowsPerPage?: number;
 }) => {
+    const { getEdges } = useReactFlow();
     const translate = useTranslate();
 
     const [page, setPage] = useState(0);
@@ -53,6 +61,26 @@ const DatasetSchemaTable = ({
         },
     ];
 
+    const { selection, setSelection } = useContext(LineageSelectionContext);
+    const selectedFields =
+        selection.nodeWithColumns.get(nodeId) ?? new Map<string, Set<string>>();
+
+    const onFieldClick = (e: MouseEvent, fieldName: string) => {
+        const selection = getNearestColumnRelations(
+            getEdges(),
+            nodeId,
+            fieldName,
+        );
+        setSelection(selection);
+        e.stopPropagation();
+    };
+
+    const onFieldDoubleClick = (e: MouseEvent, fieldName: string) => {
+        const selection = getAllColumnRelations(getEdges(), nodeId, fieldName);
+        setSelection(selection);
+        e.stopPropagation();
+    };
+
     return (
         <>
             <Stack
@@ -101,7 +129,14 @@ const DatasetSchemaTable = ({
                 </TableHead>
                 <TableBody>
                     {fieldsToShow.map((field) => (
-                        <TableRow key={field.name}>
+                        <TableRow
+                            key={field.name}
+                            className={`columnLineageField ${selectedFields.has(field.name) ? "selected" : ""}`}
+                            onClick={(e) => onFieldClick(e, field.name)}
+                            onDoubleClick={(e) =>
+                                onFieldDoubleClick(e, field.name)
+                            }
+                        >
                             <TableCell className="hidden">
                                 <Handle
                                     className="columnLineageHandle"

diff --git a/src/components/lineage/selection/LineageSelectionContext.tsx b/src/components/lineage/selection/LineageSelectionContext.tsx
@@ -0,0 +1,15 @@
+import { createContext } from "react";
+import { LineageSelection } from "./utils/common";
+
+const LineageSelectionValue = {
+    selection: {
+        nodeWithColumns: new Map<string, Set<string>>(),
+        edges: new Set<string>(),
+    },
+    setSelection: (newValue: LineageSelection): void => {},
+    resetSelection: (): void => {},
+};
+
+const LineageSelectionContext = createContext(LineageSelectionValue);
+
+export default LineageSelectionContext;
diff --git a/src/components/lineage/selection/index.ts b/src/components/lineage/selection/index.ts