From 2d485cc43599cfb96f7d02c668f94f791824568d Mon Sep 17 00:00:00 2001 From: Max Mykal Date: Mon, 6 Oct 2025 10:05:07 -0700 Subject: [PATCH] clean up lineage layout --- .../src/components/Lineage/LineageContext.ts | 4 - .../src/components/Lineage/LineageLayout.tsx | 404 +----------------- .../components/Lineage/LineageLayoutBase.tsx | 367 ++++++++++++++++ .../Lineage/LineageLayoutContainer.tsx | 43 ++ .../Lineage/stories/ModelLineage.tsx | 2 - .../components/VirtualList/FilterableList.tsx | 3 + 6 files changed, 432 insertions(+), 391 deletions(-) create mode 100644 web/common/src/components/Lineage/LineageLayoutBase.tsx create mode 100644 web/common/src/components/Lineage/LineageLayoutContainer.tsx diff --git a/web/common/src/components/Lineage/LineageContext.ts b/web/common/src/components/Lineage/LineageContext.ts index 6f4ee7e165..9da54dcbee 100644 --- a/web/common/src/components/Lineage/LineageContext.ts +++ b/web/common/src/components/Lineage/LineageContext.ts @@ -30,8 +30,6 @@ export interface LineageContextValue< setSelectedNodeId: React.Dispatch> // Layout - isBuildingLayout: boolean - setIsBuildingLayout: React.Dispatch> zoom: number setZoom: React.Dispatch> @@ -66,8 +64,6 @@ export function getInitial< nodes: [], nodesMap: {}, setNodesMap: () => {}, - isBuildingLayout: false, - setIsBuildingLayout: () => {}, currentNode: null, } } diff --git a/web/common/src/components/Lineage/LineageLayout.tsx b/web/common/src/components/Lineage/LineageLayout.tsx index 2ad31c4b9e..e01e8ae9e9 100644 --- a/web/common/src/components/Lineage/LineageLayout.tsx +++ b/web/common/src/components/Lineage/LineageLayout.tsx @@ -1,52 +1,25 @@ import { - Background, - BackgroundVariant, - Controls, - type EdgeChange, type EdgeTypes, - type NodeChange, type NodeTypes, - ReactFlow, ReactFlowProvider, type SetCenter, - getConnectedEdges, - getIncomers, - getOutgoers, - useReactFlow, - useViewport, - applyNodeChanges, - applyEdgeChanges, } from '@xyflow/react' -import '@xyflow/react/dist/style.css' -import './Lineage.css' - -import { debounce } from 'lodash' -import { CircuitBoard, Crosshair, LocateFixed, RotateCcw } from 'lucide-react' import React from 'react' -import { cn } from '@/utils' import { type LineageContextHook } from './LineageContext' -import { LineageControlButton } from './LineageControlButton' -import { LineageControlIcon } from './LineageControlIcon' + import { - DEFAULT_ZOOM, - type LineageEdge, type LineageEdgeData, type LineageNode, type LineageNodeData, - MAX_ZOOM, - MIN_ZOOM, - NODES_TRESHOLD, - NODES_TRESHOLD_ZOOM, type NodeId, type EdgeId, - ZOOM_THRESHOLD, type PortId, } from './utils' -import { VerticalContainer } from '../VerticalContainer/VerticalContainer' -import { MessageContainer } from '../MessageContainer/MessageContainer' -import { LoadingContainer } from '../LoadingContainer/LoadingContainer' + +import { LineageLayoutBase } from './LineageLayoutBase' +import { LineageLayoutContainer } from './LineageLayoutContainer' export function LineageLayout< TNodeData extends LineageNodeData = LineageNodeData, @@ -61,6 +34,7 @@ export function LineageLayout< controls, nodesDraggable, nodesConnectable, + isBuildingLayout, useLineage, onNodeClick, onNodeDoubleClick, @@ -72,6 +46,7 @@ export function LineageLayout< TEdgeID, TPortID > + isBuildingLayout?: boolean nodeTypes?: NodeTypes edgeTypes?: EdgeTypes className?: string @@ -91,360 +66,19 @@ export function LineageLayout< }) { return ( - + + + ) } - -function LineageLayoutBase< - TNodeData extends LineageNodeData = LineageNodeData, - TEdgeData extends LineageEdgeData = LineageEdgeData, - TNodeID extends string = NodeId, - TEdgeID extends string = EdgeId, - TPortID extends string = PortId, ->({ - nodeTypes, - edgeTypes, - className, - controls, - nodesDraggable = false, - nodesConnectable = false, - useLineage, - onNodeClick, - onNodeDoubleClick, -}: { - useLineage: LineageContextHook< - TNodeData, - TEdgeData, - TNodeID, - TEdgeID, - TPortID - > - nodesDraggable?: boolean - nodesConnectable?: boolean - nodeTypes?: NodeTypes - edgeTypes?: EdgeTypes - className?: string - controls?: - | React.ReactNode - | (({ setCenter }: { setCenter: SetCenter }) => React.ReactNode) - onNodeClick?: ( - event: React.MouseEvent, - node: LineageNode, - ) => void - onNodeDoubleClick?: ( - event: React.MouseEvent, - node: LineageNode, - ) => void -}) { - const { zoom: viewportZoom } = useViewport() - const { setCenter } = useReactFlow() - - const { - isBuildingLayout, - currentNode, - zoom, - nodes: initialNodes, - edges: initialEdges, - nodesMap, - showOnlySelectedNodes, - selectedNodeId, - setZoom, - setSelectedNodeId, - setShowOnlySelectedNodes, - setSelectedNodes, - 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( - (zoom: number = DEFAULT_ZOOM) => { - if (currentNode) { - setCenter(currentNode.position.x, currentNode.position.y, { - zoom, - duration: 0, - }) - } - }, - [currentNode, setCenter], - ) - - const zoomToSelectedNode = React.useCallback( - (zoom: number = DEFAULT_ZOOM) => { - const node = selectedNodeId ? nodesMap[selectedNodeId] : null - if (node) { - setCenter(node.position.x, node.position.y, { - zoom, - duration: 0, - }) - } - }, - [nodesMap, selectedNodeId, setCenter], - ) - - const getAllIncomers = React.useCallback( - ( - node: LineageNode, - visited: Set = new Set(), - ): LineageNode[] => { - if (visited.has(node.id)) return [] - - visited.add(node.id) - - return Array.from( - new Set>([ - node, - ...getIncomers(node, nodes, edges) - .map(n => getAllIncomers(n, visited)) - .flat(), - ]), - ) - }, - [nodes, edges], - ) - - const getAllOutgoers = React.useCallback( - ( - node: LineageNode, - visited: Set = new Set(), - ): LineageNode[] => { - if (visited.has(node.id)) return [] - - visited.add(node.id) - - return Array.from( - new Set>([ - node, - ...getOutgoers(node, nodes, edges) - .map(n => getAllOutgoers(n, visited)) - .flat(), - ]), - ) - }, - [nodes, edges], - ) - - React.useEffect(() => { - setNodes(initialNodes) - }, [initialNodes]) - - React.useEffect(() => { - setEdges(initialEdges) - }, [initialEdges]) - - React.useEffect(() => { - if (selectedNodeId == null) { - setShowOnlySelectedNodes(false) - setSelectedNodes(new Set()) - setSelectedEdges(new Set()) - - return - } - - const node = selectedNodeId ? nodesMap[selectedNodeId] : null - - if (node == null) { - setSelectedNodeId(null) - return - } - - const incomers = getAllIncomers(node) - const outgoers = getAllOutgoers(node) - const connectedNodes = [...incomers, ...outgoers] - - if (currentNode) { - connectedNodes.push(currentNode) - } - - const connectedEdges = getConnectedEdges< - LineageNode, - LineageEdge - >(connectedNodes, edges) - const selectedNodes = new Set(connectedNodes.map(node => node.id)) - const selectedEdges = new Set( - connectedEdges.reduce((acc, edge) => { - if ([edge.source, edge.target].every(id => selectedNodes.has(id))) { - edge.zIndex = 2 - acc.add(edge.id) - } else { - edge.zIndex = 1 - } - return acc - }, new Set()), - ) - - setSelectedNodes(selectedNodes) - setSelectedEdges(selectedEdges) - }, [ - currentNode, - selectedNodeId, - setSelectedNodes, - setSelectedEdges, - getAllIncomers, - getAllOutgoers, - setShowOnlySelectedNodes, - setSelectedNodeId, - ]) - - React.useEffect(() => { - if (selectedNodeId) { - zoomToSelectedNode(zoom) - } else { - zoomToCurrentNode(zoom) - } - }, [zoomToCurrentNode, zoomToSelectedNode]) - - React.useEffect(() => { - updateZoom(viewportZoom) - }, [updateZoom, viewportZoom]) - - React.useEffect(() => { - if (currentNode?.id) { - setSelectedNodeId(currentNode.id) - } else { - const node = nodes.length > 0 ? nodes[nodes.length - 1] : null - - if (node) { - setCenter(node.position.x, node.position.y, { - zoom: zoom, - duration: 0, - }) - } - } - }, [currentNode?.id, setSelectedNodeId, setCenter]) - - return ( - - {isBuildingLayout && ( - - - Building layout... - - - )} - , - LineageEdge - > - className="shrink-0" - nodes={nodes} - edges={edges} - nodeTypes={nodeTypes} - edgeTypes={edgeTypes} - onNodesChange={onNodesChange} - onEdgesChange={onEdgesChange} - nodesDraggable={nodesDraggable} - nodesConnectable={nodesConnectable} - zoomOnDoubleClick={false} - panOnScroll={true} - zoomOnScroll={true} - minZoom={nodes.length > NODES_TRESHOLD ? NODES_TRESHOLD_ZOOM : MIN_ZOOM} - maxZoom={MAX_ZOOM} - fitView={false} - nodeOrigin={[0.5, 0.5]} - onlyRenderVisibleElements - onNodeClick={onNodeClick} - onNodeDoubleClick={onNodeDoubleClick} - > - {zoom > ZOOM_THRESHOLD && ( - - )} - - {currentNode && ( - zoomToCurrentNode(DEFAULT_ZOOM)} - disabled={isBuildingLayout} - > - - - )} - {selectedNodeId && ( - <> - setShowOnlySelectedNodes(!showOnlySelectedNodes)} - disabled={isBuildingLayout} - > - - - zoomToSelectedNode(DEFAULT_ZOOM)} - disabled={isBuildingLayout} - > - - - - )} - {controls && typeof controls === 'function' - ? controls({ setCenter }) - : controls} - - - - ) -} diff --git a/web/common/src/components/Lineage/LineageLayoutBase.tsx b/web/common/src/components/Lineage/LineageLayoutBase.tsx new file mode 100644 index 0000000000..af47a82b29 --- /dev/null +++ b/web/common/src/components/Lineage/LineageLayoutBase.tsx @@ -0,0 +1,367 @@ +import { + Background, + BackgroundVariant, + Controls, + type EdgeChange, + type EdgeTypes, + type NodeChange, + type NodeTypes, + ReactFlow, + type SetCenter, + getConnectedEdges, + getIncomers, + getOutgoers, + useReactFlow, + useViewport, + applyNodeChanges, + applyEdgeChanges, +} from '@xyflow/react' + +import '@xyflow/react/dist/style.css' +import './Lineage.css' + +import { debounce } from 'lodash' +import { CircuitBoard, Crosshair, LocateFixed, RotateCcw } from 'lucide-react' +import React from 'react' + +import { type LineageContextHook } from './LineageContext' +import { LineageControlButton } from './LineageControlButton' +import { LineageControlIcon } from './LineageControlIcon' +import { + DEFAULT_ZOOM, + type LineageEdge, + type LineageEdgeData, + type LineageNode, + type LineageNodeData, + MAX_ZOOM, + MIN_ZOOM, + NODES_TRESHOLD, + NODES_TRESHOLD_ZOOM, + type NodeId, + type EdgeId, + ZOOM_THRESHOLD, + type PortId, +} from './utils' + +import '@xyflow/react/dist/style.css' +import './Lineage.css' +import { cn } from '@/utils' + +export function LineageLayoutBase< + TNodeData extends LineageNodeData = LineageNodeData, + TEdgeData extends LineageEdgeData = LineageEdgeData, + TNodeID extends string = NodeId, + TEdgeID extends string = EdgeId, + TPortID extends string = PortId, +>({ + nodeTypes, + edgeTypes, + className, + controls, + nodesDraggable = false, + nodesConnectable = false, + useLineage, + onNodeClick, + onNodeDoubleClick, +}: { + useLineage: LineageContextHook< + TNodeData, + TEdgeData, + TNodeID, + TEdgeID, + TPortID + > + nodesDraggable?: boolean + nodesConnectable?: boolean + nodeTypes?: NodeTypes + edgeTypes?: EdgeTypes + className?: string + controls?: + | React.ReactNode + | (({ setCenter }: { setCenter: SetCenter }) => React.ReactNode) + onNodeClick?: ( + event: React.MouseEvent, + node: LineageNode, + ) => void + onNodeDoubleClick?: ( + event: React.MouseEvent, + node: LineageNode, + ) => void +}) { + const { zoom: viewportZoom } = useViewport() + const { setCenter } = useReactFlow() + + const { + currentNode, + zoom, + nodes: initialNodes, + edges: initialEdges, + nodesMap, + showOnlySelectedNodes, + selectedNodeId, + setZoom, + setSelectedNodeId, + setShowOnlySelectedNodes, + setSelectedNodes, + 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( + (zoom: number = DEFAULT_ZOOM) => { + if (currentNode) { + setCenter(currentNode.position.x, currentNode.position.y, { + zoom, + duration: 0, + }) + } + }, + [currentNode, setCenter], + ) + + const zoomToSelectedNode = React.useCallback( + (zoom: number = DEFAULT_ZOOM) => { + const node = selectedNodeId ? nodesMap[selectedNodeId] : null + if (node) { + setCenter(node.position.x, node.position.y, { + zoom, + duration: 0, + }) + } + }, + [nodesMap, selectedNodeId, setCenter], + ) + + const getAllIncomers = React.useCallback( + ( + node: LineageNode, + visited: Set = new Set(), + ): LineageNode[] => { + if (visited.has(node.id)) return [] + + visited.add(node.id) + + return Array.from( + new Set>([ + node, + ...getIncomers(node, nodes, edges) + .map(n => getAllIncomers(n, visited)) + .flat(), + ]), + ) + }, + [nodes, edges], + ) + + const getAllOutgoers = React.useCallback( + ( + node: LineageNode, + visited: Set = new Set(), + ): LineageNode[] => { + if (visited.has(node.id)) return [] + + visited.add(node.id) + + return Array.from( + new Set>([ + node, + ...getOutgoers(node, nodes, edges) + .map(n => getAllOutgoers(n, visited)) + .flat(), + ]), + ) + }, + [nodes, edges], + ) + + React.useEffect(() => { + setNodes(initialNodes) + }, [initialNodes]) + + React.useEffect(() => { + setEdges(initialEdges) + }, [initialEdges]) + + React.useEffect(() => { + if (selectedNodeId == null) { + setShowOnlySelectedNodes(false) + setSelectedNodes(new Set()) + setSelectedEdges(new Set()) + + return + } + + const node = selectedNodeId ? nodesMap[selectedNodeId] : null + + if (node == null) { + setSelectedNodeId(null) + return + } + + const incomers = getAllIncomers(node) + const outgoers = getAllOutgoers(node) + const connectedNodes = [...incomers, ...outgoers] + + if (currentNode) { + connectedNodes.push(currentNode) + } + + const connectedEdges = getConnectedEdges< + LineageNode, + LineageEdge + >(connectedNodes, edges) + const selectedNodes = new Set(connectedNodes.map(node => node.id)) + const selectedEdges = new Set( + connectedEdges.reduce((acc, edge) => { + if ([edge.source, edge.target].every(id => selectedNodes.has(id))) { + edge.zIndex = 2 + acc.add(edge.id) + } else { + edge.zIndex = 1 + } + return acc + }, new Set()), + ) + + setSelectedNodes(selectedNodes) + setSelectedEdges(selectedEdges) + }, [ + currentNode, + selectedNodeId, + setSelectedNodes, + setSelectedEdges, + getAllIncomers, + getAllOutgoers, + setShowOnlySelectedNodes, + setSelectedNodeId, + ]) + + React.useEffect(() => { + if (selectedNodeId) { + zoomToSelectedNode(zoom) + } else { + zoomToCurrentNode(zoom) + } + }, [zoomToCurrentNode, zoomToSelectedNode]) + + React.useEffect(() => { + updateZoom(viewportZoom) + }, [updateZoom, viewportZoom]) + + React.useEffect(() => { + if (currentNode?.id) { + setSelectedNodeId(currentNode.id) + } else { + const node = nodes.length > 0 ? nodes[nodes.length - 1] : null + + if (node) { + setCenter(node.position.x, node.position.y, { + zoom: zoom, + duration: 0, + }) + } + } + }, [currentNode?.id, setSelectedNodeId, setCenter]) + + return ( + , + LineageEdge + > + className={cn('shrink-0', className)} + nodes={nodes} + edges={edges} + nodeTypes={nodeTypes} + edgeTypes={edgeTypes} + onNodesChange={onNodesChange} + onEdgesChange={onEdgesChange} + nodesDraggable={nodesDraggable} + nodesConnectable={nodesConnectable} + zoomOnDoubleClick={false} + panOnScroll={true} + zoomOnScroll={true} + minZoom={nodes.length > NODES_TRESHOLD ? NODES_TRESHOLD_ZOOM : MIN_ZOOM} + maxZoom={MAX_ZOOM} + fitView={false} + nodeOrigin={[0.5, 0.5]} + onlyRenderVisibleElements + onNodeClick={onNodeClick} + onNodeDoubleClick={onNodeDoubleClick} + > + {zoom > ZOOM_THRESHOLD && ( + + )} + + {currentNode && ( + zoomToCurrentNode(DEFAULT_ZOOM)} + > + + + )} + {selectedNodeId && ( + <> + setShowOnlySelectedNodes(!showOnlySelectedNodes)} + > + + + zoomToSelectedNode(DEFAULT_ZOOM)} + > + + + + )} + {controls && typeof controls === 'function' + ? controls({ setCenter }) + : controls} + + + ) +} diff --git a/web/common/src/components/Lineage/LineageLayoutContainer.tsx b/web/common/src/components/Lineage/LineageLayoutContainer.tsx new file mode 100644 index 0000000000..e3385a3294 --- /dev/null +++ b/web/common/src/components/Lineage/LineageLayoutContainer.tsx @@ -0,0 +1,43 @@ +import { cn } from '@/utils' + +import React from 'react' + +import { VerticalContainer } from '../VerticalContainer/VerticalContainer' +import { MessageContainer } from '../MessageContainer/MessageContainer' +import { LoadingContainer } from '../LoadingContainer/LoadingContainer' + +export function LineageLayoutContainer({ + isBuildingLayout, + loadingMessage = 'Building layout...', + className, + children, +}: { + isBuildingLayout?: boolean + loadingMessage?: string + className?: string + children: React.ReactNode +}) { + return ( + + {isBuildingLayout && ( + + + {loadingMessage} + + + )} + {children} + + ) +} diff --git a/web/common/src/components/Lineage/stories/ModelLineage.tsx b/web/common/src/components/Lineage/stories/ModelLineage.tsx index 2902350919..800f292c4b 100644 --- a/web/common/src/components/Lineage/stories/ModelLineage.tsx +++ b/web/common/src/components/Lineage/stories/ModelLineage.tsx @@ -359,7 +359,6 @@ export const ModelLineage = ({ selectedNodes, selectedEdges, selectedNodeId, - isBuildingLayout, zoom, edges, nodes, @@ -372,7 +371,6 @@ export const ModelLineage = ({ setSelectedNodes, setSelectedEdges, setSelectedNodeId, - setIsBuildingLayout, setZoom, setEdges, setNodesMap, diff --git a/web/common/src/components/VirtualList/FilterableList.tsx b/web/common/src/components/VirtualList/FilterableList.tsx index d22bfca784..16a243eb0e 100644 --- a/web/common/src/components/VirtualList/FilterableList.tsx +++ b/web/common/src/components/VirtualList/FilterableList.tsx @@ -58,6 +58,9 @@ export function FilterableList({ } inputSize="xs" className="FilterableList__Input w-full" + onClick={(e: React.MouseEvent) => { + e.stopPropagation() + }} />