From f0cda8464f34a21f45f18fa557e980fb2f105d8e Mon Sep 17 00:00:00 2001
From: Karan Hotchandani <33024356+karanh37@users.noreply.github.com>
Date: Wed, 29 May 2024 12:10:02 +0530
Subject: [PATCH] Fix#16404 - Show Node level lineage by default (#16445)

* default to node layer

* update cypress

* code cleanup

* fix cypress
---
 .../ui/cypress/e2e/Flow/Lineage.spec.ts       |  23 ++--
 .../EntityLineage/CustomControls.test.tsx     |   2 -
 .../EntityLineage/CustomNodeV1.test.tsx       |   1 -
 .../LineageLayers/LineageLayers.test.tsx      |   1 -
 .../LineageLayers/LineageLayers.tsx           |   8 +-
 .../LineageLayers/lineage-layers.less         |   4 +
 .../EntityLineage/entity-lineage.style.less   |   1 -
 .../components/Lineage/Lineage.component.tsx  |   3 +
 .../ui/src/constants/Lineage.constants.ts     |   2 +-
 .../LineageProvider.interface.tsx             |   3 -
 .../LineageProvider/LineageProvider.tsx       | 120 +++++++-----------
 .../resources/ui/src/styles/variables.less    |   1 +
 .../ui/src/utils/EntityLineageUtils.tsx       |  22 +++-
 13 files changed, 96 insertions(+), 95 deletions(-)

diff --git a/openmetadata-ui/src/main/resources/ui/cypress/e2e/Flow/Lineage.spec.ts b/openmetadata-ui/src/main/resources/ui/cypress/e2e/Flow/Lineage.spec.ts
index 7252fa9a604..d5cd9baecb5 100644
--- a/openmetadata-ui/src/main/resources/ui/cypress/e2e/Flow/Lineage.spec.ts
+++ b/openmetadata-ui/src/main/resources/ui/cypress/e2e/Flow/Lineage.spec.ts
@@ -199,15 +199,18 @@ const addPipelineBetweenNodes = (
   }
 };
 
