Fix #2797 Lineage: Issue while adding downstream Table/Dashboard/Pipeline (#3506)

openmetadata-ui/src/main/resources/ui/src/components/EntityInfoDrawer/EntityInfoDrawer.component.tsx

@@ -153,7 +153,11 @@ const EntityInfoDrawer = ({
       <header className="tw-flex tw-justify-between">
         <p className="tw-flex">
           <span className="tw-mr-2">{getEntityIcon(selectedNode.type)}</span>
-          {getHeaderLabel(selectedNode.name, selectedNode.type, isMainNode)}
+          {getHeaderLabel(
+            selectedNode.displayName ?? selectedNode.name,
+            selectedNode.type,
+            isMainNode
+          )}
         </p>
         <div className="tw-flex">
           <svg

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

@@ -133,6 +133,11 @@ const Entitylineage: FunctionComponent<EntityLineageProp> = ({
     status: 'initial',
   });
 
+  /**
+   * take node as input and check if node is main entity or not
+   * @param node
+   * @returns class `leaf-node core` for main node and `leaf-node` for leaf node
+   */
   const getNodeClass = (node: FlowElement) => {
     return `${
       node.id.includes(lineageData.entity?.id) && !isEditMode
@@ -141,14 +146,29 @@ const Entitylineage: FunctionComponent<EntityLineageProp> = ({
     }`;
   };
 
+  /**
+   * take entity as input and set it as selected entity
+   * @param entity
+   */
   const selectedEntityHandler = (entity: EntityReference) => {
     setSelectedEntity(entity);
   };
+
+  /**
+   * take state and value to set selected node
+   * @param state
+   * @param value
+   */
   const selectNodeHandler = (state: boolean, value: SelectedNode) => {
     setIsDrawerOpen(state);
     setSelectedNode(value);
   };
 
+  /**
+   *
+   * @param node
+   * @returns lable for given node
+   */
   const getNodeLable = (node: EntityReference) => {
     return (
       <>
@@ -177,6 +197,11 @@ const Entitylineage: FunctionComponent<EntityLineageProp> = ({
     );
   };
 
+  /**
+   *
+   * @param data selected edge
+   * @param confirmDelete confirmation state for deleting seslected edge
+   */
   const removeEdgeHandler = (data: SelectedEdge, confirmDelete: boolean) => {
     if (confirmDelete) {
       const edgeData: EdgeData = {
@@ -215,6 +240,11 @@ const Entitylineage: FunctionComponent<EntityLineageProp> = ({
     }
   };
 
+  /**
+   * take edge data and set it as selected edge
+   * @param evt
+   * @param data
+   */
   const onEdgeClick = (
     evt: React.MouseEvent<HTMLButtonElement>,
     data: CustomEdgeData
@@ -245,6 +275,10 @@ const Entitylineage: FunctionComponent<EntityLineageProp> = ({
     });
   };
 
+  /**
+   *
+   * @returns unique flow elements
+   */
   const setElementsHandle = () => {
     const flag: { [x: string]: boolean } = {};
     const uniqueElements: Elements = [];
@@ -276,6 +310,10 @@ const Entitylineage: FunctionComponent<EntityLineageProp> = ({
     getLayoutedElements(setElementsHandle())
   );
 
+  /**
+   * take boolean value as input and reset selected node
+   * @param value
+   */
   const closeDrawer = (value: boolean) => {
     setIsDrawerOpen(value);
     setElements((prevElements) => {
@@ -293,21 +331,47 @@ const Entitylineage: FunctionComponent<EntityLineageProp> = ({
     setSelectedNode({} as SelectedNode);
   };
 
+  /**
+   * take list of elements to remove it from the graph
+   * @param elementsToRemove
+   * @returns updated elements list
+   */
   const onElementsRemove = (elementsToRemove: Elements) =>
     setElements((els) => removeElements(elementsToRemove, els));
 
+  /**
+   * take edge or connection to add new element in the graph
+   * @param params
+   */
   const onConnect = (params: Edge | Connection) => {
     setStatus('waiting');
     setLoading(true);
     const { target, source } = params;
 
-    const downstreamNode = lineageData.downstreamEdges?.find((d) =>
+    const nodes = [
+      ...(lineageData.nodes as EntityReference[]),
+      lineageData.entity,
+    ];
+
+    const sourceDownstreamNode = lineageData.downstreamEdges?.find((d) =>
       source?.includes(d.toEntity as string)
     );
 
-    let targetNode = lineageData.nodes?.find((n) => target?.includes(n.id));
+    const sourceUpStreamNode = lineageData.upstreamEdges?.find((u) =>
+      source?.includes(u.fromEntity as string)
+    );
 
-    let sourceNode = lineageData.nodes?.find((n) => source?.includes(n.id));
+    const targetDownStreamNode = lineageData.downstreamEdges?.find((d) =>
+      target?.includes(d.toEntity as string)
+    );
+
+    const targetUpStreamNode = lineageData.upstreamEdges?.find((u) =>
+      target?.includes(u.fromEntity as string)
+    );
+
+    let targetNode = nodes?.find((n) => target?.includes(n.id));
+
+    let sourceNode = nodes?.find((n) => source?.includes(n.id));
 
     if (isUndefined(targetNode)) {
       targetNode = target?.includes(lineageData.entity?.id)
@@ -332,9 +396,11 @@ const Entitylineage: FunctionComponent<EntityLineageProp> = ({
         },
       },
     };
+
     setElements((els) =>
       addEdge({ ...params, arrowHeadType: ArrowHeadType.ArrowClosed }, els)
     );
+
     setTimeout(() => {
       addLineageHandler(newEdge)
         .then(() => {
@@ -350,8 +416,9 @@ const Entitylineage: FunctionComponent<EntityLineageProp> = ({
                   ]
                 : lineageData.nodes,
               downstreamEdges:
-                !isUndefined(downstreamNode) ||
-                sourceNode?.id === lineageData.entity?.id
+                !isUndefined(sourceUpStreamNode) ||
+                !isUndefined(targetUpStreamNode) ||
+                targetNode?.id === selectedEntity.id
                   ? [
                       ...(lineageData.downstreamEdges as EntityEdge[]),
                       {
@@ -361,8 +428,9 @@ const Entitylineage: FunctionComponent<EntityLineageProp> = ({
                     ]
                   : lineageData.downstreamEdges,
               upstreamEdges:
-                isUndefined(downstreamNode) &&
-                sourceNode?.id !== lineageData.entity?.id
+                !isUndefined(sourceDownstreamNode) ||
+                !isUndefined(targetDownStreamNode) ||
+                sourceNode?.id === selectedEntity.id
                   ? [
                       ...(lineageData.upstreamEdges as EntityEdge[]),
                       {
@@ -384,6 +452,10 @@ const Entitylineage: FunctionComponent<EntityLineageProp> = ({
     }, 500);
   };
 
+  /**
+   * take element and perform onClick logic
+   * @param el
+   */
   const onElementClick = (el: FlowElement) => {
     const node = [
       ...(lineageData.nodes as Array<EntityReference>),
@@ -393,6 +465,7 @@ const Entitylineage: FunctionComponent<EntityLineageProp> = ({
       selectNodeHandler(true, {
         name: node?.name as string,
         id: el.id,
+        displayName: node?.displayName,
         type: node?.type as string,
         entityId: node?.id as string,
       });
@@ -413,6 +486,10 @@ const Entitylineage: FunctionComponent<EntityLineageProp> = ({
     }
   };
 
+  /**
+   * this method is used for table entity to show table columns
+   * @param tableColumns
+   */
   const onNodeExpand = (tableColumns?: Column[]) => {
     const elements = getLayoutedElements(setElementsHandle());
     setElements(
@@ -433,6 +510,10 @@ const Entitylineage: FunctionComponent<EntityLineageProp> = ({
     );
   };
 
+  /**
+   * take node and get the columns for that node
+   * @param expandNode
+   */
   const getTableColumns = (expandNode?: EntityReference) => {
     if (expandNode) {
       getTableDetails(expandNode.id, ['columns'])
@@ -454,16 +535,28 @@ const Entitylineage: FunctionComponent<EntityLineageProp> = ({
     }
   };
 
+  /**
+   * take node and remove it from the graph
+   * @param node
+   */
   const removeNodeHandler = (node: FlowElement) => {
     setElements((es) => es.filter((n) => n.id !== node.id));
     setNewAddedNode({} as FlowElement);
   };
 
+  /**
+   * handle node drag event
+   * @param event
+   */
   const onDragOver = (event: DragEvent) => {
     event.preventDefault();
     event.dataTransfer.dropEffect = 'move';
   };
 
+  /**
+   * handle node drop event
+   * @param event
+   */
   const onDrop = (event: DragEvent) => {
     event.preventDefault();
 
@@ -518,6 +611,9 @@ const Entitylineage: FunctionComponent<EntityLineageProp> = ({
     }
   };
 
+  /**
+   * handle onNode select logic
+   */
   const onEntitySelect = () => {
     if (!isEmpty(selectedEntity)) {
       const isExistingNode = elements.some((n) =>
@@ -630,10 +726,12 @@ const Entitylineage: FunctionComponent<EntityLineageProp> = ({
         <div
           className={classNames(
             'tw-relative tw-h-full tw--ml-4 tw--mr-7 tw--mt-4'
-          )}>
+          )}
+          data-testid="lineage-container">
           <div className="tw-w-full tw-h-full" ref={reactFlowWrapper}>
             <ReactFlowProvider>
               <ReactFlow
+                data-testid="react-flow-component"
                 edgeTypes={{ buttonedge: CustomEdge }}
                 elements={elements as Elements}
                 elementsSelectable={!isEditMode}
@@ -729,7 +827,6 @@ const Entitylineage: FunctionComponent<EntityLineageProp> = ({
                 ) : null}
               </ReactFlow>
             </ReactFlowProvider>
-            :
           </div>
           <EntityInfoDrawer
             isMainNode={selectedNode.name === lineageData.entity?.name}

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

@@ -20,6 +20,7 @@ import {
 export interface SelectedNode {
   name: string;
   type: string;
+  displayName?: string;
   id?: string;
   entityId: string;
 }

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

@@ -0,0 +1,200 @@
+/*
+ *  Copyright 2021 Collate
+ *  Licensed under the Apache License, Version 2.0 (the "License");
+ *  you may not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *  http://www.apache.org/licenses/LICENSE-2.0
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+
+import {
+  findByTestId,
+  findByText,
+  queryByTestId,
+  render,
+} from '@testing-library/react';
+import React from 'react';
+import { MemoryRouter } from 'react-router-dom';
+import EntityLineage from './EntityLineage.component';
+
+/**
+ * mock implementation of ResizeObserver
+ */
+window.ResizeObserver = jest.fn().mockImplementation(() => ({
+  observe: jest.fn(),
+  unobserve: jest.fn(),
+  disconnect: jest.fn(),
+}));
+
+const mockLineageData = {
+  entity: {
+    id: 'efcc334a-41c8-483e-b779-464a88a7ece3',
+    type: 'table',
+    name: 'bigquery_gcp.shopify.raw_product_catalog',
+    description: '1234',
+    displayName: 'raw_product_catalog',
+    href: 'http://localhost:8585/api/v1/tables/efcc334a-41c8-483e-b779-464a88a7ece3',
+  },
+  nodes: [
+    {
+      description: 'dim_address ETL pipeline',
+      displayName: 'dim_address etl',
+      id: 'c14d78eb-dc17-4bc4-b54b-227318116da3',
+      type: 'pipeline',
+      name: 'sample_airflow.dim_address_etl',
+    },
+    {
+      description: '',
+      displayName: 'deck.gl Demo',
+      id: '7408172f-bd78-4c60-a270-f9d5ed1490ab',
+      type: 'dashboard',
+      name: 'sample_superset.10',
+    },
+    {
+      description: '',
+      displayName: 'dim_address',
+      id: 'c3cb016a-dc6e-4d22-9aa5-be8b32999a6b',
+      type: 'table',
+      name: 'bigquery_gcp.shopify.dim_address',
+    },
+    {
+      description: 'diim_location ETL pipeline',
+      displayName: 'dim_location etl',
+      id: 'bb1c2c56-9b0e-4f8e-920d-02819e5ee288',
+      type: 'pipeline',
+      name: 'sample_airflow.dim_location_etl',
+    },
+    {
+      description: '',
+      displayName: 'dim_api_client',
+      id: 'abb6567e-fbd9-47d9-95f6-29a80a5a0a52',
+      type: 'table',
+      name: 'bigquery_gcp.shopify.dim_api_client',
+    },
+  ],
+  upstreamEdges: [
+    {
+      fromEntity: 'c14d78eb-dc17-4bc4-b54b-227318116da3',
+      toEntity: 'efcc334a-41c8-483e-b779-464a88a7ece3',
+    },
+    {
+      fromEntity: 'bb1c2c56-9b0e-4f8e-920d-02819e5ee288',
+      toEntity: '7408172f-bd78-4c60-a270-f9d5ed1490ab',
+    },
+    {
+      fromEntity: '7408172f-bd78-4c60-a270-f9d5ed1490ab',
+      toEntity: 'abb6567e-fbd9-47d9-95f6-29a80a5a0a52',
+    },
+  ],
+  downstreamEdges: [
+    {
+      fromEntity: 'efcc334a-41c8-483e-b779-464a88a7ece3',
+      toEntity: '7408172f-bd78-4c60-a270-f9d5ed1490ab',
+    },
+    {
+      fromEntity: 'c14d78eb-dc17-4bc4-b54b-227318116da3',
+      toEntity: 'c3cb016a-dc6e-4d22-9aa5-be8b32999a6b',
+    },
+    {
+      fromEntity: '7408172f-bd78-4c60-a270-f9d5ed1490ab',
+      toEntity: 'abb6567e-fbd9-47d9-95f6-29a80a5a0a52',
+    },
+  ],
+};
+
+const mockEntityLineageProp = {
+  entityLineage: mockLineageData,
+  lineageLeafNodes: {
+    upStreamNode: [],
+    downStreamNode: [],
+  },
+  isNodeLoading: {
+    id: 'id1',
+    state: false,
+  },
+  deleted: false,
+  loadNodeHandler: jest.fn(),
+  addLineageHandler: jest.fn(),
+  removeLineageHandler: jest.fn(),
+  entityLineageHandler: jest.fn(),
+};
+
+jest.mock('../../utils/EntityLineageUtils', () => ({
+  dragHandle: jest.fn(),
+  getDataLabel: jest
+    .fn()
+    .mockReturnValue(<span data-testid="lineage-entity">datalabel</span>),
+  getDeletedLineagePlaceholder: jest
+    .fn()
+    .mockReturnValue(
+      <p>Lineage data is not available for deleted entities.</p>
+    ),
+  getHeaderLabel: jest.fn().mockReturnValue(<p>Header label</p>),
+  getLayoutedElements: jest.fn().mockReturnValue([]),
+  getLineageData: jest.fn().mockReturnValue([]),
+  getModalBodyText: jest.fn(),
+  onLoad: jest.fn(),
+  onNodeContextMenu: jest.fn(),
+  onNodeMouseEnter: jest.fn(),
+  onNodeMouseLeave: jest.fn(),
+  onNodeMouseMove: jest.fn(),
+}));
+
+jest.mock('../../utils/TableUtils', () => ({
+  getEntityIcon: jest.fn(),
+}));
+
+jest.mock('../../auth-provider/AuthProvider', () => ({
+  useAuthContext: jest.fn().mockReturnValue({ isAuthDisabled: true }),
+}));
+
+jest.mock('../../hooks/authHooks', () => ({
+  useAuth: jest.fn().mockReturnValue({
+    userPermissions: {},
+    isAdminUser: false,
+  }),
+}));
+
+describe('Test EntityLineage Component', () => {
+  it('Check if EntityLineage is rendering all the nodes', async () => {
+    const { container } = render(<EntityLineage {...mockEntityLineageProp} />, {
+      wrapper: MemoryRouter,
+    });
+
+    const lineageContainer = await findByTestId(container, 'lineage-container');
+    const reactFlowElement = await findByTestId(
+      container,
+      'react-flow-component'
+    );
+
+    expect(reactFlowElement).toBeInTheDocument();
+
+    expect(lineageContainer).toBeInTheDocument();
+  });
+
+  it('Check if EntityLineage has deleted as true', async () => {
+    const { container } = render(
+      <EntityLineage {...mockEntityLineageProp} deleted />,
+      {
+        wrapper: MemoryRouter,
+      }
+    );
+
+    const lineageContainer = queryByTestId(container, 'lineage-container');
+    const reactFlowElement = queryByTestId(container, 'react-flow-component');
+    const deletedMessage = await findByText(
+      container,
+      /Lineage data is not available for deleted entities/i
+    );
+
+    expect(deletedMessage).toBeInTheDocument();
+
+    expect(reactFlowElement).not.toBeInTheDocument();
+
+    expect(lineageContainer).not.toBeInTheDocument();
+  });
+});

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

@@ -11,6 +11,11 @@
  *  limitations under the License.
  */
 
+import {
+  faChevronLeft,
+  faChevronRight,
+} from '@fortawesome/free-solid-svg-icons';
+import { FontAwesomeIcon } from '@fortawesome/react-fontawesome';
 import dagre from 'dagre';
 import { LeafNodes, LineagePos, LoadingNodeState } from 'Models';
 import React, { MouseEvent as ReactMouseEvent } from 'react';
@@ -46,11 +51,6 @@ import { EntityReference } from '../generated/type/entityReference';
 import { getPartialNameFromFQN } from './CommonUtils';
 import { isLeafNode } from './EntityUtils';
 import { getEntityLink } from './TableUtils';
-import {
-  faChevronLeft,
-  faChevronRight,
-} from '@fortawesome/free-solid-svg-icons';
-import { FontAwesomeIcon } from '@fortawesome/react-fontawesome';
 
 export const getHeaderLabel = (
   v = '',
@@ -121,7 +121,10 @@ export const getLineageData = (
   ) => void
 ) => {
   const [x, y] = [0, 0];
-  const nodes = entityLineage['nodes'];
+  const nodes = [
+    ...(entityLineage['nodes'] as EntityReference[]),
+    entityLineage['entity'],
+  ];
   let upstreamEdges: Array<LineageEdge & { isMapped: boolean }> =
     entityLineage['upstreamEdges']?.map((up) => ({ isMapped: false, ...up })) ||
     [];
@@ -290,9 +293,16 @@ export const getLineageData = (
     return downNodesArr;
   };
 
-  getUpStreamData(mainNode);
+  /**
+   * Get upstream and downstream of each node and store it in
+   * UPStreamNodes
+   * DOWNStreamNodes
+   */
+  nodes?.forEach((node) => {
+    getUpStreamData(node);
 
-  getDownStreamData(mainNode);
+    getDownStreamData(node);
+  });
 
   const lineageData = [
     {

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

@@ -228,8 +228,14 @@ export const getEntityOverview = (
       return overview;
     }
     case EntityType.DASHBOARD: {
-      const { owner, tags, dashboardUrl, service, fullyQualifiedName } =
-        entityDetail;
+      const {
+        owner,
+        tags,
+        dashboardUrl,
+        service,
+        fullyQualifiedName,
+        displayName,
+      } = entityDetail;
       const ownerValue = getOwnerFromId(owner?.id);
       const tier = getTierFromTableTags(tags || []);
 
@@ -260,7 +266,7 @@ export const getEntityOverview = (
         },
         {
           name: `${serviceType} url`,
-          value: fullyQualifiedName?.split('.')[1] as string,
+          value: displayName || (fullyQualifiedName?.split('.')[1] as string),
           url: dashboardUrl as string,
           isLink: true,
           isExternal: true,