diff --git a/vscode/extension/tests/lineage_settings.spec.ts b/vscode/extension/tests/lineage_settings.spec.ts index c3237f13dc..c2bd4a1e17 100644 --- a/vscode/extension/tests/lineage_settings.spec.ts +++ b/vscode/extension/tests/lineage_settings.spec.ts @@ -1,4 +1,5 @@ import { test, expect } from './fixtures' +import type { FrameLocator, Page } from '@playwright/test' import fs from 'fs-extra' import { openLineageView, @@ -8,6 +9,31 @@ import { } from './utils' import { createPythonInterpreterSettingsSpecifier } from './utils_code_server' +/** + * Find the iframe that hosts the lineage UI (the one containing the + * Settings cog button). Returns null if it can't be located. + */ +async function findLineageFrame(page: Page): Promise { + const iframes = page.locator('iframe') + const iframeCount = await iframes.count() + + for (let i = 0; i < iframeCount; i++) { + const contentFrame = iframes.nth(i).contentFrame() + if (!contentFrame) continue + const activeFrame = contentFrame.locator('#active-frame').contentFrame() + if (!activeFrame) continue + try { + await activeFrame + .getByRole('button', { name: 'Settings' }) + .waitFor({ timeout: 1000 }) + return activeFrame + } catch { + continue + } + } + return null +} + test('Settings button is visible in the lineage view', async ({ page, sharedCodeServer, @@ -35,30 +61,64 @@ test('Settings button is visible in the lineage view', async ({ // Open lineage await openLineageView(page) - const iframes = page.locator('iframe') - const iframeCount = await iframes.count() - let settingsCount = 0 + const lineageFrame = await findLineageFrame(page) + expect(lineageFrame).not.toBeNull() +}) - for (let i = 0; i < iframeCount; i++) { - const iframe = iframes.nth(i) - const contentFrame = iframe.contentFrame() - if (contentFrame) { - const activeFrame = contentFrame.locator('#active-frame').contentFrame() - if (activeFrame) { - try { - await activeFrame - .getByRole('button', { - name: 'Settings', - }) - .waitFor({ timeout: 1000 }) - settingsCount++ - } catch { - // Continue to next iframe if this one doesn't have the error - continue - } - } - } - } +test('Only Direct Neighbors toggle filters the lineage graph', async ({ + page, + sharedCodeServer, + tempDir, +}) => { + await fs.copy(SUSHI_SOURCE_PATH, tempDir) + await createPythonInterpreterSettingsSpecifier(tempDir) - expect(settingsCount).toBeGreaterThan(0) + await openServerPage(page, tempDir, sharedCodeServer) + await page.waitForSelector('text=models') + + await page + .getByRole('treeitem', { name: 'models', exact: true }) + .locator('a') + .click() + await page + .getByRole('treeitem', { name: 'waiters.py', exact: true }) + .locator('a') + .click() + await waitForLoadedSQLMesh(page) + + await openLineageView(page) + + const lineageFrame = await findLineageFrame(page) + expect(lineageFrame).not.toBeNull() + if (!lineageFrame) return + + // Wait for the graph to render at least one node + await lineageFrame.locator('.react-flow__node').first().waitFor() + const nodesBefore = await lineageFrame.locator('.react-flow__node').count() + expect(nodesBefore).toBeGreaterThan(0) + + // Open the settings menu and toggle "Only Direct Neighbors" + await lineageFrame.getByRole('button', { name: 'Settings' }).click() + const toggle = lineageFrame.getByRole('button', { + name: 'Only Direct Neighbors', + }) + await toggle.waitFor() + await toggle.click() + + // After enabling, the visible node set must be a subset of the original. + // We assert a strict drop only when the original graph had room to shrink + // (i.e. more than the worst-case direct-neighbor count of 1 + parents + children). + await page.waitForTimeout(250) // let React Flow re-layout + const nodesAfter = await lineageFrame.locator('.react-flow__node').count() + expect(nodesAfter).toBeLessThanOrEqual(nodesBefore) + expect(nodesAfter).toBeGreaterThan(0) // main node is always shown + + // Toggle off → graph returns to the full size + await lineageFrame.getByRole('button', { name: 'Settings' }).click() + await lineageFrame + .getByRole('button', { name: 'Only Direct Neighbors' }) + .click() + await page.waitForTimeout(250) + const nodesRestored = await lineageFrame.locator('.react-flow__node').count() + expect(nodesRestored).toBe(nodesBefore) }) diff --git a/vscode/react/src/components/graph/ModelLineage.tsx b/vscode/react/src/components/graph/ModelLineage.tsx index 3d157d3869..df8745140c 100644 --- a/vscode/react/src/components/graph/ModelLineage.tsx +++ b/vscode/react/src/components/graph/ModelLineage.tsx @@ -209,6 +209,8 @@ function ModelColumnLineage(): JSX.Element { withConnected, withImpacted, withSecondary, + withOnlyDirect, + directNeighbors, hasBackground, activeEdges, connectedNodes, @@ -217,6 +219,7 @@ function ModelColumnLineage(): JSX.Element { handleError, setActiveNodes, setWithColumns, + setWithOnlyDirect, } = useLineageFlow() const { setCenter } = useReactFlow() @@ -252,6 +255,8 @@ function ModelColumnLineage(): JSX.Element { withConnected, withImpacted, withSecondary, + withOnlyDirect, + directNeighbors, ) const newEdges = getUpdatedEdges( allEdges, @@ -264,6 +269,8 @@ function ModelColumnLineage(): JSX.Element { withConnected, withImpacted, withSecondary, + withOnlyDirect, + directNeighbors, ) const createLayout = createGraphLayout({ nodesMap, @@ -324,6 +331,8 @@ function ModelColumnLineage(): JSX.Element { withConnected, withImpacted, withSecondary, + withOnlyDirect, + directNeighbors, ) const newEdges = getUpdatedEdges( @@ -337,6 +346,8 @@ function ModelColumnLineage(): JSX.Element { withConnected, withImpacted, withSecondary, + withOnlyDirect, + directNeighbors, ) setEdges(newEdges) @@ -353,6 +364,8 @@ function ModelColumnLineage(): JSX.Element { withConnected, withImpacted, withSecondary, + withOnlyDirect, + directNeighbors, withColumns, mainNode, ]) @@ -395,6 +408,8 @@ function ModelColumnLineage(): JSX.Element { void + withOnlyDirect: boolean + onWithOnlyDirectChange: (value: boolean) => void } +const itemClass = clsx( + 'group flex w-full items-center px-2 py-1 text-sm', + 'text-[var(--vscode-button-foreground)]', + 'hover:bg-[var(--vscode-button-background)] bg-[var(--vscode-button-hoverBackground)]', +) + export function SettingsControl({ showColumns, onWithColumnsChange, + withOnlyDirect, + onWithOnlyDirectChange, }: SettingsControlProps): JSX.Element { return ( onWithColumnsChange(!showColumns)} > Show Columns @@ -44,6 +50,19 @@ export function SettingsControl({ /> )} + onWithOnlyDirectChange(!withOnlyDirect)} + > + Only Direct Neighbors + {withOnlyDirect && ( + ) diff --git a/vscode/react/src/components/graph/context.tsx b/vscode/react/src/components/graph/context.tsx index 9ab4f0722e..274c6dc1e8 100644 --- a/vscode/react/src/components/graph/context.tsx +++ b/vscode/react/src/components/graph/context.tsx @@ -7,6 +7,7 @@ import { } from 'react' import { getNodeMap, hasActiveEdge, hasActiveEdgeConnector } from './help' import { type Node } from 'reactflow' +import { isNil } from '@/utils/index' import type { Lineage } from '@/domain/lineage' import type { ModelSQLMeshModel } from '@/domain/sqlmesh-model' import type { Column } from '@/domain/column' @@ -43,6 +44,8 @@ interface LineageFlow { hasBackground: boolean withImpacted: boolean withSecondary: boolean + withOnlyDirect: boolean + directNeighbors: Set manuallySelectedColumn?: [ModelSQLMeshModel, Column] highlightedNodes: HighlightedNodes nodesMap: Record @@ -55,6 +58,7 @@ interface LineageFlow { setHasBackground: React.Dispatch> setWithImpacted: React.Dispatch> setWithSecondary: React.Dispatch> + setWithOnlyDirect: React.Dispatch> setConnections: React.Dispatch>> hasActiveEdge: (edge: [string | undefined, string | undefined]) => boolean addActiveEdges: (edges: Array<[string, string]>) => void @@ -85,6 +89,8 @@ export const LineageFlowContext = createContext({ withConnected: false, withImpacted: true, withSecondary: false, + withOnlyDirect: false, + directNeighbors: new Set(), hasBackground: true, mainNode: undefined, activeEdges: new Map(), @@ -103,6 +109,7 @@ export const LineageFlowContext = createContext({ setWithImpacted: () => false, setWithSecondary: () => false, setWithConnected: () => false, + setWithOnlyDirect: () => false, hasActiveEdge: () => false, addActiveEdges: () => {}, removeActiveEdges: () => {}, @@ -161,6 +168,7 @@ export default function LineageFlowProvider({ const [hasBackground, setHasBackground] = useState(true) const [withImpacted, setWithImpacted] = useState(true) const [withSecondary, setWithSecondary] = useState(false) + const [withOnlyDirect, setWithOnlyDirect] = useState(false) const nodesMap = useMemo( () => @@ -264,6 +272,39 @@ export default function LineageFlowProvider({ [nodesConnections], ) + // Reverse adjacency index: parent -> children. Built once per `lineage` + // change so per-model `directNeighbors` lookups stay O(parents + children) + // instead of scanning the whole graph on every mainNode switch. + const childrenByParent = useMemo(() => { + const map = new Map() + for (const [child, info] of Object.entries(lineage) as Array< + [ModelEncodedFQN, Lineage] + >) { + for (const parent of info?.models ?? []) { + const existing = map.get(parent) + if (existing) { + existing.push(child) + } else { + map.set(parent, [child]) + } + } + } + return map + }, [lineage]) + + const directNeighbors = useMemo(() => { + const set = new Set() + if (isNil(mainNode)) return set + set.add(mainNode) + for (const parent of lineage[mainNode]?.models ?? []) { + set.add(parent) + } + for (const child of childrenByParent.get(mainNode) ?? []) { + set.add(child) + } + return set + }, [mainNode, lineage, childrenByParent]) + const selectedEdges = useMemo( () => Array.from(selectedNodes) @@ -292,6 +333,8 @@ export default function LineageFlowProvider({ withConnected, withImpacted, withSecondary, + withOnlyDirect, + directNeighbors, showControls, hasBackground, nodesMap, @@ -304,6 +347,7 @@ export default function LineageFlowProvider({ setWithConnected, setWithImpacted, setWithSecondary, + setWithOnlyDirect, setHasBackground, setSelectedNodes, setMainNode, diff --git a/vscode/react/src/components/graph/help.ts b/vscode/react/src/components/graph/help.ts index 93e5c4db45..8b9c8881f9 100644 --- a/vscode/react/src/components/graph/help.ts +++ b/vscode/react/src/components/graph/help.ts @@ -562,6 +562,8 @@ export function getUpdatedEdges( withConnected: boolean = false, withImpacted: boolean = false, withSecondary: boolean = false, + withOnlyDirect: boolean = false, + directNeighbors: Set = new Set(), ): Edge[] { const tempEdges = edges.map(edge => { const isActiveEdge = hasActiveEdge(activeEdges, [ @@ -593,10 +595,15 @@ export function getUpdatedEdges( hasEdge(selectedEdges, edge.id) && activeNodes.has(edge.source) && activeNodes.has(edge.target) + const shouldHideNonDirect = + withOnlyDirect && + (isFalse(directNeighbors.has(edge.source)) || + isFalse(directNeighbors.has(edge.target))) if ( isFalse(shouldHideImpacted) && isFalse(shouldHideSecondary) && + isFalse(shouldHideNonDirect) && (isFalse(hasSelections) || isVisibleEdge) ) { edge.hidden = false @@ -653,6 +660,8 @@ export function getUpdatedNodes( withConnected: boolean, withImpacted: boolean, withSecondary: boolean, + withOnlyDirect: boolean = false, + directNeighbors: Set = new Set(), ): Node[] { return nodes.map(node => { node.hidden = true @@ -667,8 +676,14 @@ export function getUpdatedNodes( isFalse(withSecondary) && isFalse(hasSelections) const shouldHideSecondary = isSecondaryNode && withoutSecondaryNodes const shouldHideImpacted = isImpactedNode && withoutImpactedNodes + const shouldHideNonDirect = + withOnlyDirect && isFalse(directNeighbors.has(node.id)) - if (isFalse(shouldHideImpacted) && isFalse(shouldHideSecondary)) { + if ( + isFalse(shouldHideImpacted) && + isFalse(shouldHideSecondary) && + isFalse(shouldHideNonDirect) + ) { node.hidden = isFalse(isActiveNode) }