Fix column expand freeze

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

@@ -46,6 +46,7 @@ import { SearchIndex } from '../../../enums/search.enum';
 import { LineageDirection } from '../../../generated/api/lineage/entityCountLineageRequest';
 import useCustomLocation from '../../../hooks/useCustomLocation/useCustomLocation';
 import { useFqn } from '../../../hooks/useFqn';
+import { useLineageStore } from '../../../hooks/useLineageStore';
 import { QueryFieldInterface } from '../../../pages/ExplorePage/ExplorePage.interface';
 import { exportLineageByEntityCountAsync } from '../../../rest/lineageAPI';
 import { getQuickFilterQuery } from '../../../utils/ExploreUtils';
@@ -83,13 +84,17 @@ const CustomControls: FC<{
     setSelectedQuickFilters,
     nodes,
     selectedQuickFilters,
-    lineageConfig,
     onExportClick,
-    onLineageConfigUpdate,
-    onLineageEditClick,
+  } = useLineageProvider();
+
+  const {
+    lineageConfig,
     isEditMode,
+    setLineageConfig,
+    toggleEditMode,
     platformView,
-  } = useLineageProvider();
+  } = useLineageStore();
+
   const [filterSelectionActive, setFilterSelectionActive] = useState(false);
   const [dialogVisible, setDialogVisible] = useState(false);
   const [nodeDepthAnchorEl, setNodeDepthAnchorEl] =
