Change lineage positioning algorithm (#18897)

* use elk algorithm to position nodes

* change positioning

* fix spacing

* do not reset zoom value

* minor pw fix

* force click on lineage edge

openmetadata-ui/src/main/resources/ui/package.json

@@ -85,6 +85,7 @@
     "eventemitter3": "^5.0.1",
     "fast-json-patch": "^3.1.1",
     "history": "4.5.1",
+    "elkjs": "^0.9.3",
     "html-react-parser": "^1.4.14",
     "https-browserify": "^1.0.0",
     "i18next": "^21.10.0",

openmetadata-ui/src/main/resources/ui/playwright/e2e/Pages/Lineage.spec.ts

@@ -92,58 +92,70 @@ for (const EntityClass of entities) {
       defaultEntity
     );
 
-    await test.step('Should create lineage for the entity', async () => {
-      await redirectToHomePage(page);
-      await currentEntity.visitEntityPage(page);
-      await visitLineageTab(page);
-      await verifyColumnLayerInactive(page);
-      await editLineage(page);
-      await performZoomOut(page);
-      for (const entity of entities) {
-        await connectEdgeBetweenNodes(page, currentEntity, entity);
-      }
+    try {
+      await test.step('Should create lineage for the entity', async () => {
+        await redirectToHomePage(page);
+        await currentEntity.visitEntityPage(page);
+        await visitLineageTab(page);
+        await verifyColumnLayerInactive(page);
+        await editLineage(page);
+        await performZoomOut(page);
+        for (const entity of entities) {
+          await connectEdgeBetweenNodes(page, currentEntity, entity);
+        }
 
-      await redirectToHomePage(page);
-      await currentEntity.visitEntityPage(page);
-      await visitLineageTab(page);
-      await page
-        .locator('.react-flow__controls-fitview')
-        .dispatchEvent('click');
+        await redirectToHomePage(page);
+        await currentEntity.visitEntityPage(page);
+        await visitLineageTab(page);
+        await page.click('[data-testid="edit-lineage"]');
+        await page
+          .locator('.react-flow__controls-fitview')
+          .dispatchEvent('click');
 
-      for (const entity of entities) {
-        await verifyNodePresent(page, entity);
-      }
-    });
+        for (const entity of entities) {
+          await verifyNodePresent(page, entity);
+        }
+        await page.click('[data-testid="edit-lineage"]');
+      });
 
-    await test.step('Should create pipeline between entities', async () => {
-      await editLineage(page);
-      await performZoomOut(page);
+      await test.step('Should create pipeline between entities', async () => {
+        await redirectToHomePage(page);
+        await currentEntity.visitEntityPage(page);
+        await visitLineageTab(page);
+        await editLineage(page);
+        await page
+          .locator('.react-flow__controls-fitview')
+          .dispatchEvent('click');
 
-      for (const entity of entities) {
-        await applyPipelineFromModal(page, currentEntity, entity, pipeline);
-      }
-    });
+        for (const entity of entities) {
+          await applyPipelineFromModal(page, currentEntity, entity, pipeline);
+        }
+      });
 
-    await test.step('Verify Lineage Export CSV', async () => {
-      await redirectToHomePage(page);
-      await currentEntity.visitEntityPage(page);
-      await visitLineageTab(page);
-      await verifyExportLineageCSV(page, currentEntity, entities, pipeline);
-    });
+      await test.step('Verify Lineage Export CSV', async () => {
+        await redirectToHomePage(page);
+        await currentEntity.visitEntityPage(page);
+        await visitLineageTab(page);
+        await verifyExportLineageCSV(page, currentEntity, entities, pipeline);
+      });
 
-    await test.step('Remove lineage between nodes for the entity', async () => {
-      await redirectToHomePage(page);
-      await currentEntity.visitEntityPage(page);
-      await visitLineageTab(page);
-      await editLineage(page);
-      await performZoomOut(page);
+      await test.step(
+        'Remove lineage between nodes for the entity',
+        async () => {
+          await redirectToHomePage(page);
+          await currentEntity.visitEntityPage(page);
+          await visitLineageTab(page);
+          await editLineage(page);
+          await performZoomOut(page);
 
-      for (const entity of entities) {
-        await deleteEdge(page, currentEntity, entity);
-      }
-    });
-
-    await cleanup();
+          for (const entity of entities) {
+            await deleteEdge(page, currentEntity, entity);
+          }
+        }
+      );
+    } finally {
+      await cleanup();
+    }
   });
 }
 

