Fix : UI manual Lineage Editor Issues (#4532)

* Fix : UI manual Lineage Editor Issues

* Keep node if only edge is deleted.

* Adding hidden handler

* Add invisible handle on custom node

* Fix funtion name typo

* Fix node overlapping issue

* Fix #3508 Manual Lineage Editor: Do not reorganize the graph as the user is connecting the nodes

* Fix code smell

* Minor Fix

* Styling fix

* Fix Flaky state issue

* Refactor onConnect Method

* Fix duplicate edge and node issue

* Fix Failing Unit test

* Fix confirmation modal source and target node name issue

* Add check for isNode in Element Click Handler

* Add makeEdge Helper

* Add JSDoc for helper methods

* Remove onElementsRemove prop

* Refactor node remove button

* Move util method to util file

* Allow users to delete edge and node separately

* Add unit test

* Fix Node Styling

* Minor Fix

* Add invisble edges

openmetadata-ui/src/main/resources/ui/src/components/EntityLineage/CustomEdge.component.test.tsx

@@ -11,7 +11,12 @@
  *  limitations under the License.
  */
 
-import { findByTestId, queryByTestId, render } from '@testing-library/react';
+import {
+  findAllByTestId,
+  findByTestId,
+  queryByTestId,
+  render,
+} from '@testing-library/react';
 import React from 'react';
 import { ArrowHeadType, EdgeProps, Position } from 'react-flow-renderer';
 import { MemoryRouter } from 'react-router-dom';
@@ -40,6 +45,7 @@ const mockCustomEdgeProp = {
       id: 'node1',
     },
   },