@@ -297,7 +302,7 @@ const CustomControls: FC<{
 
   const handleDialogSave = (newConfig: LineageConfig) => {
     // Implement save logic here
-    onLineageConfigUpdate?.(newConfig);
+    setLineageConfig?.(newConfig);
     setDialogVisible(false);
   };
 
@@ -363,20 +368,12 @@ const CustomControls: FC<{
           color={isEditMode ? 'primary' : 'default'}
           data-testid="edit-lineage"
           size="large"
-          onClick={onLineageEditClick}>
+          onClick={toggleEditMode}>
           <EditIcon />
         </StyledIconButton>
       </Tooltip>
     ) : null;
-  }, [
-    hasEditAccess,
-    deleted,
-    platformView,
-    entityType,
-    isEditMode,
-    onLineageEditClick,
-    t,
-  ]);
+  }, [hasEditAccess, deleted, platformView, entityType, isEditMode, t]);
 
   const settingsButton = useMemo(() => {
     const handleSettingsClick = () => {

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

@@ -25,7 +25,7 @@ import { FOREIGN_OBJECT_SIZE } from '../../../constants/Lineage.constants';
 import { useLineageProvider } from '../../../context/LineageProvider/LineageProvider';
 import { EntityType } from '../../../enums/entity.enum';
 import { StatusType } from '../../../generated/entity/data/pipeline';
-import { LineageLayer } from '../../../generated/settings/settings';
+import { useLineageStore } from '../../../hooks/useLineageStore';
 import {
   getColumnSourceTargetHandles,
   getEdgePathData,
@@ -87,18 +87,20 @@ export const CustomEdge = ({
   const { fromEntity, toEntity, pipeline, pipelineEntityType } = edge;
 
   const {
-    tracedNodes,
-    tracedColumns,
-    isEditMode,
-    activeLayer,
     onAddPipelineClick,
     onColumnEdgeRemove,
     dataQualityLineage,
     dqHighlightedEdges,
-    selectedColumn,
-    columnsInCurrentPages,
   } = useLineageProvider();
 
+  const {
+    tracedNodes,
+    tracedColumns,
+    isEditMode,
+    isDQEnabled,
+    selectedColumn,
+  } = useLineageStore();
+
   const theme = useTheme();
 
   // Get edge path data once
@@ -126,15 +128,12 @@ export const CustomEdge = ({
 
   // Compute if should show DQ tracing
   const showDqTracing = useMemo(() => {
-    if (
-      !activeLayer.includes(LineageLayer.DataObservability) ||
-      !dataQualityLineage?.nodes
-    ) {
+    if (!isDQEnabled || !dataQualityLineage?.nodes) {
       return false;
     }
 
     return dqHighlightedEdges?.has(id);
-  }, [activeLayer, dataQualityLineage?.nodes, id, dqHighlightedEdges]);
+  }, [isDQEnabled, dataQualityLineage?.nodes, id, dqHighlightedEdges]);
 
   // Determine if column is highlighted based on traced columns
   const isColumnHighlighted = useMemo(() => {
@@ -148,57 +147,33 @@ export const CustomEdge = ({
     });
 
     return (
-      tracedColumns.includes(decodedHandles.sourceHandle ?? '') &&
-      tracedColumns.includes(decodedHandles.targetHandle ?? '')
+      tracedColumns.has(decodedHandles.sourceHandle ?? '') &&
+      tracedColumns.has(decodedHandles.targetHandle ?? '')
     );
   }, [isColumnLineage, tracedColumns, sourceHandle, targetHandle]);
 
-  const areBothColumnHandlesPresentInCurrentPage = useMemo(() => {
-    const decodedHandles = getColumnSourceTargetHandles({
-      sourceHandle,
-      targetHandle,
-    });
-
-    const allColumnsInCurrentPages = new Set(
-      Object.values(columnsInCurrentPages).flat()
-    );
-
-    return (
-      allColumnsInCurrentPages.has(decodedHandles.sourceHandle ?? '') &&
-      allColumnsInCurrentPages.has(decodedHandles.targetHandle ?? '')
-    );
-  }, [columnsInCurrentPages, sourceHandle, targetHandle]);
-
   // Calculate edge style with memoization
   const updatedStyle = useMemo(() => {
     const isNodeTraced =
-      tracedNodes.includes(edge.fromEntity.id) &&
-      tracedNodes.includes(edge.toEntity.id);
+      tracedNodes.has(edge.fromEntity.id) && tracedNodes.has(edge.toEntity.id);
 
     let stroke = '';
-    let display = 'block';
     let opacity = 1;
 
     // For nodes edges
     if (isNodeTraced) {
       stroke = theme.palette.primary.main;
-    } else if (!(tracedNodes.length === 0) || !(tracedColumns.length === 0)) {
+    } else if (!(tracedNodes.size === 0) || !(tracedColumns.size === 0)) {
       opacity = 0.3;
     }
 
     // For columns edges
     if (isColumnLineage) {
-      display = 'none';
-      const noTracing = tracedNodes.length === 0 && tracedColumns.length === 0;
-
       if (isColumnHighlighted) {
-        display = 'block';
         stroke = selectedColumn
           ? theme.palette.allShades.indigo[600]
           : theme.palette.primary.main;
         opacity = 1;
-      } else if (noTracing && areBothColumnHandlesPresentInCurrentPage) {
-        display = 'block';
       }
     }
 
@@ -210,21 +185,19 @@ export const CustomEdge = ({
       ...style,
       stroke,
       opacity,
-      display,
     };
   }, [
     tracedNodes,
     edge.fromEntity.id,
     edge.toEntity.id,
-    tracedColumns.length,
+    tracedColumns.size,
     isColumnLineage,
     showDqTracing,
     style,
     theme.palette.primary.main,
     theme.palette.allShades.indigo,
     theme.palette.allShades.error,
     isColumnHighlighted,
-    areBothColumnHandlesPresentInCurrentPage,
     selectedColumn,
   ]);
 
@@ -255,12 +228,9 @@ export const CustomEdge = ({
 
   // Calculate pipeline status for styling
   const currentPipelineStatus = useMemo(() => {
-    const isPipelineActiveNow = activeLayer.includes(
-      LineageLayer.DataObservability
-    );
     const pipelineData = pipeline?.pipelineStatus;
 
-    if (pipelineData && isPipelineActiveNow) {
+    if (pipelineData && isDQEnabled) {
       switch (pipelineData.executionStatus) {
         case StatusType.Failed:
           return 'red';
@@ -275,7 +245,7 @@ export const CustomEdge = ({
     }
 
     return '';
-  }, [pipeline?.pipelineStatus, activeLayer]);
+  }, [pipeline?.pipelineStatus, isDQEnabled]);
 
   // Calculate blinking class for pipeline nodes
   const blinkingClass = useMemo(() => {