diff --git a/web/common/src/components/Lineage/LineageLayout.tsx b/web/common/src/components/Lineage/LineageLayout.tsx index 411ace4e65..4b2e06b0b3 100644 --- a/web/common/src/components/Lineage/LineageLayout.tsx +++ b/web/common/src/components/Lineage/LineageLayout.tsx @@ -2,7 +2,9 @@ import { Background, BackgroundVariant, Controls, + type EdgeChange, type EdgeTypes, + type NodeChange, type NodeTypes, ReactFlow, ReactFlowProvider, @@ -12,6 +14,8 @@ import { getOutgoers, useReactFlow, useViewport, + applyNodeChanges, + applyEdgeChanges, } from '@xyflow/react' import '@xyflow/react/dist/style.css' @@ -55,6 +59,8 @@ export function LineageLayout< edgeTypes, className, controls, + nodesDraggable, + nodesConnectable, useLineage, onNodeClick, onNodeDoubleClick, @@ -69,6 +75,8 @@ export function LineageLayout< nodeTypes?: NodeTypes edgeTypes?: EdgeTypes className?: string + nodesDraggable?: boolean + nodesConnectable?: boolean controls?: | React.ReactNode | (({ setCenter }: { setCenter: SetCenter }) => React.ReactNode) @@ -86,6 +94,8 @@ export function LineageLayout< + nodesDraggable?: boolean + nodesConnectable?: boolean nodeTypes?: NodeTypes edgeTypes?: EdgeTypes className?: string @@ -140,8 +154,8 @@ function LineageLayoutBase< isBuildingLayout, currentNode, zoom, - nodes, - edges, + nodes: initialNodes, + edges: initialEdges, nodesMap, showOnlySelectedNodes, selectedNodeId, @@ -152,6 +166,32 @@ function LineageLayoutBase< setSelectedEdges, } = useLineage() + const [nodes, setNodes] = React.useState(initialNodes) + const [edges, setEdges] = React.useState(initialEdges) + + const onNodesChange = React.useCallback( + (changes: NodeChange>[]) => { + setNodes( + applyNodeChanges>(changes, nodes), + ) + }, + [nodes, setNodes], + ) + + const onEdgesChange = React.useCallback( + ( + changes: EdgeChange>[], + ) => { + setEdges( + applyEdgeChanges>( + changes, + edges, + ), + ) + }, + [edges, setEdges], + ) + const updateZoom = React.useMemo(() => debounce(setZoom, 200), [setZoom]) const zoomToCurrentNode = React.useCallback( @@ -221,6 +261,14 @@ function LineageLayoutBase< [nodes, edges], ) + React.useEffect(() => { + setNodes(initialNodes) + }, [initialNodes]) + + React.useEffect(() => { + setEdges(initialEdges) + }, [initialEdges]) + React.useEffect(() => { if (selectedNodeId == null) { setShowOnlySelectedNodes(false) @@ -290,8 +338,6 @@ function LineageLayoutBase< React.useEffect(() => { if (currentNode?.id) { setSelectedNodeId(currentNode.id) - } else if (selectedNodeId) { - // setSelectedNodeId(selectedNodeId); } else { const node = nodes.length > 0 ? nodes[nodes.length - 1] : null @@ -332,8 +378,10 @@ function LineageLayoutBase< edges={edges} nodeTypes={nodeTypes} edgeTypes={edgeTypes} - nodesDraggable={false} - nodesConnectable={false} + onNodesChange={onNodesChange} + onEdgesChange={onEdgesChange} + nodesDraggable={nodesDraggable} + nodesConnectable={nodesConnectable} zoomOnDoubleClick={false} panOnScroll={true} zoomOnScroll={true} diff --git a/web/common/src/components/Lineage/stories/Lineage.stories.tsx b/web/common/src/components/Lineage/stories/Lineage.stories.tsx index 4ad8ca9f8b..6e16bed61e 100644 --- a/web/common/src/components/Lineage/stories/Lineage.stories.tsx +++ b/web/common/src/components/Lineage/stories/Lineage.stories.tsx @@ -17,33 +17,56 @@ export const LineageModel = () => { > { const [zoom, setZoom] = React.useState(ZOOM_THRESHOLD) const [isBuildingLayout, setIsBuildingLayout] = React.useState(false) + const [nodesDraggable, setNodesDraggable] = React.useState(false) const [edges, setEdges] = React.useState< LineageEdge[] >([]) @@ -388,6 +389,7 @@ export const ModelLineage = ({ nodeTypes={nodeTypes} edgeTypes={edgeTypes} className={className} + nodesDraggable={nodesDraggable} controls={ <> + setNodesDraggable(prev => !prev)} + disabled={isBuildingLayout} + > + + } />