-const expandCols = (nodeFqn, hasShowMore) => {
-  cy.get(
-    `[data-testid="lineage-node-${nodeFqn}"] [data-testid="expand-cols-btn"]`
-  ).click({ force: true });
-  if (hasShowMore) {
-    cy.get(
-      `[data-testid="lineage-node-${nodeFqn}"] [data-testid="show-more-cols-btn"]`
-    ).click({ force: true });
-  }
+const activateColumnLayer = () => {
+  cy.get('[data-testid="lineage-layer-btn"]').click();
+  cy.get('[data-testid="lineage-layer-column-btn"]').click();
+};
+
+const verifyColumnLayerInactive = () => {
+  cy.get('[data-testid="lineage-layer-btn"]').click(); // Open Layer popover
+  cy.get('[data-testid="lineage-layer-column-btn"]').should(
+    'not.have.class',
+    'active'
+  );
+  cy.get('[data-testid="lineage-layer-btn"]').click(); // Close Layer popover
 };
 
 const addColumnLineage = (fromNode, toNode, exitEditMode = true) => {
@@ -243,6 +246,7 @@ describe('Lineage verification', { tags: 'DataAssets' }, () => {
       });
 
       cy.get('[data-testid="lineage"]').click();
+      verifyColumnLayerInactive();
       cy.get('[data-testid="edit-lineage"]').click();
 
       performZoomOut();
@@ -329,6 +333,7 @@ describe('Lineage verification', { tags: 'DataAssets' }, () => {
       const targetEntity = LINEAGE_ITEMS[i];
       if (targetEntity.columns.length > 0) {
         addPipelineBetweenNodes(sourceEntity, targetEntity);
+        activateColumnLayer();
         // Add column lineage
         addColumnLineage(sourceEntity, targetEntity);
         cy.get('[data-testid="edit-lineage"]').click();
diff --git a/openmetadata-ui/src/main/resources/ui/src/components/Entity/EntityLineage/CustomControls.test.tsx b/openmetadata-ui/src/main/resources/ui/src/components/Entity/EntityLineage/CustomControls.test.tsx
index 1b997cf1294..dddc1e27657 100644
--- a/openmetadata-ui/src/main/resources/ui/src/components/Entity/EntityLineage/CustomControls.test.tsx
+++ b/openmetadata-ui/src/main/resources/ui/src/components/Entity/EntityLineage/CustomControls.test.tsx
@@ -20,7 +20,6 @@ import CustomControlsComponent from './CustomControls.component';
 const mockOnOptionSelect = jest.fn();
 const mockOnLineageConfigUpdate = jest.fn();
 const mockOnEditLineageClick = jest.fn();
-const mockOnExpandColumnClick = jest.fn();
 const mockHandleFullScreenViewClick = jest.fn();
 const mockOnExitFullScreenViewClick = jest.fn();
 const mockOnZoomHandler = jest.fn();
@@ -50,7 +49,6 @@ jest.mock('reactflow', () => ({
 
 jest.mock('../../../context/LineageProvider/LineageProvider', () => ({
   useLineageProvider: jest.fn().mockImplementation(() => ({
-    toggleColumnView: mockOnExpandColumnClick,
     onLineageEditClick: mockOnEditLineageClick,
     onExportClick: mockOnExportClick,
     activeLayer: [LineageLayerView.COLUMN],
diff --git a/openmetadata-ui/src/main/resources/ui/src/components/Entity/EntityLineage/CustomNodeV1.test.tsx b/openmetadata-ui/src/main/resources/ui/src/components/Entity/EntityLineage/CustomNodeV1.test.tsx
index dea7a40979b..5c98ab7603d 100644
--- a/openmetadata-ui/src/main/resources/ui/src/components/Entity/EntityLineage/CustomNodeV1.test.tsx
+++ b/openmetadata-ui/src/main/resources/ui/src/components/Entity/EntityLineage/CustomNodeV1.test.tsx
@@ -87,7 +87,6 @@ jest.mock('../../../context/LineageProvider/LineageProvider', () => ({
       downstreamEdges: [],
     },
     activeLayer: [LineageLayerView.COLUMN],
-    expandedNodes: [],
     fetchPipelineStatus: jest.fn(),
     onColumnClick: onMockColumnClick,
   })),
diff --git a/openmetadata-ui/src/main/resources/ui/src/components/Entity/EntityLineage/LineageLayers/LineageLayers.test.tsx b/openmetadata-ui/src/main/resources/ui/src/components/Entity/EntityLineage/LineageLayers/LineageLayers.test.tsx
index 4a92dad514b..30f1f8ed954 100644
--- a/openmetadata-ui/src/main/resources/ui/src/components/Entity/EntityLineage/LineageLayers/LineageLayers.test.tsx
+++ b/openmetadata-ui/src/main/resources/ui/src/components/Entity/EntityLineage/LineageLayers/LineageLayers.test.tsx
@@ -61,7 +61,6 @@ jest.mock('../../../../context/LineageProvider/LineageProvider', () => ({
       downstreamEdges: [],
     },
     activeLayer: [],
-    expandedNodes: [],
     fetchPipelineStatus: jest.fn(),
     onColumnClick: onMockColumnClick,
     onUpdateLayerView: onMockUpdateLayerView,
diff --git a/openmetadata-ui/src/main/resources/ui/src/components/Entity/EntityLineage/LineageLayers/LineageLayers.tsx b/openmetadata-ui/src/main/resources/ui/src/components/Entity/EntityLineage/LineageLayers/LineageLayers.tsx
index a445309c3ae..dd7797ba5eb 100644
--- a/openmetadata-ui/src/main/resources/ui/src/components/Entity/EntityLineage/LineageLayers/LineageLayers.tsx
+++ b/openmetadata-ui/src/main/resources/ui/src/components/Entity/EntityLineage/LineageLayers/LineageLayers.tsx
@@ -24,7 +24,7 @@ import { getEntityIcon } from '../../../../utils/TableUtils';
 import './lineage-layers.less';
 
 const LineageLayers = () => {
-  const { activeLayer, onUpdateLayerView } = useLineageProvider();
+  const { activeLayer, onUpdateLayerView, isEditMode } = useLineageProvider();
 
   const onButtonClick = (value: LineageLayerView) => {
     const index = activeLayer.indexOf(value);
@@ -43,6 +43,7 @@ const LineageLayers = () => {
             className={classNames('lineage-layer-button h-15', {
               active: activeLayer.includes(LineageLayerView.COLUMN),
             })}
+            data-testid="lineage-layer-column-btn"
             onClick={() => onButtonClick(LineageLayerView.COLUMN)}>
             <div className="lineage-layer-btn">
               <div className="layer-icon">
@@ -59,6 +60,7 @@ const LineageLayers = () => {
                 LineageLayerView.DATA_OBSERVARABILITY
               ),
             })}
+            data-testid="lineage-layer-observability-btn"
             onClick={() =>
               onButtonClick(LineageLayerView.DATA_OBSERVARABILITY)
             }>
@@ -78,7 +80,9 @@ const LineageLayers = () => {
       trigger="click">
       <Button
         ghost
-        className="layers-btn h-15"
+        className={classNames('layers-btn h-15', {
+          'layers-btn-edit-mode': isEditMode,
+        })}
         data-testid="lineage-layer-btn"
         type="primary">
         <div className="lineage-layer-btn">
diff --git a/openmetadata-ui/src/main/resources/ui/src/components/Entity/EntityLineage/LineageLayers/lineage-layers.less b/openmetadata-ui/src/main/resources/ui/src/components/Entity/EntityLineage/LineageLayers/lineage-layers.less
index 30a4db23297..18f6882c72a 100644
--- a/openmetadata-ui/src/main/resources/ui/src/components/Entity/EntityLineage/LineageLayers/lineage-layers.less
+++ b/openmetadata-ui/src/main/resources/ui/src/components/Entity/EntityLineage/LineageLayers/lineage-layers.less
@@ -51,3 +51,7 @@
 .layers-btn {
   background-color: @white !important;
 }
+
+.layers-btn-edit-mode {
+  margin-left: @lineage-sidebar-width;
+}
diff --git a/openmetadata-ui/src/main/resources/ui/src/components/Entity/EntityLineage/entity-lineage.style.less b/openmetadata-ui/src/main/resources/ui/src/components/Entity/EntityLineage/entity-lineage.style.less
index cd3877a64c7..39115765704 100644
--- a/openmetadata-ui/src/main/resources/ui/src/components/Entity/EntityLineage/entity-lineage.style.less
+++ b/openmetadata-ui/src/main/resources/ui/src/components/Entity/EntityLineage/entity-lineage.style.less
@@ -13,7 +13,6 @@
 
 @import (reference) url('../../../styles/variables.less');
 
-@lineage-sidebar-width: 110px;
 @lineage-breadcrumb-panel: 50px;
 @lineage-toolbar-height: 54px;
 
diff --git a/openmetadata-ui/src/main/resources/ui/src/components/Lineage/Lineage.component.tsx b/openmetadata-ui/src/main/resources/ui/src/components/Lineage/Lineage.component.tsx
index e6dfba576e9..8cce8f4c6d4 100644
--- a/openmetadata-ui/src/main/resources/ui/src/components/Lineage/Lineage.component.tsx
+++ b/openmetadata-ui/src/main/resources/ui/src/components/Lineage/Lineage.component.tsx
@@ -156,6 +156,9 @@ const Lineage = ({
               deleteKeyCode={null}
               edgeTypes={customEdges}
               edges={edges}
+              fitViewOptions={{
+                padding: 48,
+              }}
               maxZoom={MAX_ZOOM_VALUE}
               minZoom={MIN_ZOOM_VALUE}
               nodeTypes={nodeTypes}
diff --git a/openmetadata-ui/src/main/resources/ui/src/constants/Lineage.constants.ts b/openmetadata-ui/src/main/resources/ui/src/constants/Lineage.constants.ts
index f58b1250e1f..c9d9db34fe3 100644
--- a/openmetadata-ui/src/main/resources/ui/src/constants/Lineage.constants.ts
+++ b/openmetadata-ui/src/main/resources/ui/src/constants/Lineage.constants.ts
@@ -19,7 +19,7 @@ import { SearchIndex } from '../enums/search.enum';
 import { Source } from '../generated/type/entityLineage';
 
 export const FOREIGN_OBJECT_SIZE = 40;
-export const ZOOM_VALUE = 0.75;
+export const ZOOM_VALUE = 0.65;
 export const MIN_ZOOM_VALUE = 0.1;
 export const MAX_ZOOM_VALUE = 2.5;
 export const ZOOM_SLIDER_STEP = 0.1;
diff --git a/openmetadata-ui/src/main/resources/ui/src/context/LineageProvider/LineageProvider.interface.tsx b/openmetadata-ui/src/main/resources/ui/src/context/LineageProvider/LineageProvider.interface.tsx
index f47ff1d6070..07fc54b8130 100644
--- a/openmetadata-ui/src/main/resources/ui/src/context/LineageProvider/LineageProvider.interface.tsx
+++ b/openmetadata-ui/src/main/resources/ui/src/context/LineageProvider/LineageProvider.interface.tsx
@@ -53,7 +53,6 @@ export interface LineageContextType {
   reactFlowInstance?: ReactFlowInstance;
   nodes: Node[];
   edges: Edge[];
-  expandedNodes: string[];
   tracedNodes: string[];
   tracedColumns: string[];
   lineageConfig: LineageConfig;
@@ -67,7 +66,6 @@ export interface LineageContextType {
   selectedNode: SourceType;
   upstreamDownstreamData: UpstreamDownstreamData;
   selectedColumn: string;
-  expandAllColumns: boolean;
   activeLayer: LineageLayerView[];
   onInitReactFlow: (reactFlowInstance: ReactFlowInstance) => void;
   onPaneClick: () => void;
@@ -83,7 +81,6 @@ export interface LineageContextType {
   onNodeCollapse: (node: Node | NodeProps, direction: EdgeTypeEnum) => void;
   onNodesChange: (changes: NodeChange[]) => void;
   onEdgesChange: (changes: EdgeChange[]) => void;
-  toggleColumnView: () => void;
   loadChildNodesHandler: (
     node: SourceType,
     direction: EdgeTypeEnum
diff --git a/openmetadata-ui/src/main/resources/ui/src/context/LineageProvider/LineageProvider.tsx b/openmetadata-ui/src/main/resources/ui/src/context/LineageProvider/LineageProvider.tsx
index 339f051797d..d84fb8fcefd 100644
--- a/openmetadata-ui/src/main/resources/ui/src/context/LineageProvider/LineageProvider.tsx
+++ b/openmetadata-ui/src/main/resources/ui/src/context/LineageProvider/LineageProvider.tsx
@@ -59,7 +59,6 @@ import {
 import { SourceType } from '../../components/SearchedData/SearchedData.interface';
 import {
   ELEMENT_DELETE_STATE,
-  ZOOM_TRANSITION_DURATION,
   ZOOM_VALUE,
 } from '../../constants/Lineage.constants';
 import { mockDatasetData } from '../../constants/mockTourData.constants';
@@ -78,6 +77,7 @@ import { useFqn } from '../../hooks/useFqn';
 import { getLineageDataByFQN, updateLineageEdge } from '../../rest/lineageAPI';
 import {
   addLineageHandler,
+  centerNodePosition,
   createEdges,
   createNewEdge,
   createNodes,
@@ -124,14 +124,10 @@ const LineageProvider = ({ children }: LineageProviderProps) => {
   const [selectedNode, setSelectedNode] = useState<SourceType>(
     {} as SourceType
   );
-  const [activeLayer, setActiveLayer] = useState<LineageLayerView[]>([
-    LineageLayerView.COLUMN,
-  ]);
+  const [activeLayer, setActiveLayer] = useState<LineageLayerView[]>([]);
   const [activeNode, setActiveNode] = useState<Node>();
   const [selectedColumn, setSelectedColumn] = useState<string>('');
   const [showAddEdgeModal, setShowAddEdgeModal] = useState<boolean>(false);
-  const [expandedNodes, setExpandedNodes] = useState<string[]>([]);
-  const [expandAllColumns, setExpandAllColumns] = useState(false);
   const [selectedEdge, setSelectedEdge] = useState<Edge>();
   const [entityLineage, setEntityLineage] = useState<EntityLineageResponse>({
     nodes: [],
@@ -679,10 +675,6 @@ const LineageProvider = ({ children }: LineageProviderProps) => {
   }, []);
 
   const onInitReactFlow = (reactFlowInstance: ReactFlowInstance) => {
-    setTimeout(() => {
-      onLoad(reactFlowInstance);
-    }, 0);
-
     setReactFlowInstance(reactFlowInstance);
   };
 
@@ -698,43 +690,6 @@ const LineageProvider = ({ children }: LineageProviderProps) => {
     setZoomValue(value);
   }, []);
 
-  const toggleColumnView = useCallback(() => {
-    const updatedVal = !expandAllColumns;
-    setExpandAllColumns(updatedVal);
-    setNodes((prevNodes) => {
-      const updatedNode = prevNodes.map((node) => {
-        const nodeId = node.data.node.id;
-
-        // Update the expandedNodes state based on the toggle value
-        if (updatedVal && !expandedNodes.includes(nodeId)) {
-          setExpandedNodes((prevExpandedNodes) => [
-            ...prevExpandedNodes,
-            nodeId,
-          ]);
-        } else if (!updatedVal) {
-          setExpandedNodes((prevExpandedNodes) =>
-            prevExpandedNodes.filter((id) => id !== nodeId)
-          );
-        }
-
-        return node;
-      });
-
-      const { edge, node } = getLayoutedElements(
-        {
-          node: updatedNode,
-          edge: edges,
-        },
-        EntityLineageDirection.LEFT_RIGHT,
-        updatedVal
-      );
-
-      setEdges(edge);
-
-      return node;
-    });
-  }, [expandAllColumns, expandedNodes, edges]);
-
   const onRemove = useCallback(async () => {
     try {
       setDeletionState({ ...ELEMENT_DELETE_STATE, loading: true });
@@ -1066,14 +1021,42 @@ const LineageProvider = ({ children }: LineageProviderProps) => {
   );
 
   const selectNode = (node: Node) => {
-    const { position } = node;
-    // moving selected node in center
-    reactFlowInstance?.setCenter(position.x, position.y, {
-      duration: ZOOM_TRANSITION_DURATION,
-      zoom: zoomValue,
-    });
+    centerNodePosition(node, reactFlowInstance);
   };
 
+  const repositionLayout = useCallback(
+    (activateNode = false) => {
+      const isColView = activeLayer.includes(LineageLayerView.COLUMN);
+      const { node, edge } = getLayoutedElements(
+        {
+          node: nodes,
+          edge: edges,
+        },
+        EntityLineageDirection.LEFT_RIGHT,
+        isColView
+      );
+
+      setNodes(node);
+      setEdges(edge);
+
+      const rootNode = node.find((n) => n.data.isRootNode);
+      if (!rootNode) {
+        if (activateNode && reactFlowInstance) {
+          onLoad(reactFlowInstance); // Call fitview in case of pipeline
+        }
+
+        return;
+      }
+
+      // Center the root node in the view
+      centerNodePosition(rootNode, reactFlowInstance);
+      if (activateNode) {
+        onNodeClick(rootNode);
+      }
+    },
+    [reactFlowInstance, activeLayer, nodes, edges, onNodeClick]
+  );
+
   const redrawLineage = useCallback(
     (lineageData: EntityLineageResponse) => {
       const allNodes = uniqWith(
@@ -1087,7 +1070,8 @@ const LineageProvider = ({ children }: LineageProviderProps) => {
       const updatedNodes = createNodes(
         allNodes,
         lineageData.edges ?? [],
-        decodedFqn
+        decodedFqn,
+        activeLayer.includes(LineageLayerView.COLUMN)
       );
       const updatedEdges = createEdges(
         allNodes,
@@ -1105,11 +1089,11 @@ const LineageProvider = ({ children }: LineageProviderProps) => {
       );
       setUpstreamDownstreamData(data);
 
-      if (activeNode) {
+      if (activeNode && !isEditMode) {
         selectNode(activeNode);
       }
     },
-    [decodedFqn, activeNode]
+    [decodedFqn, activeNode, activeLayer, isEditMode]
   );
 
   useEffect(() => {
@@ -1164,19 +1148,15 @@ const LineageProvider = ({ children }: LineageProviderProps) => {
   }, [isEditMode, deletePressed, backspacePressed, activeNode, selectedEdge]);
 
   useEffect(() => {
-    const { node, edge } = getLayoutedElements(
-      {
-        node: nodes,
-        edge: edges,
-      },
-      EntityLineageDirection.LEFT_RIGHT,
-      activeLayer.includes(LineageLayerView.COLUMN)
-    );
-
-    setNodes(node);
-    setEdges(edge);
+    repositionLayout();
   }, [activeLayer]);
 
+  useEffect(() => {
+    if (reactFlowInstance?.viewportInitialized) {
+      repositionLayout(true); // Activate the root node
+    }
+  }, [reactFlowInstance?.viewportInitialized]);
+
   const activityFeedContextValues = useMemo(() => {
     return {
       isDrawerOpen,
@@ -1191,10 +1171,8 @@ const LineageProvider = ({ children }: LineageProviderProps) => {
       selectedColumn,
       zoomValue,
       status,
-      expandedNodes,
       tracedNodes,
       tracedColumns,
-      expandAllColumns,
       upstreamDownstreamData,
       init,
       activeLayer,
@@ -1210,7 +1188,6 @@ const LineageProvider = ({ children }: LineageProviderProps) => {
       onZoomUpdate,
       updateEntityType,
       onDrawerClose,
-      toggleColumnView,
       loadChildNodesHandler,
       fetchLineageData,
       removeNodeHandler,
@@ -1236,10 +1213,8 @@ const LineageProvider = ({ children }: LineageProviderProps) => {
     selectedColumn,
     zoomValue,
     status,
-    expandedNodes,
     tracedNodes,
     tracedColumns,
-    expandAllColumns,
     upstreamDownstreamData,
     init,
     activeLayer,
@@ -1257,7 +1232,6 @@ const LineageProvider = ({ children }: LineageProviderProps) => {
     updateEntityType,
     loadChildNodesHandler,
     fetchLineageData,
-    toggleColumnView,
     removeNodeHandler,
     onNodeClick,
     onEdgeClick,
diff --git a/openmetadata-ui/src/main/resources/ui/src/styles/variables.less b/openmetadata-ui/src/main/resources/ui/src/styles/variables.less
index 4368d88d6ee..d069d641ae6 100644
--- a/openmetadata-ui/src/main/resources/ui/src/styles/variables.less
+++ b/openmetadata-ui/src/main/resources/ui/src/styles/variables.less
@@ -118,6 +118,7 @@
 @welcome-page-height: calc(100vh - 112px);
 @data-product-page-height: calc(100vh - 156px);
 @glossary-page-tab-height: calc(100vh - 206px);
+@lineage-sidebar-width: 110px;
 
 // Severity
 @severity-1: #9c0700;
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 2b25dafb73b..3638cde979b 100644
--- a/openmetadata-ui/src/main/resources/ui/src/utils/EntityLineageUtils.tsx
+++ b/openmetadata-ui/src/main/resources/ui/src/utils/EntityLineageUtils.tsx
@@ -68,6 +68,7 @@ import {
   LINEAGE_EXPORT_HEADERS,
   NODE_HEIGHT,
   NODE_WIDTH,
+  ZOOM_TRANSITION_DURATION,
   ZOOM_VALUE,
 } from '../constants/Lineage.constants';
 import {
@@ -120,6 +121,22 @@ export const onLoad = (reactFlowInstance: ReactFlowInstance) => {
   reactFlowInstance.fitView();
   reactFlowInstance.zoomTo(ZOOM_VALUE);
 };
+
+export const centerNodePosition = (
+  node: Node,
+  reactFlowInstance?: ReactFlowInstance
+) => {
+  const { position, width } = node;
+  reactFlowInstance?.setCenter(
+    position.x + (width ?? 1 / 2),
+    position.y + NODE_HEIGHT / 2,
+    {
+      zoom: ZOOM_VALUE,
+      duration: ZOOM_TRANSITION_DURATION,
+    }
+  );
+};
+
 /* eslint-disable-next-line */
 export const onNodeMouseEnter = (_event: ReactMouseEvent, _node: Node) => {
   return;
@@ -631,7 +648,8 @@ const getNodeType = (
 export const createNodes = (
   nodesData: EntityReference[],
   edgesData: EdgeDetails[],
-  entityFqn: string
+  entityFqn: string,
+  isExpanded = false
 ) => {
   const uniqueNodesData = removeDuplicateNodes(nodesData).sort((a, b) =>
     getEntityName(a).localeCompare(getEntityName(b))
@@ -651,7 +669,7 @@ export const createNodes = (
   // Add nodes to the graph
   uniqueNodesData.forEach((node) => {
     const { childrenHeight } = getEntityChildrenAndLabel(node as SourceType);
-    const nodeHeight = childrenHeight + 220;
+    const nodeHeight = isExpanded ? childrenHeight + 220 : NODE_HEIGHT;
     graph.setNode(node.id, { width: NODE_WIDTH, height: nodeHeight });
   });