From 5e6f79b6852d10a26b0250292420ca645e5acca6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=9C=D0=B0=D1=80=D1=82=D1=8B=D0=BD=D0=BE=D0=B2=20=D0=9C?= =?UTF-8?q?=D0=B0=D0=BA=D1=81=D0=B8=D0=BC=20=D0=A1=D0=B5=D1=80=D0=B3=D0=B5?= =?UTF-8?q?=D0=B5=D0=B2=D0=B8=D1=87?= Date: Tue, 4 Mar 2025 16:32:51 +0300 Subject: [PATCH] [DOP-24183] Highlight column lineage on click --- src/components/lineage/LineageGraph.tsx | 121 ++++++++----- src/components/lineage/LineageView.tsx | 2 +- .../lineage/converters/getGraphEdges.tsx | 16 +- .../lineage/edges/ColumnLineageEdge.css | 5 + .../lineage/edges/ColumnLineageEdge.tsx | 4 +- .../nodes/dataset_node/DatasetNode.css | 8 + .../nodes/dataset_node/DatasetNode.tsx | 5 +- .../nodes/dataset_node/DatasetSchemaTable.tsx | 41 ++++- .../selection/LineageSelectionContext.tsx | 15 ++ src/components/lineage/selection/index.ts | 3 - .../lineage/selection/useLineageSelection.ts | 162 ------------------ .../selection/useLineageSelectionProvider.tsx | 45 +++++ .../selection/utils/columnSelection.ts | 73 ++++++++ .../lineage/selection/utils/common.ts | 161 +++++++++++++++++ .../lineage/selection/utils/edgeSelection.ts | 54 ++++++ .../lineage/selection/utils/nodeSelection.ts | 33 ++++ 16 files changed, 526 insertions(+), 222 deletions(-) create mode 100644 src/components/lineage/edges/ColumnLineageEdge.css create mode 100644 src/components/lineage/selection/LineageSelectionContext.tsx delete mode 100644 src/components/lineage/selection/index.ts delete mode 100644 src/components/lineage/selection/useLineageSelection.ts create mode 100644 src/components/lineage/selection/useLineageSelectionProvider.tsx create mode 100644 src/components/lineage/selection/utils/columnSelection.ts create mode 100644 src/components/lineage/selection/utils/common.ts create mode 100644 src/components/lineage/selection/utils/edgeSelection.ts create mode 100644 src/components/lineage/selection/utils/nodeSelection.ts diff --git a/src/components/lineage/LineageGraph.tsx b/src/components/lineage/LineageGraph.tsx index 680d476..709ea54 100644 --- 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()]]), + 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 ( - fitView()} - {...selectionHandlers} - {...props} - > - - - - + + fitView()} + onEdgeClick={onEdgeClick} + onEdgeDoubleClick={onEdgeDoubleClick} + onNodeClick={onNodeClick} + onNodeDoubleClick={onNodeDoubleClick} + onPaneClick={onPaneClick} + {...props} + > + + + + + ); }; diff --git a/src/components/lineage/LineageView.tsx b/src/components/lineage/LineageView.tsx index 04968f0..8009ec9 100644 --- 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} - > + /> )} diff --git a/src/components/lineage/converters/getGraphEdges.tsx b/src/components/lineage/converters/getGraphEdges.tsx index 82d4722..8074e3b 100644 --- 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 new file mode 100644 index 0000000..f253d6a --- /dev/null +++ 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 index 4525677..ebfdadc 100644 --- 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 index 4417fa7..cfbb2bf 100644 --- 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 index e839129..b730187 100644 --- 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): ReactElement => { { smart_count: props.data.schemaCount }, )} - + ) : null } diff --git a/src/components/lineage/nodes/dataset_node/DatasetSchemaTable.tsx b/src/components/lineage/nodes/dataset_node/DatasetSchemaTable.tsx index b902f7f..c6f2df7 100644 --- 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>(); + + 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 ( <> {fieldsToShow.map((field) => ( - + onFieldClick(e, field.name)} + onDoubleClick={(e) => + onFieldDoubleClick(e, field.name) + } + > >(), + edges: new Set(), + }, + 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 deleted file mode 100644 index cda0a5b..0000000 --- a/src/components/lineage/selection/index.ts +++ /dev/null @@ -1,3 +0,0 @@ -import useLineageSelection from "./useLineageSelection"; - -export { useLineageSelection }; diff --git a/src/components/lineage/selection/useLineageSelection.ts b/src/components/lineage/selection/useLineageSelection.ts deleted file mode 100644 index 1b46485..0000000 --- a/src/components/lineage/selection/useLineageSelection.ts +++ /dev/null @@ -1,162 +0,0 @@ -import { Edge, Node, useReactFlow } from "@xyflow/react"; -import { useCallback, MouseEvent } from "react"; - -const getAllConnection = ( - startNodeIds: string[], - edgesByNodeId: Map, - edgeAttribute: "source" | "target", -): { nodes: Set; edges: Set } => { - const result = { - nodes: new Set(), - edges: new Set(), - }; - - let startFrom = new Set(startNodeIds); - while (edgesByNodeId.size > 0) { - const visitedNodes = new Set(); - const visitedEdges = new Set(); - - startFrom.forEach((nodeId) => { - const edges = edgesByNodeId.get(nodeId); - if (edges) { - edges.forEach((edge: Edge) => { - visitedNodes.add(edge[edgeAttribute]); - visitedEdges.add(edge.id); - }); - edgesByNodeId.delete(nodeId); - } - }); - if (visitedNodes.size == 0 && visitedEdges.size == 0) { - break; - } - - visitedNodes.forEach((nodeId) => { - result.nodes.add(nodeId); - }); - visitedEdges.forEach((edgeId) => { - result.edges.add(edgeId); - }); - startFrom = visitedNodes; - } - - return result; -}; - -const useLineageSelection = () => { - const { getNodes, getEdges, setEdges, setNodes } = useReactFlow(); - - const setSelection = useCallback( - (nodesToSelect: Set, edgesToSelect: Set) => { - 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); - }, - [], - ); - - const getAllRelations = useCallback( - (nodeIds: string[]): { nodes: Set; edges: Set } => { - const result = { - nodes: new Set(nodeIds), - edges: new Set(), - }; - - const rawEdges = getEdges(); - - // Convert list to maps to use O(1) lookups - const edgesBySource = new Map(); - const edgesByTarget = new Map(); - rawEdges.forEach((edge) => { - const sameSource = edgesBySource.get(edge.source) ?? []; - sameSource.push(edge); - edgesBySource.set(edge.source, sameSource); - - const sameTarget = edgesByTarget.get(edge.target) ?? []; - sameTarget.push(edge); - edgesByTarget.set(edge.target, sameTarget); - }); - - const upstreams = getAllConnection( - nodeIds, - edgesByTarget, - "source", - ); - const downstreams = getAllConnection( - nodeIds, - edgesBySource, - "target", - ); - - upstreams.nodes.forEach((nodeId) => result.nodes.add(nodeId)); - upstreams.edges.forEach((edgeId) => result.edges.add(edgeId)); - downstreams.nodes.forEach((nodeId) => result.nodes.add(nodeId)); - downstreams.edges.forEach((edgeId) => result.edges.add(edgeId)); - - return result; - }, - [], - ); - - const selectNode = useCallback((nodeId: string) => { - setSelection(new Set([nodeId]), new Set()); - }, []); - - 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 { nodes, edges } = getAllRelations([ - currentEdge.source, - currentEdge.target, - ]); - setSelection(nodes, edges); - e.stopPropagation(); - }, - [], - ); - - const onNodeClick = useCallback((e: MouseEvent, currentNode: Node) => { - selectNode(currentNode.id); - e.stopPropagation(); - }, []); - - const onNodeDoubleClick = useCallback( - (e: MouseEvent, currentNode: Node) => { - const { nodes, edges } = getAllRelations([currentNode.id]); - setSelection(nodes, edges); - e.stopPropagation(); - }, - [], - ); - - const onPaneClick = useCallback((e: MouseEvent) => { - setSelection(new Set(), new Set()); - e.stopPropagation(); - }, []); - - return { - setSelection, - onEdgeClick, - onEdgeDoubleClick, - onNodeClick, - onNodeDoubleClick, - onPaneClick, - }; -}; - -export default useLineageSelection; diff --git a/src/components/lineage/selection/useLineageSelectionProvider.tsx b/src/components/lineage/selection/useLineageSelectionProvider.tsx new file mode 100644 index 0000000..f9fefa0 --- /dev/null +++ b/src/components/lineage/selection/useLineageSelectionProvider.tsx @@ -0,0 +1,45 @@ +import { useReactFlow } from "@xyflow/react"; +import { useCallback, useState } from "react"; +import { LineageSelection } from "./utils/common"; + +const useLineageSelectionProvider = () => { + const { getNodes, setNodes, getEdges, setEdges } = useReactFlow(); + + const [selection, setSelectionState] = useState({ + nodeWithColumns: new Map>(), + edges: new Set(), + }); + + const setSelection = useCallback((newValue: LineageSelection) => { + setSelectionState(newValue); + const newNodes = getNodes().map((node) => { + return { + ...node, + selected: newValue.nodeWithColumns.has(node.id), + }; + }); + const newEdges = getEdges().map((edge) => { + return { + ...edge, + selected: newValue.edges.has(edge.id), + }; + }); + setNodes(newNodes); + setEdges(newEdges); + }, []); + + const resetSelection = useCallback(() => { + setSelectionState({ + nodeWithColumns: new Map>(), + edges: new Set(), + }); + }, []); + + return { + selection, + setSelection, + resetSelection, + }; +}; + +export default useLineageSelectionProvider; diff --git a/src/components/lineage/selection/utils/columnSelection.ts b/src/components/lineage/selection/utils/columnSelection.ts new file mode 100644 index 0000000..99ce0cc --- /dev/null +++ b/src/components/lineage/selection/utils/columnSelection.ts @@ -0,0 +1,73 @@ +import { Edge } from "@xyflow/react"; +import { + getAllConnections, + mergeSelection, + LineageSelection, + splitEdges, +} from "./common"; + +export const getNearestColumnRelations = ( + edges: Edge[], + nodeId: string, + fieldName: string, +): LineageSelection => { + // For specific node and column return only nearest edges and their columns. + + const connectedEdges = edges.filter( + (edge) => + (edge.source === nodeId && edge.sourceHandle === fieldName) || + (edge.target === nodeId && edge.targetHandle === fieldName), + ); + + const result: LineageSelection = { + nodeWithColumns: new Map(), + edges: new Set(connectedEdges.map((edge) => edge.id)), + }; + + connectedEdges.forEach((edge) => { + const columns = + result.nodeWithColumns.get(edge.source) ?? new Set(); + if (edge.sourceHandle) { + columns.add(edge.sourceHandle); + } + result.nodeWithColumns.set(edge.source, columns); + }); + + connectedEdges.forEach((edge) => { + const columns = + result.nodeWithColumns.get(edge.target) ?? new Set(); + if (edge.targetHandle) { + columns.add(edge.targetHandle); + } + result.nodeWithColumns.set(edge.target, columns); + }); + + return result; +}; + +export const getAllColumnRelations = ( + edges: Edge[], + nodeId: string, + fieldName: string, +): LineageSelection => { + // For specific node and column return all connected edges and nodes, recursively. + + const { edgesBySource, edgesByTarget } = splitEdges(edges); + + // walk `source -> edge - target` + const downstreams = getAllConnections( + nodeId, + new Set([fieldName]), + edgesBySource, + "outgoing", + ); + // same, but in opposite direction + const upstreams = getAllConnections( + nodeId, + new Set([fieldName]), + edgesByTarget, + "incoming", + ); + + return mergeSelection(upstreams, downstreams); +}; diff --git a/src/components/lineage/selection/utils/common.ts b/src/components/lineage/selection/utils/common.ts new file mode 100644 index 0000000..5624127 --- /dev/null +++ b/src/components/lineage/selection/utils/common.ts @@ -0,0 +1,161 @@ +import { Edge } from "@xyflow/react"; + +const mergeSets = (set1: Set, set2: Set): Set => { + // [...set1, ...set2] is not available in old browsers, using a workaround + const result = new Set(); + set1.forEach((value) => result.add(value)); + set2.forEach((value) => result.add(value)); + return result; +}; + +const mergeMaps = ( + map1: Map>, + map2: Map>, +): Map> => { + // [...map1, ...map2] is not available in old browsers, using a workaround. + // Also values are collections which should be merged too. + const result = new Map>(); + + map1.forEach((newValues, key) => { + const existingValues = result.get(key) ?? new Set(); + result.set(key, mergeSets(existingValues, newValues)); + }); + map2.forEach((newValues, key) => { + const existingValues = result.get(key) ?? new Set(); + result.set(key, mergeSets(existingValues, newValues)); + }); + return result; +}; + +export type LineageSelection = { + // empty column set means select the entire node + nodeWithColumns: Map>; + edges: Set; +}; + +export const splitEdges = ( + edges: Edge[], +): { + edgesBySource: Map; + edgesByTarget: Map; +} => { + // Convert list to maps to use O(1) lookups + const edgesBySource = new Map(); + const edgesByTarget = new Map(); + edges.forEach((edge) => { + const sameSource = edgesBySource.get(edge.source) ?? []; + sameSource.push(edge); + edgesBySource.set(edge.source, sameSource); + + const sameTarget = edgesByTarget.get(edge.target) ?? []; + sameTarget.push(edge); + edgesByTarget.set(edge.target, sameTarget); + }); + + return { edgesBySource, edgesByTarget }; +}; + +export const getAllConnections = ( + startNodesId: string, + startNodeColumns: Set, + edgesByNodeId: Map, + direction: "incoming" | "outgoing", +): LineageSelection => { + // Iterate over edges connected to specific node. + // If startNodeColumns is passed, select only edges with specified sourceHandle/targetHandle. + // Go to connected source/target nodes by these edges. Recursively. + // Return all visited nodes, columns and edges. + + let edgeNodeAttribute: "source" | "target"; + let directNodeHandle: "sourceHandle" | "targetHandle"; + let reverseNodeHandle: "sourceHandle" | "targetHandle"; + if (direction == "incoming") { + edgeNodeAttribute = "source"; + directNodeHandle = "targetHandle"; + reverseNodeHandle = "sourceHandle"; + } else { + edgeNodeAttribute = "target"; + directNodeHandle = "sourceHandle"; + reverseNodeHandle = "targetHandle"; + } + + const result: LineageSelection = { + nodeWithColumns: new Map([[startNodesId, startNodeColumns]]), + edges: new Set(), + }; + + let startFrom = result.nodeWithColumns; + while (edgesByNodeId.size > 0) { + const visitedNodeColumns = new Map>(); + const visitedEdges = new Set(); + + startFrom.forEach((columns, nodeId) => { + const edges = edgesByNodeId.get(nodeId); + if (!edges) { + return; + } + edges.forEach((edge: Edge) => { + const nodeId = edge[edgeNodeAttribute]; + const columnToSearch = edge[directNodeHandle]; + const columnToInclude = edge[reverseNodeHandle]; + + const visitedColumns = + visitedNodeColumns.get(nodeId) ?? new Set(); + + if ( + columns.size > 0 && + columnToSearch && + columnToInclude && + columns.has(columnToSearch) + ) { + visitedEdges.add(edge.id); + visitedColumns.add(columnToInclude); + visitedNodeColumns.set(nodeId, visitedColumns); + } else if (columns.size == 0) { + visitedEdges.add(edge.id); + visitedNodeColumns.set(nodeId, visitedColumns); + } + }); + edgesByNodeId.delete(nodeId); + }); + if (visitedNodeColumns.size == 0 && visitedEdges.size == 0) { + break; + } + + result.nodeWithColumns = mergeMaps( + result.nodeWithColumns, + visitedNodeColumns, + ); + result.edges = mergeSets(result.edges, visitedEdges); + startFrom = visitedNodeColumns; + } + + return result; +}; + +export const mergeSelection = ( + selection1: LineageSelection, + selection2: LineageSelection, +): LineageSelection => { + return { + nodeWithColumns: mergeMaps( + selection1.nodeWithColumns, + selection2.nodeWithColumns, + ), + edges: mergeSets(selection1.edges, selection2.edges), + }; +}; + +export const isSubgraphSelected = (selection: LineageSelection) => { + if (selection.edges.size > 0) { + return true; + } + + let someColumnsSelected = false; + selection.nodeWithColumns.forEach((columns) => { + if (columns.size > 0) { + someColumnsSelected = true; + } + }); + return someColumnsSelected; +}; diff --git a/src/components/lineage/selection/utils/edgeSelection.ts b/src/components/lineage/selection/utils/edgeSelection.ts new file mode 100644 index 0000000..972cae7 --- /dev/null +++ b/src/components/lineage/selection/utils/edgeSelection.ts @@ -0,0 +1,54 @@ +import { Edge } from "@xyflow/react"; +import { + getAllConnections, + mergeSelection, + LineageSelection, + splitEdges, +} from "./common"; + +export const getNearestEdgeRelations = (edge: Edge): LineageSelection => { + // For specific edge return connected nodes and columns (if any). + + return { + nodeWithColumns: new Map>([ + [ + edge.source, + new Set([edge.sourceHandle].filter((field) => field != null)), + ], + [ + edge.target, + new Set([edge.targetHandle].filter((field) => field != null)), + ], + ]), + edges: new Set([edge.id]), + }; +}; + +export const getAllEdgeRelations = ( + edges: Edge[], + edge: Edge, +): LineageSelection => { + // For specific edge return all connected nodes and edges, recursively. + + const { edgesBySource, edgesByTarget } = splitEdges(edges); + + // walk `source -> edge -> target` + // for columnLineageEdge, select nodes and columns connected to it + const downstreams = getAllConnections( + edge.target, + new Set([edge.targetHandle].filter((field) => field != null)), + edgesBySource, + "outgoing", + ); + // same, but in opposite direction + const upstreams = getAllConnections( + edge.source, + new Set([edge.sourceHandle].filter((field) => field != null)), + edgesByTarget, + "incoming", + ); + + const result = mergeSelection(downstreams, upstreams); + result.edges.add(edge.id); + return result; +}; diff --git a/src/components/lineage/selection/utils/nodeSelection.ts b/src/components/lineage/selection/utils/nodeSelection.ts new file mode 100644 index 0000000..ebd52d4 --- /dev/null +++ b/src/components/lineage/selection/utils/nodeSelection.ts @@ -0,0 +1,33 @@ +import { Edge } from "@xyflow/react"; +import { + getAllConnections, + mergeSelection, + LineageSelection, + splitEdges, +} from "./common"; + +export const getAllNodeRelations = ( + edges: Edge[], + nodeId: string, +): LineageSelection => { + // For specific node return all connected edges and nodes, recursively. + + // Convert list to maps to use O(1) lookups + const { edgesBySource, edgesByTarget } = splitEdges(edges); + + // walk `source -> edge -> target` + const downstreams = getAllConnections( + nodeId, + new Set(), + edgesBySource, + "outgoing", + ); + const upstreams = getAllConnections( + nodeId, + new Set(), + edgesByTarget, + "incoming", + ); + + return mergeSelection(downstreams, upstreams); +};