From a76b9bdcec25ad616b8f1985f09c0327cf30feac Mon Sep 17 00:00:00 2001
From: Ashish Gupta <ashish@getcollate.io>
Date: Mon, 10 Feb 2025 13:41:41 +0530
Subject: [PATCH] modify the appeariance of self connecting edge lineage
 (#19714)

Co-authored-by: Chirag Madlani <12962843+chirag-madlani@users.noreply.github.com>
---
 .../EntityLineage/CustomEdge.component.tsx    |  49 ++++----
 .../ui/src/utils/EntityLineageUtils.tsx       | 110 ++++++++++++++++++
 2 files changed, 132 insertions(+), 27 deletions(-)

diff --git a/openmetadata-ui/src/main/resources/ui/src/components/Entity/EntityLineage/CustomEdge.component.tsx b/openmetadata-ui/src/main/resources/ui/src/components/Entity/EntityLineage/CustomEdge.component.tsx
index 81e57860c2b..b7ede63f5ff 100644
--- a/openmetadata-ui/src/main/resources/ui/src/components/Entity/EntityLineage/CustomEdge.component.tsx
+++ b/openmetadata-ui/src/main/resources/ui/src/components/Entity/EntityLineage/CustomEdge.component.tsx
@@ -15,7 +15,7 @@ import Icon from '@ant-design/icons/lib/components/Icon';
 import { Button, Tag } from 'antd';
 import classNames from 'classnames';
 import React, { Fragment, useCallback, useMemo } from 'react';
-import { EdgeProps, getBezierPath } from 'reactflow';
+import { EdgeProps } from 'reactflow';
 import { ReactComponent as IconEditCircle } from '../../../assets/svg/ic-edit-circle.svg';
 import { ReactComponent as FunctionIcon } from '../../../assets/svg/ic-function.svg';
 import { ReactComponent as IconTimesCircle } from '../../../assets/svg/ic-times-circle.svg';
@@ -27,7 +27,10 @@ import { EntityType } from '../../../enums/entity.enum';
 import { StatusType } from '../../../generated/entity/data/pipeline';
 import { LineageLayer } from '../../../generated/settings/settings';
 import { useApplicationStore } from '../../../hooks/useApplicationStore';
-import { getColumnSourceTargetHandles } from '../../../utils/EntityLineageUtils';
+import {
+  getColumnSourceTargetHandles,
+  getEdgePathData,
+} from '../../../utils/EntityLineageUtils';
 import { getEntityName } from '../../../utils/EntityUtils';
 import EntityPopOverCard from '../../common/PopOverCard/EntityPopOverCard';
 import { CustomEdgeData } from './EntityLineage.interface';
@@ -69,6 +72,8 @@ export const CustomEdge = ({
   markerEnd,
   data,
   selected,
+  source,
+  target,
 }: EdgeProps) => {
   const {
     edge,
@@ -96,6 +101,21 @@ export const CustomEdge = ({
 
   const { theme } = useApplicationStore();
 
+  const {
+    edgePath,
+    edgeCenterX,
+    edgeCenterY,
+    invisibleEdgePath,
+    invisibleEdgePath1,
+  } = getEdgePathData(source, target, offset, {
+    sourceX,
+    sourceY,
+    targetX,
+    targetY,
+    sourcePosition,
+    targetPosition,
+  });
+
   const showDqTracing = useMemo(() => {
     return (
       (activeLayer.includes(LineageLayer.DataObservability) &&
@@ -122,31 +142,6 @@ export const CustomEdge = ({
     );
   }, [isColumnLineage, tracedColumns, sourceHandle, targetHandle]);
 
-  const [edgePath, edgeCenterX, edgeCenterY] = getBezierPath({
-    sourceX,
-    sourceY,
-    sourcePosition,
-    targetX,
-    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 updatedStyle = useMemo(() => {
     const isNodeTraced =
       tracedNodes.includes(edge.fromEntity.id) &&
diff --git a/openmetadata-ui/src/main/resources/ui/src/utils/EntityLineageUtils.tsx b/openmetadata-ui/src/main/resources/ui/src/utils/EntityLineageUtils.tsx
index 3256397d340..f7073b87ba0 100644
--- a/openmetadata-ui/src/main/resources/ui/src/utils/EntityLineageUtils.tsx
+++ b/openmetadata-ui/src/main/resources/ui/src/utils/EntityLineageUtils.tsx
@@ -30,6 +30,7 @@ import React, { MouseEvent as ReactMouseEvent } from 'react';
 import {
   Connection,
   Edge,
+  getBezierPath,
   getConnectedEdges,
   getIncomers,
   getOutgoers,
@@ -1451,3 +1452,112 @@ export const getColumnFunctionValue = (
 
   return column?.function;
 };
+
+interface EdgeAlignmentPathDataProps {
+  sourceX: number;
+  sourceY: number;
+  targetX: number;
+  targetY: number;
+  sourcePosition: Position;
+  targetPosition: Position;
+}
+
+export const isSelfConnectingEdge = (source: string, target: string) => {
+  return source === target;
+};
+
+const getSelfConnectingEdgePath = ({
+  sourceX,
+  sourceY,
+  targetX,
+  targetY,
+}: EdgeAlignmentPathDataProps) => {
+  const radiusX = (sourceX - targetX) * 0.6;
+  const radiusY = 50;
+
+  return `M ${sourceX - 5} ${sourceY} A ${radiusX} ${radiusY} 0 1 0 ${
+    targetX + 2
+  } ${targetY}`;
+};
+
+export const getEdgePathAlignmentData = (
+  source: string,
+  target: string,
+  edgePathData: {
+    sourceX: number;
+    sourceY: number;
+    targetX: number;
+    targetY: number;
+  }
+) => {
+  if (isSelfConnectingEdge(source, target)) {
+    // modify the edge path data as per the self connecting edges behavior
+    return {
+      sourceX: edgePathData.sourceX - 5,
+      sourceY: edgePathData.sourceY - 80,
+      targetX: edgePathData.targetX + 2,
+      targetY: edgePathData.targetY - 80,
+    };
+  }
+
+  return edgePathData;
+};
+
+const getEdgePath = (
+  edgePath: string,
+  source: string,
+  target: string,
+  alignmentPathData: EdgeAlignmentPathDataProps
+) => {
+  return isSelfConnectingEdge(source, target)
+    ? getSelfConnectingEdgePath(alignmentPathData)
+    : edgePath;
+};
+
+export const getEdgePathData = (
+  source: string,
+  target: string,
+  offset: number,
+  edgePathData: EdgeAlignmentPathDataProps
+) => {
+  const { sourceX, sourceY, targetX, targetY } = getEdgePathAlignmentData(
+    source,
+    target,
+    edgePathData
+  );
+  const { sourcePosition, targetPosition } = edgePathData;
+
+  const [edgePath, edgeCenterX, edgeCenterY] = getBezierPath({
+    sourceX,
+    sourceY,
+    sourcePosition,
+    targetX,
+    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,
+  });
+
+  return {
+    edgePath: getEdgePath(edgePath, source, target, edgePathData), // pass the initial data edgePathData, as edge modification will be done based on the initial data
+    edgeCenterX,
+    edgeCenterY,
+    invisibleEdgePath,
+    invisibleEdgePath1,
+  };
+};