lineage: fix expand collapse operation on nodes (#21309)

This commit is contained in:
Karan Hotchandani 2025-06-12 23:09:20 +05:30 committed by GitHub
parent e55c836cb6
commit 074c2efa10
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
11 changed files with 254 additions and 89 deletions

View File

@ -13,7 +13,9 @@
import test, { expect } from '@playwright/test';
import { get } from 'lodash';
import { GlobalSettingOptions } from '../../constant/settings';
import { ContainerClass } from '../../support/entity/ContainerClass';
import { DashboardClass } from '../../support/entity/DashboardClass';
import { MetricClass } from '../../support/entity/MetricClass';
import { MlModelClass } from '../../support/entity/MlModelClass';
import { SearchIndexClass } from '../../support/entity/SearchIndexClass';
import { TableClass } from '../../support/entity/TableClass';
@ -27,6 +29,7 @@ import {
import {
addPipelineBetweenNodes,
fillLineageConfigForm,
performCollapse,
performExpand,
performZoomOut,
verifyColumnLayerActive,
@ -48,6 +51,8 @@ test.describe('Lineage Settings Tests', () => {
const dashboard = new DashboardClass();
const mlModel = new MlModelClass();
const searchIndex = new SearchIndexClass();
const container = new ContainerClass();
const metric = new MetricClass();
try {
await Promise.all([
@ -56,12 +61,16 @@ test.describe('Lineage Settings Tests', () => {
dashboard.create(apiContext),
mlModel.create(apiContext),
searchIndex.create(apiContext),
container.create(apiContext),
metric.create(apiContext),
]);
await addPipelineBetweenNodes(page, table, topic);
await addPipelineBetweenNodes(page, topic, dashboard);
await addPipelineBetweenNodes(page, dashboard, mlModel);
await addPipelineBetweenNodes(page, mlModel, searchIndex);
await addPipelineBetweenNodes(page, searchIndex, container);
await addPipelineBetweenNodes(page, container, metric);
await test.step(
'Lineage config should throw error if upstream depth is less than 0',
@ -168,7 +177,17 @@ test.describe('Lineage Settings Tests', () => {
await verifyNodePresent(page, topic);
await verifyNodePresent(page, mlModel);
await performExpand(page, mlModel, false, searchIndex);
await performExpand(page, topic, true);
await performExpand(page, searchIndex, false, container);
await performExpand(page, container, false, metric);
await performExpand(page, topic, true, table);
// perform collapse
await performCollapse(page, mlModel, false, [
searchIndex,
container,
metric,
]);
await performCollapse(page, dashboard, true, [table, topic]);
}
);

View File

@ -275,6 +275,7 @@ test('Verify column lineage between table and topic', async ({ browser }) => {
await table.visitEntityPage(page);
await visitLineageTab(page);
await activateColumnLayer(page);
await page.click('[data-testid="edit-lineage"]');
await removeColumnLineage(page, sourceCol, targetCol);
@ -389,6 +390,8 @@ test('Verify function data in edge drawer', async ({ browser }) => {
await page.reload();
await lineageReq;
await activateColumnLayer(page);
await page
.locator(
`[data-testid="column-edge-${btoa(sourceColName)}-${btoa(

View File

@ -236,6 +236,30 @@ export const performExpand = async (
}
};
export const performCollapse = async (
page: Page,
node: EntityClass,
upstream: boolean,
hiddenEntity: EntityClass[]
) => {
const nodeFqn = get(node, 'entityResponseData.fullyQualifiedName');
const handleDirection = upstream ? 'left' : 'right';
const collapseBtn = page
.locator(`[data-testid="lineage-node-${nodeFqn}"]`)
.locator(`.react-flow__handle-${handleDirection}`)
.getByTestId('minus-icon');
await collapseBtn.click();
for (const entity of hiddenEntity) {
const hiddenNodeFqn = get(entity, 'entityResponseData.fullyQualifiedName');
const hiddenNode = page.locator(
`[data-testid="lineage-node-${hiddenNodeFqn}"]`
);
await expect(hiddenNode).not.toBeVisible();
}
};
export const verifyNodePresent = async (page: Page, node: EntityClass) => {
const nodeFqn = get(node, 'entityResponseData.fullyQualifiedName');
const name = get(node, 'entityResponseData.name');

View File

@ -208,3 +208,22 @@ export const getColumnContent = (
</div>
);
};
export function getNodeClassNames({
isSelected,
showDqTracing,
isTraced,
}: {
isSelected: boolean;
showDqTracing: boolean;
isTraced: boolean;
}) {
return classNames(
'lineage-node p-0',
isSelected ? 'custom-node-header-active' : 'custom-node-header-normal',
{
'data-quality-failed-custom-node-header': showDqTracing,
'custom-node-header-tracing': isTraced,
}
);
}

View File

@ -11,7 +11,6 @@
* limitations under the License.
*/
import classNames from 'classnames';
import React, { memo, useCallback, useEffect, useMemo, useState } from 'react';
import { Handle, NodeProps, Position } from 'reactflow';
import { useLineageProvider } from '../../../context/LineageProvider/LineageProvider';
@ -20,7 +19,11 @@ import { LineageDirection } from '../../../generated/api/lineage/lineageDirectio
import { LineageLayer } from '../../../generated/configuration/lineageSettings';
import LineageNodeRemoveButton from '../../Lineage/LineageNodeRemoveButton';
import './custom-node.less';
import { getCollapseHandle, getExpandHandle } from './CustomNode.utils';
import {
getCollapseHandle,
getExpandHandle,
getNodeClassNames,
} from './CustomNode.utils';
import './entity-lineage.style.less';
import {
ExpandCollapseHandlesProps,
@ -100,7 +103,8 @@ const ExpandCollapseHandles = memo(
isDownstreamNode,
isUpstreamNode,
isRootNode,
expandPerformed,
upstreamExpandPerformed,
downstreamExpandPerformed,
upstreamLineageLength,
onCollapse,
onExpand,
@ -114,18 +118,21 @@ const ExpandCollapseHandles = memo(
{hasOutgoers &&
(isDownstreamNode || isRootNode) &&
getCollapseHandle(LineageDirection.Downstream, onCollapse)}
{!hasOutgoers &&
!expandPerformed &&
!downstreamExpandPerformed &&
getExpandHandle(LineageDirection.Downstream, () =>
onExpand(LineageDirection.Downstream)
)}
{hasIncomers &&
(isUpstreamNode || isRootNode) &&
getCollapseHandle(LineageDirection.Upstream, () =>
onCollapse(LineageDirection.Upstream)
)}
{!hasIncomers &&
!expandPerformed &&
!upstreamExpandPerformed &&
upstreamLineageLength > 0 &&
getExpandHandle(LineageDirection.Upstream, () =>
onExpand(LineageDirection.Upstream)
@ -166,17 +173,23 @@ const CustomNodeV1 = (props: NodeProps) => {
id,
fullyQualifiedName,
upstreamLineage = [],
expandPerformed = false,
upstreamExpandPerformed = false,
downstreamExpandPerformed = false,
} = node;
const [isTraced, setIsTraced] = useState<boolean>(false);
const [isTraced, setIsTraced] = useState(false);
const showDqTracing = useMemo(() => {
return (
(activeLayer.includes(LineageLayer.DataObservability) &&
dataQualityLineage?.nodes?.some((dqNode) => dqNode.id === id)) ??
false
);
}, [activeLayer, dataQualityLineage, id]);
const showDqTracing = useMemo(
() =>
activeLayer.includes(LineageLayer.DataObservability) &&
dataQualityLineage?.nodes?.some((dqNode) => dqNode.id === id),
[activeLayer, dataQualityLineage, id]
);
const containerClass = getNodeClassNames({
isSelected,
showDqTracing: showDqTracing ?? false,
isTraced,
});
const onExpand = useCallback(
(direction: LineageDirection) => {
@ -197,33 +210,20 @@ const CustomNodeV1 = (props: NodeProps) => {
return label;
}
const renderRemoveBtn =
isSelected && isEditMode && !isRootNode ? (
<LineageNodeRemoveButton onRemove={() => removeNodeHandler(props)} />
) : null;
return (
<>
<LineageNodeLabelV1 node={node} />
{renderRemoveBtn}
{isSelected && isEditMode && !isRootNode && (
<LineageNodeRemoveButton onRemove={() => removeNodeHandler(props)} />
)}
</>
);
}, [node.id, isNewNode, label, isSelected, isEditMode, isRootNode]);
const containerClass = useMemo(() => {
return classNames(
'lineage-node p-0',
isSelected ? 'custom-node-header-active' : 'custom-node-header-normal',
{
'data-quality-failed-custom-node-header': showDqTracing,
'custom-node-header-tracing': isTraced,
}
);
}, [isSelected, showDqTracing, isTraced]);
const expandCollapseProps = useMemo(
const expandCollapseProps = useMemo<ExpandCollapseHandlesProps>(
() => ({
expandPerformed,
upstreamExpandPerformed,
downstreamExpandPerformed,
hasIncomers,
hasOutgoers,
isDownstreamNode,
@ -235,7 +235,8 @@ const CustomNodeV1 = (props: NodeProps) => {
onExpand,
}),
[
expandPerformed,
upstreamExpandPerformed,
downstreamExpandPerformed,
hasIncomers,
hasOutgoers,
isDownstreamNode,
@ -248,6 +249,11 @@ const CustomNodeV1 = (props: NodeProps) => {
]
);
const handlesElement = useMemo(
() => <ExpandCollapseHandles {...expandCollapseProps} />,
[expandCollapseProps]
);
useEffect(() => {
setIsTraced(tracedNodes.includes(id));
}, [tracedNodes, id]);
@ -257,9 +263,7 @@ const CustomNodeV1 = (props: NodeProps) => {
className={containerClass}
data-testid={`lineage-node-${fullyQualifiedName}`}>
<NodeHandles
expandCollapseHandles={
<ExpandCollapseHandles {...expandCollapseProps} />
}
expandCollapseHandles={handlesElement}
id={id}
isConnectable={isConnectable}
nodeType={nodeType}

View File

@ -88,7 +88,8 @@ export interface ExpandCollapseHandlesProps {
isDownstreamNode: boolean;
isUpstreamNode: boolean;
isRootNode: boolean;
expandPerformed: boolean;
upstreamExpandPerformed: boolean;
downstreamExpandPerformed: boolean;
upstreamLineageLength: number;
onCollapse: (direction?: LineageDirection) => void;
onExpand: (direction: LineageDirection) => void;

View File

@ -21,12 +21,7 @@ import React, {
} from 'react';
import { useTranslation } from 'react-i18next';
import { useHistory } from 'react-router-dom';
import ReactFlow, {
Background,
MiniMap,
Panel,
ReactFlowProvider,
} from 'reactflow';
import ReactFlow, { Background, MiniMap, Panel } from 'reactflow';
import {
MAX_ZOOM_VALUE,
MIN_ZOOM_VALUE,
@ -73,7 +68,6 @@ const Lineage = ({
onNodeDrop,
onNodesChange,
onEdgesChange,
entityLineage,
onPaneClick,
onConnect,
onInitReactFlow,
@ -162,7 +156,7 @@ const Lineage = ({
data-testid="lineage-container"
id="lineage-container" // ID is required for export PNG functionality
ref={reactFlowWrapper}>
{entityLineage && (
{init ? (
<>
{isPlatformLineage ? null : (
<CustomControlsComponent className="absolute top-1 right-1 p-xs" />
@ -178,10 +172,6 @@ const Lineage = ({
isFullScreen ? onExitFullScreenViewClick : undefined
}
/>
</>
)}
{init ? (
<ReactFlowProvider>
<ReactFlow
elevateEdgesOnSelect
className="custom-react-flow"
@ -222,7 +212,7 @@ const Lineage = ({
<LineageLayers entity={entity} entityType={entityType} />
</Panel>
</ReactFlow>
</ReactFlowProvider>
</>
) : (
<div className="loading-card">
<Loader />

View File

@ -81,6 +81,7 @@ export interface LineageEntityReference extends EntityReference {
parentId?: string;
childrenLength?: number;
};
expandPerformed?: boolean;
upstreamExpandPerformed?: boolean;
downstreamExpandPerformed?: boolean;
direction?: LineageDirection;
}

View File

@ -17,7 +17,7 @@ import { MOCK_EXPLORE_SEARCH_RESULTS } from '../Explore/Explore.mock';
import Lineage from './Lineage.component';
import { EntityLineageResponse } from './Lineage.interface';
let entityLineage: EntityLineageResponse | undefined = {
const entityLineage: EntityLineageResponse | undefined = {
entity: {
name: 'fact_sale',
fullyQualifiedName: 'sample_data.ecommerce_db.shopify.fact_sale',
@ -90,6 +90,7 @@ jest.mock('../../context/LineageProvider/LineageProvider', () => ({
activeLayer: [],
entityLineage: entityLineage,
updateEntityData: jest.fn(),
init: true,
})),
}));
@ -126,19 +127,7 @@ describe('Lineage', () => {
const customControlsComponent = screen.getByText('Controls Component');
const lineageComponent = screen.getByTestId('lineage-container');
expect(customControlsComponent).toBeInTheDocument();
expect(lineageComponent).toBeInTheDocument();
});
it('does not render CustomControlsComponent when entityLineage is falsy', () => {
const mockPropsWithoutEntity = {
...mockProps,
entity: undefined,
};
entityLineage = undefined;
render(<Lineage {...mockPropsWithoutEntity} />);
const customControlsComponent = screen.getByText('Controls Component');
expect(customControlsComponent).toBeInTheDocument();
});
});

View File

@ -23,6 +23,7 @@ import React, {
useContext,
useEffect,
useMemo,
useRef,
useState,
} from 'react';
import { useTranslation } from 'react-i18next';
@ -55,6 +56,7 @@ import {
EntityLineageResponse,
LineageData,
LineageEntityReference,
NodeData,
} from '../../components/Lineage/Lineage.interface';
import LineageNodeRemoveButton from '../../components/Lineage/LineageNodeRemoveButton';
import { SourceType } from '../../components/SearchedData/SearchedData.interface';
@ -152,6 +154,11 @@ const LineageProvider = ({ children }: LineageProviderProps) => {
{} as SourceType
);
const [activeLayer, setActiveLayer] = useState<LineageLayer[]>([]);
// Added this ref to compare the previous active layer with the current active layer.
// We need to redraw the lineage if the column level lineage is added or removed.
const prevActiveLayerRef = useRef<LineageLayer[]>([]);
const [activeNode, setActiveNode] = useState<Node>();
const [expandAllColumns, setExpandAllColumns] = useState(false);
const [selectedColumn, setSelectedColumn] = useState<string>('');
@ -264,6 +271,7 @@ const LineageProvider = ({ children }: LineageProviderProps) => {
allNodes,
lineageData.edges ?? [],
decodedFqn,
activeLayer.includes(LineageLayer.ColumnLevelLineage),
isFirstTime ? true : undefined
);
@ -365,7 +373,7 @@ const LineageProvider = ({ children }: LineageProviderProps) => {
setLineageData(res);
const { nodes, edges, entity } = parseLineageData(res, '');
const { nodes, edges, entity } = parseLineageData(res, '', decodedFqn);
const updatedEntityLineage = {
nodes,
edges,
@ -409,7 +417,7 @@ const LineageProvider = ({ children }: LineageProviderProps) => {
});
setLineageData(res);
const { nodes, edges, entity } = parseLineageData(res, fqn);
const { nodes, edges, entity } = parseLineageData(res, fqn, decodedFqn);
const updatedEntityLineage = {
nodes,
edges,
@ -510,9 +518,19 @@ const LineageProvider = ({ children }: LineageProviderProps) => {
direction,
});
const currentNodes: Record<string, NodeData> = {};
entityLineage.nodes?.forEach((node) => {
currentNodes[node.fullyQualifiedName ?? ''] = {
entity: node,
paging: (node as LineageEntityReference).paging ?? {
entityDownstreamCount: 0,
entityUpstreamCount: 0,
},
};
});
const concatenatedLineageData = {
nodes: {
...lineageData?.nodes,
...currentNodes,
...res.nodes,
},
downstreamEdges: {
@ -527,20 +545,56 @@ const LineageProvider = ({ children }: LineageProviderProps) => {
const { nodes: newNodes, edges: newEdges } = parseLineageData(
concatenatedLineageData,
node.fullyQualifiedName ?? ''
node.fullyQualifiedName ?? '',
decodedFqn
);
const uniqueNodes = [...(entityLineage.nodes ?? [])];
for (const nNode of newNodes ?? []) {
if (
!uniqueNodes.some(
(n) => n.fullyQualifiedName === nNode.fullyQualifiedName
)
) {
uniqueNodes.push(nNode);
}
}
const updatedEntityLineage = {
entity: entityLineage.entity,
nodes: uniqWith(
[...(entityLineage.nodes ?? []), ...newNodes],
isEqual
),
nodes: uniqueNodes,
edges: uniqWith(
[...(entityLineage.edges ?? []), ...newEdges],
isEqual
),
};
// remove the nodes and edges from the lineageData
const visibleNodes: Record<string, NodeData> = {};
uniqueNodes.forEach((node: EntityReference) => {
visibleNodes[node.fullyQualifiedName ?? ''] = {
entity: node,
paging: (node as LineageEntityReference).paging ?? {
entityDownstreamCount: 0,
entityUpstreamCount: 0,
},
};
});
const currentNode = updatedEntityLineage.nodes.find(
(n) => n.fullyQualifiedName === node.fullyQualifiedName
);
if (currentNode) {
if (direction === LineageDirection.Upstream) {
(currentNode as LineageEntityReference).upstreamExpandPerformed =
true;
} else {
(currentNode as LineageEntityReference).downstreamExpandPerformed =
true;
}
}
updateLineageData(updatedEntityLineage, {
shouldRedraw: true,
centerNode: false,
@ -887,7 +941,8 @@ const LineageProvider = ({ children }: LineageProviderProps) => {
const { nodes: newNodes, edges: newEdges } = parseLineageData(
concatenatedLineageData,
parentNode.data.node.fullyQualifiedName
parentNode.data.node.fullyQualifiedName,
decodedFqn
);
updateLineageData(
@ -1100,7 +1155,12 @@ const LineageProvider = ({ children }: LineageProviderProps) => {
);
const { edges: createdEdges, columnsHavingLineage } =
createEdgesAndEdgeMaps(allNodes, allEdges, decodedFqn);
createEdgesAndEdgeMaps(
allNodes,
allEdges,
decodedFqn,
activeLayer.includes(LineageLayer.ColumnLevelLineage)
);
setEdges(createdEdges);
setColumnsHavingLineage(columnsHavingLineage);
@ -1328,6 +1388,31 @@ const LineageProvider = ({ children }: LineageProviderProps) => {
);
});
// Find the node in updatedNodes by ID and set expandPerformed: false
const currentNodeId = (node as Node).id;
const nodeToUpdate = updatedNodes.find((n) => n.id === currentNodeId);
if (nodeToUpdate) {
if (direction === LineageDirection.Upstream) {
(nodeToUpdate as LineageEntityReference).upstreamExpandPerformed =
false;
} else {
(nodeToUpdate as LineageEntityReference).downstreamExpandPerformed =
false;
}
}
// remove the nodes and edges from the lineageData
const visibleNodes: Record<string, NodeData> = {};
updatedNodes.forEach((node) => {
visibleNodes[node.fullyQualifiedName ?? ''] = {
entity: node,
paging: (node as LineageEntityReference).paging ?? {
entityDownstreamCount: 0,
entityUpstreamCount: 0,
},
};
});
updateLineageData(
{
...entityLineage,
@ -1503,7 +1588,20 @@ const LineageProvider = ({ children }: LineageProviderProps) => {
}, [isEditMode, deletePressed, backspacePressed, activeNode, selectedEdge]);
useEffect(() => {
repositionLayout();
const prevActiveLayer = prevActiveLayerRef.current;
const prevHadColumn = prevActiveLayer.includes(
LineageLayer.ColumnLevelLineage
);
const currHasColumn = activeLayer.includes(LineageLayer.ColumnLevelLineage);
if (prevHadColumn !== currHasColumn) {
redraw();
} else {
repositionLayout();
}
prevActiveLayerRef.current = activeLayer;
}, [activeLayer, expandAllColumns]);
useEffect(() => {

View File

@ -16,7 +16,15 @@ import { graphlib, layout } from '@dagrejs/dagre';
import { AxiosError } from 'axios';
import ELK, { ElkExtendedEdge, ElkNode } from 'elkjs/lib/elk.bundled.js';
import { t } from 'i18next';
import { get, isEmpty, isNil, isUndefined, uniqueId } from 'lodash';
import {
get,
isEmpty,
isEqual,
isNil,
isUndefined,
uniqueId,
uniqWith,
} from 'lodash';
import { EntityTags, LoadingState } from 'Models';
import React, { MouseEvent as ReactMouseEvent } from 'react';
import {
@ -839,6 +847,7 @@ export const createEdgesAndEdgeMaps = (
nodes: EntityReference[],
edges: EdgeDetails[],
entityFqn: string,
isColumnLayerActive: boolean,
hidden?: boolean
) => {
const lineageEdgesV1: Edge[] = [];
@ -851,10 +860,6 @@ export const createEdgesAndEdgeMaps = (
const sourceId = edge.fromEntity.id;
const targetId = edge.toEntity.id;
// Update edge maps for fast lookup
outgoingMap.set(sourceId, (outgoingMap.get(sourceId) ?? 0) + 1);
incomingMap.set(targetId, (incomingMap.get(targetId) ?? 0) + 1);
const sourceType = nodes.find((n) => sourceId === n.id);
const targetType = nodes.find((n) => targetId === n.id);
@ -862,7 +867,11 @@ export const createEdgesAndEdgeMaps = (
return;
}
if (!isUndefined(edge.columns)) {
// Update edge maps for fast lookup
outgoingMap.set(sourceId, (outgoingMap.get(sourceId) ?? 0) + 1);
incomingMap.set(targetId, (incomingMap.get(targetId) ?? 0) + 1);
if (!isUndefined(edge.columns) && isColumnLayerActive) {
edge.columns?.forEach((e) => {
const toColumn = e.toColumn ?? '';
if (toColumn && e.fromColumns?.length) {
@ -1133,7 +1142,7 @@ export const getConnectedNodesEdges = (
return {
nodes: outgoers,
edges: connectedEdges,
edges: uniqWith(connectedEdges, isEqual),
nodeFqn: childNodeFqn,
};
};
@ -1406,9 +1415,16 @@ const processNodeArray = (
.map((node: NodeData) => ({
...node.entity,
paging: node.paging,
expandPerformed:
(node.entity as LineageEntityReference).expandPerformed ||
node.entity.fullyQualifiedName === entityFqn,
upstreamExpandPerformed:
(node.entity as LineageEntityReference).upstreamExpandPerformed !==
undefined
? (node.entity as LineageEntityReference).upstreamExpandPerformed
: node.entity.fullyQualifiedName === entityFqn,
downstreamExpandPerformed:
(node.entity as LineageEntityReference).downstreamExpandPerformed !==
undefined
? (node.entity as LineageEntityReference).downstreamExpandPerformed
: node.entity.fullyQualifiedName === entityFqn,
}))
.flat();
};
@ -1509,7 +1525,8 @@ const processPagination = (
export const parseLineageData = (
data: LineageData,
entityFqn: string
entityFqn: string, // This contains fqn of node or entity that is being viewed in lineage page
rootFqn: string // This contains the fqn of the entity that is being viewed in lineage page
): {
nodes: LineageEntityReference[];
edges: EdgeDetails[];
@ -1518,7 +1535,7 @@ export const parseLineageData = (
const { nodes, downstreamEdges, upstreamEdges } = data;
// Process nodes
const nodesArray = processNodeArray(nodes, entityFqn);
const nodesArray = processNodeArray(nodes, rootFqn);
const processedNodes: LineageEntityReference[] = [...nodesArray];
// Process edges