openmetadata-ui/src/main/resources/ui/playwright/utils/lineage.ts

@@ -132,7 +132,9 @@ export const dragAndDropNode = async (
   await page.hover(originSelector);
   await page.mouse.down();
   const box = (await destinationElement.boundingBox())!;
-  await page.mouse.move(box.x + box.width / 2, box.y + box.height / 2);
+  const x = (box.x + box.width / 2) * 0.25; // 0.25 as zoom factor
+  const y = (box.y + box.height / 2) * 0.25; // 0.25 as zoom factor
+  await page.mouse.move(x, y);
   await destinationElement.hover();
   await page.mouse.up();
 };
@@ -348,7 +350,8 @@ export const applyPipelineFromModal = async (
 
   await page
     .locator(`[data-testid="edge-${fromNodeFqn}-${toNodeFqn}"]`)
-    .dispatchEvent('click');
+    .click({ force: true });
+
   await page.locator('[data-testid="add-pipeline"]').dispatchEvent('click');
 
   const waitForSearchResponse = page.waitForResponse(

openmetadata-ui/src/main/resources/ui/src/components/Entity/EntityLineage/LineageNodeLabelV1.tsx

@@ -85,12 +85,8 @@ const LineageNodeLabelV1 = ({ node }: Pick<LineageNodeLabelProps, 'node'>) => {
   return (
     <div className="w-76">
       <div className="m-0 p-x-md p-y-xs">
-        <div className="d-flex gap-2 items-center m-b-xs">
-          <Space
-            wrap
-            align="start"
-            className="lineage-breadcrumb w-full"
-            size={4}>
+        {breadcrumbs.length > 0 && (
+          <div className="d-flex gap-2 items-center m-b-xs lineage-breadcrumb">
             {breadcrumbs.map((breadcrumb, index) => (
               <React.Fragment key={breadcrumb.name}>
                 <Typography.Text
@@ -105,8 +101,9 @@ const LineageNodeLabelV1 = ({ node }: Pick<LineageNodeLabelProps, 'node'>) => {
                 )}
               </React.Fragment>
             ))}
-          </Space>
-        </div>
+          </div>
+        )}
+
         <EntityLabel node={node} />
       </div>
     </div>

openmetadata-ui/src/main/resources/ui/src/components/Entity/EntityLineage/NodeChildren/NodeChildren.component.tsx

@@ -113,6 +113,10 @@ const NodeChildren = ({ node, isConnectable }: NodeChildrenProps) => {
     }
   }, [children]);
 
+  useEffect(() => {
+    setShowAllColumns(expandAllColumns);
+  }, [expandAllColumns]);
+
   const renderRecord = useCallback(
     (record: Column) => {
       const isColumnTraced = tracedColumns.includes(

openmetadata-ui/src/main/resources/ui/src/components/Entity/EntityLineage/custom-node.less

@@ -30,6 +30,7 @@
 .ant-btn.ant-btn-background-ghost.expand-btn {
   background-color: white;
   box-shadow: none;
+
   &:hover {
     background-color: white;
   }
@@ -43,6 +44,7 @@
   border: 1px solid @lineage-border;
   border-radius: 10px;
   overflow: hidden;
+
   .profiler-item {
     width: 36px;
     height: 36px;
@@ -50,16 +52,20 @@
     border-radius: 4px;
     line-height: 36px;
     font-size: 14px;
+
     &.green {
       border: 1px solid @green-5;
     }
+
     &.amber {
       border: 1px solid @yellow-4;
     }
+
     &.red {
       border: 1px solid @red-5;
     }
   }
+
   .column-container {
     min-height: 48px;
     padding: 12px;
@@ -68,19 +74,24 @@
     .lineage-collapse-column.ant-collapse {
       border: none;
       border-radius: 0;
+
       .ant-collapse-header {
         padding: 0;
         font-size: 12px;
+
         .custom-node-column-container {
           background-color: @lineage-collapse-header;
         }
+
         .lineage-column-node-handle {
           background-color: @lineage-collapse-header;
         }
       }
+
       .ant-collapse-content-box {
         padding: 4px;
       }
+
       .ant-collapse-item {
         border: none;
         border-radius: 0;
@@ -114,6 +125,7 @@
   .lineage-node-handle {
     border-color: @primary-color;
   }
+
   .lineage-node {
     border-color: @primary-color !important;
   }
@@ -137,15 +149,19 @@
   .lineage-node {
     border-color: @primary-color !important;
   }
+
   .lineage-node-handle {
     border-color: @primary-color;
+
     svg {
       color: @primary-color;
     }
   }
+
   .label-container {
     background: @primary-1;
   }
+
   .column-container {
     background: @primary-1;
     border-top: 1px solid @border-color;
@@ -171,15 +187,19 @@
     &.lineage-node {
       border-color: @red-3 !important;
     }
+
     .lineage-node-handle {
       border-color: @red-3;
+
       svg {
         color: @red-3;
       }
     }
+
     .label-container {
       background: fade(@red-3, 10%);
     }
+
     .column-container {
       background: fade(@red-3, 10%);
       border-top: 1px solid @border-color;
@@ -191,15 +211,18 @@
   .label-container {
     background: @primary-1;
   }
+
   .column-container {
     background: @primary-1;
     border-top: 1px solid @border-color;
   }
 }
+
 .data-quality-failed-custom-node-header.custom-node-header-active {
   .label-container {
     background: fade(@red-3, 10%);
   }
+
   .column-container {
     background: fade(@red-3, 10%);
     border-top: 1px solid @red-3;
@@ -214,6 +237,7 @@
   .lineage-node-handle.react-flow__handle-left {
     left: -22px;
   }
+
   .lineage-node-handle.react-flow__handle-right {
     right: -22px;
   }
@@ -227,6 +251,7 @@
   border-color: @lineage-border !important;
   background: @white !important;
   top: 43px !important; // Need to show handles on top half
+
   svg {
     color: @text-grey-muted;
   }
@@ -241,9 +266,11 @@
   height: 25px;
   transform: none;
   border: none;
+
   &.react-flow__handle-left {
     left: 0;
   }
+
   &.react-flow__handle-right {
     right: 0;
   }
@@ -266,6 +293,7 @@
 
     .custom-node-name-icon {
       width: 14px;
+      display: flex;
     }
   }
 

openmetadata-ui/src/main/resources/ui/src/components/Entity/EntityLineage/lineage-node-label.less

@@ -15,9 +15,9 @@
     height: 28px;
     width: 28px;
   }
+
   .lineage-breadcrumb {
     .lineage-breadcrumb-item {
-      max-width: 140px;
       white-space: nowrap;
       overflow: hidden;
       text-overflow: ellipsis;

openmetadata-ui/src/main/resources/ui/src/context/LineageProvider/LineageProvider.tsx

@@ -63,11 +63,7 @@ import {
   ZOOM_VALUE,
 } from '../../constants/Lineage.constants';
 import { mockDatasetData } from '../../constants/mockTourData.constants';
-import {
-  EntityLineageDirection,
-  EntityLineageNodeType,
-  EntityType,
-} from '../../enums/entity.enum';
+import { EntityLineageNodeType, EntityType } from '../../enums/entity.enum';
 import { AddLineage } from '../../generated/api/lineage/addLineage';
 import { LineageSettings } from '../../generated/configuration/lineageSettings';
 import { LineageLayer } from '../../generated/settings/settings';
@@ -97,7 +93,7 @@ import {
   getChildMap,
   getClassifiedEdge,
   getConnectedNodesEdges,
-  getLayoutedElements,
+  getELKLayoutedElements,
   getLineageEdge,
   getLineageEdgeForAPI,
   getLoadingStatusValue,
@@ -107,6 +103,7 @@ import {
   getUpdatedColumnsFromEdge,
   getUpstreamDownstreamNodesEdges,
   onLoad,
+  positionNodesUsingElk,
   removeLineageHandler,
 } from '../../utils/EntityLineageUtils';
 import { getEntityReferenceFromEntity } from '../../utils/EntityUtils';
@@ -1067,27 +1064,29 @@ const LineageProvider = ({ children }: LineageProviderProps) => {
   );
 
   const selectNode = (node: Node) => {
-    centerNodePosition(node, reactFlowInstance);
+    centerNodePosition(node, reactFlowInstance, zoomValue);
   };
 
   const repositionLayout = useCallback(
-    (activateNode = false) => {
+    async (activateNode = false) => {
+      if (nodes.length === 0 || !reactFlowInstance) {
+        return;
+      }
+
       const isColView = activeLayer.includes(LineageLayer.ColumnLevelLineage);
-      const { node, edge } = getLayoutedElements(
-        {
-          node: nodes,
-          edge: edges,
-        },
-        EntityLineageDirection.LEFT_RIGHT,
-        isColView,
-        isEditMode || expandAllColumns,
-        columnsHavingLineage
-      );
+      const { nodes: layoutedNodes, edges: layoutedEdges } =
+        await getELKLayoutedElements(
+          nodes,
+          edges,
+          isColView,
+          isEditMode || expandAllColumns,
+          columnsHavingLineage
+        );
 
-      setNodes(node);
-      setEdges(edge);
+      setNodes(layoutedNodes);
+      setEdges(layoutedEdges);
 
-      const rootNode = node.find((n) => n.data.isRootNode);
+      const rootNode = layoutedNodes.find((n) => n.data.isRootNode);
       if (!rootNode) {
         if (activateNode && reactFlowInstance) {
           onLoad(reactFlowInstance); // Call fitview in case of pipeline
@@ -1097,12 +1096,13 @@ const LineageProvider = ({ children }: LineageProviderProps) => {
       }
 
       // Center the root node in the view
-      centerNodePosition(rootNode, reactFlowInstance);
+      centerNodePosition(rootNode, reactFlowInstance, zoomValue);
       if (activateNode) {
         onNodeClick(rootNode);
       }
     },
     [
+      zoomValue,
       reactFlowInstance,
       activeLayer,
       nodes,
@@ -1115,7 +1115,7 @@ const LineageProvider = ({ children }: LineageProviderProps) => {
   );
 
   const redrawLineage = useCallback(
-    (lineageData: EntityLineageResponse) => {
+    async (lineageData: EntityLineageResponse) => {
       const allNodes = uniqWith(
         [
           ...(lineageData.nodes ?? []),
@@ -1135,8 +1135,28 @@ const LineageProvider = ({ children }: LineageProviderProps) => {
         lineageData.edges ?? [],
         decodedFqn
       );
-      setNodes(updatedNodes);
-      setEdges(updatedEdges);
+
+      if (reactFlowInstance && reactFlowInstance.viewportInitialized) {
+        const positionedNodesEdges = await positionNodesUsingElk(
+          updatedNodes,
+          updatedEdges,
+          activeLayer.includes(LineageLayer.ColumnLevelLineage),
+          isEditMode || expandAllColumns,
+          columnsHavingLineage
+        );
+        setNodes(positionedNodesEdges.nodes);
+        setEdges(positionedNodesEdges.edges);
+        const rootNode = positionedNodesEdges.nodes.find(
+          (n) => n.data.isRootNode
+        );
+        if (rootNode) {
+          centerNodePosition(rootNode, reactFlowInstance, zoomValue);
+        }
+      } else {
+        setNodes(updatedNodes);
+        setEdges(updatedEdges);
+      }
+
       setColumnsHavingLineage(columnsHavingLineage);
 
       // Get upstream downstream nodes and edges data
@@ -1151,7 +1171,14 @@ const LineageProvider = ({ children }: LineageProviderProps) => {
         selectNode(activeNode);
       }
     },
-    [decodedFqn, activeNode, activeLayer, isEditMode]
+    [
+      decodedFqn,
+      activeNode,
+      activeLayer,
+      isEditMode,
+      reactFlowInstance,
+      zoomValue,
+    ]
   );
 
   useEffect(() => {

openmetadata-ui/src/main/resources/ui/src/utils/EntityLineageUtils.tsx

@@ -14,6 +14,7 @@
 import { CheckOutlined, SearchOutlined } from '@ant-design/icons';
 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 {
   cloneDeep,
@@ -125,14 +126,15 @@ export const onLoad = (reactFlowInstance: ReactFlowInstance) => {
 
 export const centerNodePosition = (
   node: Node,
-  reactFlowInstance?: ReactFlowInstance
+  reactFlowInstance?: ReactFlowInstance,
+  zoomValue?: number
 ) => {
   const { position, width } = node;
   reactFlowInstance?.setCenter(
     position.x + (width ?? 1 / 2),
     position.y + NODE_HEIGHT / 2,
     {
-      zoom: ZOOM_VALUE,
+      zoom: zoomValue ?? ZOOM_VALUE,
       duration: ZOOM_TRANSITION_DURATION,
     }
   );
@@ -218,6 +220,73 @@ export const getLayoutedElements = (
   return { node: uNode, edge: edgesRequired };
 };
 
+const layoutOptions = {
+  'elk.algorithm': 'layered',
+  'elk.direction': 'RIGHT',
+  'elk.layered.spacing.edgeNodeBetweenLayers': '50',
+  'elk.spacing.nodeNode': '60',
+  'elk.layered.nodePlacement.strategy': 'SIMPLE',
+};
+
+const elk = new ELK();
+
+export const getELKLayoutedElements = async (
+  nodes: Node[],
+  edges: Edge[],
+  isExpanded = true,
+  expandAllColumns = false,
+  columnsHavingLineage: string[] = []
+) => {
+  const elkNodes: ElkNode[] = nodes.map((node) => {
+    const { childrenHeight } = getEntityChildrenAndLabel(
+      node.data.node,
+      expandAllColumns,
+      columnsHavingLineage
+    );
+    const nodeHeight = isExpanded ? childrenHeight + 220 : NODE_HEIGHT;
+
+    return {
+      ...node,
+      targetPosition: 'left',
+      sourcePosition: 'right',
+      width: NODE_WIDTH,
+      height: nodeHeight,
+    };
+  });
+
+  const elkEdges: ElkExtendedEdge[] = edges.map((edge) => ({
+    id: edge.id,
+    sources: [edge.source],
+    targets: [edge.target],
+  }));
+
+  const graph = {
+    id: 'root',
+    layoutOptions: layoutOptions,
+    children: elkNodes,
+    edges: elkEdges,
+  };
+
+  try {
+    const layoutedGraph = await elk.layout(graph);
+    const updatedNodes: Node[] = nodes.map((node) => {
+      const layoutedNode = (layoutedGraph?.children ?? []).find(
+        (elkNode) => elkNode.id === node.id
+      );
+
+      return {
+        ...node,
+        position: { x: layoutedNode?.x ?? 0, y: layoutedNode?.y ?? 0 },
+        hidden: false,
+      };
+    });
+
+    return { nodes: updatedNodes, edges: edges ?? [] };
+  } catch (error) {
+    return { nodes: [], edges: [] };
+  }
+};
+
 export const getModalBodyText = (selectedEdge: Edge) => {
   const { data } = selectedEdge;
   const { fromEntity, toEntity } = data.edge as EdgeDetails;
@@ -508,7 +577,7 @@ const calculateHeightAndFlattenNode = (
       expandAllColumns ||
       columnsHavingLineage.indexOf(child.fullyQualifiedName ?? '') !== -1
     ) {
-      totalHeight += 27; // Add height for the current child
+      totalHeight += 31; // Add height for the current child
     }
     flattened.push(child);
 
@@ -682,6 +751,24 @@ const getNodeType = (
   return EntityLineageNodeType.DEFAULT;
 };
 
+export const positionNodesUsingElk = async (
+  nodes: Node[],
+  edges: Edge[],
+  isColView: boolean,
+  expandAllColumns = false,
+  columnsHavingLineage: string[] = []
+) => {
+  const obj = await getELKLayoutedElements(
+    nodes,
+    edges,
+    isColView,
+    expandAllColumns,
+    columnsHavingLineage
+  );
+
+  return obj;
+};
+
 export const createNodes = (
   nodesData: EntityReference[],
   edgesData: EdgeDetails[],
@@ -692,37 +779,8 @@ export const createNodes = (
     getEntityName(a).localeCompare(getEntityName(b))
   );
 
-  const GraphInstance = graphlib.Graph;
-  const graph = new GraphInstance();
-
-  // Set an object for the graph label
-  graph.setGraph({
-    rankdir: EntityLineageDirection.LEFT_RIGHT,
-  });
-
-  // Default to assigning a new object as a label for each new edge.
-  graph.setDefaultEdgeLabel(() => ({}));
-
-  // Add nodes to the graph
-  uniqueNodesData.forEach((node) => {
+  return uniqueNodesData.map((node) => {
     const { childrenHeight } = getEntityChildrenAndLabel(node as SourceType);
-    const nodeHeight = isExpanded ? childrenHeight + 220 : NODE_HEIGHT;
-    graph.setNode(node.id, { width: NODE_WIDTH, height: nodeHeight });
-  });
-
-  // Add edges to the graph (if you have edge information)
-  edgesData.forEach((edge) => {
-    graph.setEdge(edge.fromEntity.id, edge.toEntity.id);
-  });
-
-  // Perform the layout
-  layout(graph);
-
-  // Get the layout positions
-  const layoutPositions = graph.nodes().map((nodeId) => graph.node(nodeId));
-
-  return uniqueNodesData.map((node, index) => {
-    const position = layoutPositions[index];
     const type =
       node.type === EntityLineageNodeType.LOAD_MORE
         ? node.type
@@ -741,9 +799,11 @@ export const createNodes = (
         node,
         isRootNode: entityFqn === node.fullyQualifiedName,
       },
+      width: NODE_WIDTH,
+      height: isExpanded ? childrenHeight + 220 : NODE_HEIGHT,
       position: {
-        x: position.x - NODE_WIDTH / 2,
-        y: position.y - position.height / 2,
+        x: 0,
+        y: 0,
       },
     };
   });

openmetadata-ui/src/main/resources/ui/yarn.lock

@@ -7448,6 +7448,11 @@ electron-to-chromium@^1.4.535:
   resolved "https://registry.yarnpkg.com/electron-to-chromium/-/electron-to-chromium-1.4.537.tgz#aac4101db53066be1e49baedd000a26bc754adc9"
   integrity sha512-W1+g9qs9hviII0HAwOdehGYkr+zt7KKdmCcJcjH0mYg6oL8+ioT3Skjmt7BLoAQqXhjf40AXd+HlR4oAWMlXjA==
 
+elkjs@^0.9.3:
+  version "0.9.3"
+  resolved "https://registry.yarnpkg.com/elkjs/-/elkjs-0.9.3.tgz#16711f8ceb09f1b12b99e971b138a8384a529161"
+  integrity sha512-f/ZeWvW/BCXbhGEf1Ujp29EASo/lk1FDnETgNKwJrsVvGZhUWCZyg3xLJjAsxfOmt8KjswHmI5EwCQcPMpOYhQ==
+
 emittery@^0.7.1:
   version "0.7.2"
   resolved "https://registry.yarnpkg.com/emittery/-/emittery-0.7.2.tgz#25595908e13af0f5674ab419396e2fb394cdfa82"