+  selected: true,
 } as EdgeProps;
 
 describe('Test CustomEdge Component', () => {
@@ -49,27 +55,24 @@ describe('Test CustomEdge Component', () => {
     });
 
     const deleteButton = await findByTestId(container, 'delete-button');
-    const edgePathElement = await findByTestId(
+    const edgePathElement = await findAllByTestId(
       container,
       'react-flow-edge-path'
     );
 
     expect(deleteButton).toBeInTheDocument();
-    expect(edgePathElement).toBeInTheDocument();
+    expect(edgePathElement).toHaveLength(edgePathElement.length);
   });
 
-  it('Check if CustomEdge has selectedNode as empty object', async () => {
+  it('Check if CustomEdge has selected as false', async () => {
     const { container } = render(
-      <CustomEdge
+      <CustomEdge {...mockCustomEdgeProp} selected={false} />,
-        {...mockCustomEdgeProp}
-        data={{ ...mockCustomEdgeProp.data, selectedNode: {} }}
-      />,
       {
         wrapper: MemoryRouter,
       }
     );
 
-    const edgePathElement = await findByTestId(
+    const edgePathElement = await findAllByTestId(
       container,
       'react-flow-edge-path'
     );
@@ -77,6 +80,6 @@ describe('Test CustomEdge Component', () => {
     const deleteButton = queryByTestId(container, 'delete-button');
 
     expect(deleteButton).not.toBeInTheDocument();
-    expect(edgePathElement).toBeInTheDocument();
+    expect(edgePathElement).toHaveLength(edgePathElement.length);
   });
 });

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

@@ -34,8 +34,10 @@ export const CustomEdge = ({
   arrowHeadType,
   markerEndId,
   data,
+  selected,
 }: EdgeProps) => {
-  const { onEdgeClick, selectedNode, ...rest } = data;
+  const { onEdgeClick, ...rest } = data;
+  const offset = 4;
 
   const edgePath = getBezierPath({
     sourceX,
@@ -45,6 +47,22 @@ export const CustomEdge = ({
     targetY,
     targetPosition,
   });
+  const invisibleEdgePath = getBezierPath({
+    sourceX: sourceX + offset,
+    sourceY: sourceY + offset,
+    sourcePosition,
+    targetX: targetX + offset,
+    targetY: targetY + offset,
+    targetPosition,
+  });
+  const invisibleEdgePath1 = getBezierPath({
+    sourceX: sourceX - offset,
+    sourceY: sourceY - offset,
+    sourcePosition,
+    targetX: targetX - offset,
+    targetY: targetY - offset,
+    targetPosition,
+  });
   const markerEnd = getMarkerEnd(arrowHeadType, markerEndId);
   const [edgeCenterX, edgeCenterY] = getEdgeCenter({
     sourceX,
@@ -53,6 +71,19 @@ export const CustomEdge = ({
     targetY,
   });
 
+  const getInvisiblePath = (path: string) => {
+    return (
+      <path
+        className="react-flow__edge-path"
+        d={path}
+        data-testid="react-flow-edge-path"
+        id={id}
+        markerEnd={markerEnd}
+        style={{ ...style, strokeWidth: '6px', opacity: 0 }}
+      />
+    );
+  };
+
   return (
     <Fragment>
       <path
@@ -63,15 +94,17 @@ export const CustomEdge = ({
         markerEnd={markerEnd}
         style={style}
       />
-      {(rest as CustomEdgeData)?.source?.includes(selectedNode?.id) ||
+      {getInvisiblePath(invisibleEdgePath)}
-      (rest as CustomEdgeData)?.target?.includes(selectedNode?.id) ? (
+      {getInvisiblePath(invisibleEdgePath1)}
+
+      {selected ? (
         <foreignObject
           data-testid="delete-button"
           height={foreignObjectSize}
           requiredExtensions="http://www.w3.org/1999/xhtml"
           width={foreignObjectSize}
-          x={edgeCenterX - foreignObjectSize / 4}
+          x={edgeCenterX - foreignObjectSize / offset}
-          y={edgeCenterY - foreignObjectSize / 4}>
+          y={edgeCenterY - foreignObjectSize / offset}>
           <button
             className="tw-cursor-pointer tw-flex tw-z-9999"
             onClick={(event) => onEdgeClick?.(event, rest as CustomEdgeData)}>

openmetadata-ui/src/main/resources/ui/src/components/EntityLineage/CustomNode.component.test.tsx

@@ -88,7 +88,9 @@ const mockTableColumns = [
 ];
 
 const mockCustomNodeProp = {
+  id: 'node-1',
   type: 'default',
+  selected: false,
   isConnectable: false,
   data: {
     label: <p>label</p>,

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

@@ -13,7 +13,8 @@
 
 import classNames from 'classnames';
 import React, { CSSProperties, Fragment } from 'react';
-import { Handle, HandleProps, Position } from 'react-flow-renderer';
+import { Handle, HandleProps, NodeProps, Position } from 'react-flow-renderer';
+import { getNodeRemoveButton } from '../../utils/EntityLineageUtils';
 import { getConstraintIcon } from '../../utils/TableUtils';
 
 const handleStyles = {
@@ -29,23 +30,97 @@ const getHandle = (
   isConnectable: HandleProps['isConnectable'],
   isNewNode = false
 ) => {
+  const getLeftRightHandleStyles = () => {
+    return {
+      opacity: 0,
+      borderRadius: '0px',
+      height: '162%',
+    };
+  };
+
+  const getTopBottomHandleStyles = () => {
+    return {
+      opacity: 0,
+      borderRadius: '0px',
+      width: '110%',
+    };
+  };
+
   if (nodeType === 'output') {
     return (
-      <Handle
+      <Fragment>
-        isConnectable={isConnectable}
+        <Handle
-        position={Position.Left}
+          isConnectable={isConnectable}
-        style={{ ...handleStyles, left: '-14px' } as CSSProperties}
+          position={Position.Left}
-        type="target"
+          style={{ ...handleStyles, left: '-14px' } as CSSProperties}
-      />
+          type="target"
+        />
+        <Handle
+          isConnectable={isConnectable}
+          position={Position.Left}
+          style={{
+            ...getLeftRightHandleStyles(),
+            marginLeft: '-10px',
+          }}
+          type="target"
+        />
+        <Handle
+          isConnectable={isConnectable}
+          position={Position.Bottom}
+          style={{
+            ...getTopBottomHandleStyles(),
+            marginBottom: '-6px',
+          }}
+          type="target"
+        />
+        <Handle
+          isConnectable={isConnectable}
+          position={Position.Top}
+          style={{
+            ...getTopBottomHandleStyles(),
+            marginTop: '-6px',
+          }}
+          type="target"
+        />
+      </Fragment>
     );
   } else if (nodeType === 'input') {
     return (
-      <Handle
+      <Fragment>
-        isConnectable={isConnectable}
+        <Handle
-        position={Position.Right}
+          isConnectable={isConnectable}
-        style={{ ...handleStyles, right: '-14px' } as CSSProperties}
+          position={Position.Right}
-        type="source"
+          style={{ ...handleStyles, right: '-14px' } as CSSProperties}
-      />
+          type="source"
+        />
+        <Handle
+          isConnectable={isConnectable}
+          position={Position.Right}
+          style={{
+            ...getLeftRightHandleStyles(),
+            marginRight: '-10px',
+          }}
+          type="source"
+        />
+        <Handle
+          isConnectable={isConnectable}
+          position={Position.Bottom}
+          style={{
+            ...getTopBottomHandleStyles(),
+            marginBottom: '-6px',
+          }}
+          type="target"
+        />
+        <Handle
+          isConnectable={isConnectable}
+          position={Position.Top}
+          style={{
+            ...getTopBottomHandleStyles(),
+            marginTop: '-6px',
+          }}
+          type="target"
+        />
+      </Fragment>
     );
   } else {
     return (
@@ -74,24 +149,63 @@ const getHandle = (
           }
           type="source"
         />
+        <Handle
+          isConnectable={isConnectable}
+          position={Position.Left}
+          style={{
+            ...getLeftRightHandleStyles(),
+            marginLeft: '-10px',
+          }}
+          type="target"
+        />
+        <Handle
+          isConnectable={isConnectable}
+          position={Position.Right}
+          style={{
+            ...getLeftRightHandleStyles(),
+            marginRight: '-10px',
+          }}
+          type="source"
+        />
+        <Handle
+          isConnectable={isConnectable}
+          position={Position.Bottom}
+          style={{
+            ...getTopBottomHandleStyles(),
+            marginBottom: '-6px',
+          }}
+          type="target"
+        />
+        <Handle
+          isConnectable={isConnectable}
+          position={Position.Top}
+          style={{
+            ...getTopBottomHandleStyles(),
+            marginTop: '-6px',
+          }}
+          type="target"
+        />
       </Fragment>
     );
   }
 };
 
-/* eslint-disable-next-line */
+const CustomNode = (props: NodeProps) => {
-const CustomNode = (props: any) => {
+  const { data, type, isConnectable, selected } = props;
   /* eslint-disable-next-line */
-  const { data, type, isConnectable } = props;
+  const { label, columns, isNewNode, removeNodeHandler, isEditMode } = data;
-  /* eslint-disable-next-line */
-  const { label, columns, isNewNode } = data;
 
   return (
     <div className="tw-relative nowheel ">
       {getHandle(type, isConnectable, isNewNode)}
       {/* Node label could be simple text or reactNode */}
       <div className={classNames('tw-px-2')} data-testid="node-label">
-        {label}
+        {label}{' '}
+        {selected && isEditMode
+          ? getNodeRemoveButton(() => {
+              removeNodeHandler?.(props);
+            })
+          : null}
       </div>
 
       {columns?.length ? (

openmetadata-ui/src/main/resources/ui/src/components/EntityLineage/EntityLineage.interface.ts

@@ -11,7 +11,7 @@
  *  limitations under the License.
  */
 
-import { LeafNodes, LineagePos, LoadingNodeState } from 'Models';
+import { LeafNodes, LineagePos, LoadingNodeState, LoadingState } from 'Models';
 import {
   EntityLineage,
   EntityReference,
@@ -71,3 +71,5 @@ export interface SelectedEdge {
   source: EntityReference;
   target: EntityReference;
 }
+
+export type ElementLoadingState = Exclude<LoadingState, 'waiting'>;

openmetadata-ui/src/main/resources/ui/src/components/EntityLineage/Entitylineage.component.test.tsx

@@ -146,6 +146,7 @@ jest.mock('../../utils/EntityLineageUtils', () => ({
   onNodeMouseEnter: jest.fn(),
   onNodeMouseLeave: jest.fn(),
   onNodeMouseMove: jest.fn(),
+  getUniqueFlowElements: jest.fn().mockReturnValue([]),
 }));
 
 jest.mock('../../utils/TableUtils', () => ({

openmetadata-ui/src/main/resources/ui/src/constants/Lineage.constants.ts

@@ -1,4 +1,5 @@
 import { capitalize } from 'lodash';
+import { ElementLoadingState } from '../components/EntityLineage/EntityLineage.interface';
 import { EntityType } from '../enums/entity.enum';
 
 export const foreignObjectSize = 40;
@@ -17,3 +18,8 @@ export const positionY = 60;
 
 export const nodeWidth = 300;
 export const nodeHeight = 40;
+
+export const ELEMENT_DELETE_STATE = {
+  loading: false,
+  status: 'initial' as ElementLoadingState,
+};

openmetadata-ui/src/main/resources/ui/src/styles/tailwind.css

@@ -330,8 +330,8 @@
     @apply tw-border-main;
     box-shadow: 0 0 0 0.5px #e2dce4;
   }
-  .leaf-node.selected-node,
+  .leaf-node.selected,
-  .leaf-node.selected-node:hover {
+  .leaf-node.selected:hover {
     @apply tw-border-primary-active;
     box-shadow: 0 0 0 0.5px #7147e8;
   }

openmetadata-ui/src/main/resources/ui/src/styles/x-master.css

@@ -741,6 +741,15 @@ body .profiler-graph .recharts-active-dot circle {
 .leaf-node.core .react-flow__handle {
   background-color: #7147e8;
 }
+.react-flow__edge {
+  pointer-events: all;
+  cursor: pointer;
+}
+
+.react-flow__edge .react-flow__edge-path {
+  stroke-width: 2px;
+}
+
 .react-flow__edge.selected .react-flow__edge-path {
   stroke: #7147e8;
 }

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

@@ -59,6 +59,7 @@ import {
   prepareLabel,
 } from './CommonUtils';
 import { isLeafNode } from './EntityUtils';
+import SVGIcons from './SvgUtils';
 import { getEntityLink } from './TableUtils';
 
 export const getHeaderLabel = (
@@ -119,13 +120,14 @@ export const getLineageData = (
   loadNodeHandler: (node: EntityReference, pos: LineagePos) => void,
   lineageLeafNodes: LeafNodes,
   isNodeLoading: LoadingNodeState,
-  getNodeLable: (node: EntityReference) => React.ReactNode,
+  getNodeLabel: (node: EntityReference) => React.ReactNode,
   isEditMode: boolean,
   edgeType: string,
   onEdgeClick: (
     evt: React.MouseEvent<HTMLButtonElement>,
     data: CustomEdgeData
-  ) => void
+  ) => void,
+  removeNodeHandler: (node: Node) => void
 ) => {
   const [x, y] = [0, 0];
   const nodes = [
@@ -140,6 +142,7 @@ export const getLineageData = (
       isMapped: false,
       ...down,
     })) || [];
+
   const mainNode = entityLineage['entity'];
 
   const UPStreamNodes: Elements = [];
@@ -161,8 +164,10 @@ export const getLineageData = (
       type: 'default',
       className: 'leaf-node',
       data: {
-        label: getNodeLable(node),
+        label: getNodeLabel(node),
         entityType: node.type,
+        removeNodeHandler,
+        isEditMode,
       },
       position: {
         x: pos === 'from' ? -xVal : xVal,
@@ -171,6 +176,10 @@ export const getLineageData = (
     };
   };
 
+  const makeEdge = (edge: FlowElement) => {
+    lineageEdges.push(edge);
+  };
+
   const getNodes = (
     id: string,
     pos: LineagePos,
@@ -187,7 +196,7 @@ export const getLineageData = (
           if (node) {
             UPNodes.push(node);
             UPStreamNodes.push(makeNode(node, 'from', depth, upDepth));
-            lineageEdges.push({
+            makeEdge({
               id: `edge-${up.fromEntity}-${id}-${depth}`,
               source: `${node.id}`,
               target: edg ? edg.id : `${id}`,
@@ -229,7 +238,7 @@ export const getLineageData = (
           if (node) {
             DOWNNodes.push(node);
             DOWNStreamNodes.push(makeNode(node, 'to', depth, downDepth));
-            lineageEdges.push({
+            makeEdge({
               id: `edge-${id}-${down.toEntity}`,
               source: edg ? edg.id : `${id}`,
               target: `${node.id}`,
@@ -328,7 +337,9 @@ export const getLineageData = (
           : 'input',
       className: `leaf-node ${!isEditMode ? 'core' : ''}`,
       data: {
-        label: getNodeLable(mainNode),
+        label: getNodeLabel(mainNode),
+        isEditMode,
+        removeNodeHandler,
       },
       position: { x: x, y: y },
     },
@@ -343,6 +354,7 @@ export const getLineageData = (
             ...up,
             type: isEditMode ? 'default' : 'input',
             data: {
+              ...up.data,
               label: (
                 <div className="tw-flex">
                   <div
@@ -387,6 +399,7 @@ export const getLineageData = (
             ...down,
             type: isEditMode ? 'default' : 'output',
             data: {
+              ...down.data,
               label: (
                 <div className="tw-flex tw-justify-between">
                   <div>{down?.data?.label}</div>
@@ -444,7 +457,7 @@ export const getDataLabel = (
   } else {
     return (
       <span
-        className="tw-break-words description-text tw-self-center"
+        className="tw-break-words tw-self-center tw-w-60"
         data-testid="lineage-entity">
         {type === 'table'
           ? databaseName && schemaName
@@ -496,8 +509,8 @@ export const getLayoutedElements = (
   elements.forEach((el) => {
     if (isNode(el)) {
       dagreGraph.setNode(el.id, {
-        width: el?.__rf?.width ?? nodeWidth,
+        width: nodeWidth,
-        height: el?.__rf?.height ?? nodeHeight,
+        height: nodeHeight,
       });
     } else {
       dagreGraph.setEdge(el.source, el.target);
@@ -512,11 +525,8 @@ export const getLayoutedElements = (
       el.targetPosition = isHorizontal ? Position.Left : Position.Top;
       el.sourcePosition = isHorizontal ? Position.Right : Position.Bottom;
       el.position = {
-        x:
+        x: nodeWithPosition.x - nodeWidth / 2,
-          nodeWithPosition.x -
+        y: nodeWithPosition.y - nodeHeight / 2,
-          (el?.__rf?.width ?? nodeWidth) / 2 +
-          Math.random() / 1000,
-        y: nodeWithPosition.y - (el?.__rf?.height ?? nodeHeight) / 2,
       };
     }
 
@@ -527,27 +537,19 @@ export const getLayoutedElements = (
 export const getModalBodyText = (selectedEdge: SelectedEdge) => {
   let sourceEntity = '';
   let targetEntity = '';
+  const sourceFQN = selectedEdge.source.fullyQualifiedName || '';
+  const targetFQN = selectedEdge.target.fullyQualifiedName || '';
 
   if (selectedEdge.source.type === EntityType.TABLE) {
-    sourceEntity = getPartialNameFromTableFQN(
+    sourceEntity = getPartialNameFromTableFQN(sourceFQN, [FqnPart.Table]);
-      selectedEdge.source.name as string,
-      [FqnPart.Table]
-    );
   } else {
-    sourceEntity = getPartialNameFromFQN(selectedEdge.source.name as string, [
+    sourceEntity = getPartialNameFromFQN(sourceFQN, ['database']);
-      'database',
-    ]);
   }
 
   if (selectedEdge.target.type === EntityType.TABLE) {
-    targetEntity = getPartialNameFromTableFQN(
+    targetEntity = getPartialNameFromTableFQN(targetFQN, [FqnPart.Table]);
-      selectedEdge.target.name as string,
-      [FqnPart.Table]
-    );
   } else {
-    targetEntity = getPartialNameFromFQN(selectedEdge.target.name as string, [
+    targetEntity = getPartialNameFromFQN(targetFQN, ['database']);
-      'database',
-    ]);
   }
 
   return `Are you sure you want to remove the edge between "${
@@ -560,3 +562,32 @@ export const getModalBodyText = (selectedEdge: SelectedEdge) => {
       : targetEntity
   }"?`;
 };
+
+export const getUniqueFlowElements = (elements: FlowElement[]) => {
+  const flag: { [x: string]: boolean } = {};
+  const uniqueElements: Elements = [];
+
+  elements.forEach((elem) => {
+    if (!flag[elem.id]) {
+      flag[elem.id] = true;
+      uniqueElements.push(elem);
+    }
+  });
+
+  return uniqueElements;
+};
+
+/**
+ *
+ * @param onClick - callback
+ * @returns - Button element with attach callback
+ */
+export const getNodeRemoveButton = (onClick: () => void) => {
+  return (
+    <button
+      className="tw-absolute tw--top-4 tw--right-6 tw-cursor-pointer tw-z-9999 tw-bg-body-hover tw-rounded-full"
+      onClick={() => onClick()}>
+      <SVGIcons alt="times-circle" icon="icon-times-circle" width="16px" />
+    </button>
+  );
+};