Fix column expand freeze

[chirag-madlani]

Describe your changes:

Fixes

I worked on ... because ...

Type of change:

• Bug fix
• Improvement
• New feature
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Documentation

Checklist:

• I have read the CONTRIBUTING document.
• My PR title is Fixes <issue-number>: <short explanation>
• I have commented on my code, particularly in hard-to-understand areas.
• For JSON Schema changes: I updated the migration scripts or explained why it is not needed.

---

Summary by Gitar

• State management refactor:
  • Extracted UI state from LineageProvider Context into new useLineageStore.ts Zustand store to prevent unnecessary re-renders
• New hook:
  • Created useMapBasedNodesEdges hook with Map-based state management for O(1) node/edge operations (add, remove, update)
• Performance optimization:
  • Changed columnsHavingLineage from array to Set for O(1) lookups and replaced array operations with Map-based lookups
• React optimization:
  • Wrapped EntityTypeIcon and EntityFooter in React.memo and optimized calculateHeightAndFlattenNode with accumulator pattern

_{This will update automatically on new commits.}

---

[gitar-bot]

🚨 Bug: onPaneClick passes arrays instead of Sets to store setters

In onPaneClick (LineageProvider.tsx), setTracedNodes([]) and setTracedColumns([]) pass empty arrays, but the Zustand store now expects Set<string> for both. All other call sites were properly updated to new Set(), but these two were missed.

This will cause a TypeScript compilation error, and if it somehow reaches runtime (e.g., via any casts), every call to .has() or .size on tracedColumns/tracedNodes throughout the codebase would crash.

Fix:

const onPaneClick = useCallback(() => {
  setTracedNodes(new Set());
  setTracedColumns(new Set());
  setSelectedColumn('');
  setActiveNode(undefined);
  setSelectedNode(undefined);
}, []);

_{Was this helpful? React with 👍 / 👎}

[gitar-bot]

⚠️ Bug: toggleEditMode doesn't clear local selectedEdge state

The old onLineageEditClick in LineageProvider explicitly called setSelectedEdge(undefined) (local state) to close the drawer when toggling edit mode. The new toggleEditMode in the Zustand store sets selectedEdge: undefined on the store's copy, but LineageProvider still manages selectedEdge as local React state (line 196: const [selectedEdge, setSelectedEdge] = useState<Edge>()). The store's setter and the provider's local setter are different.

This means when a user toggles edit mode while an edge drawer is open, the drawer won't close because the local selectedEdge is never cleared.

Fix: Either:

1. Move selectedEdge fully into the Zustand store and consume it from there in the provider, OR
2. Have the CustomControls component call a provider callback (like the old onLineageEditClick) that clears local state alongside toggling the store's edit mode.

_{Was this helpful? React with 👍 / 👎}

[gitar-bot]

Code Review ⚠️ Changes requested 8 resolved / 11 findings

Solid state management refactor with two bugs: onPaneClick passes arrays instead of Sets to Zustand setters (type error / runtime crash), and toggleEditMode in the store doesn't clear the locally-managed selectedEdge, leaving the drawer open when toggling edit mode.

🚨 Bug: onPaneClick passes arrays instead of Sets to store setters

📄 openmetadata-ui/src/main/resources/ui/src/context/LineageProvider/LineageProvider.tsx:1254

In onPaneClick (LineageProvider.tsx), setTracedNodes([]) and setTracedColumns([]) pass empty arrays, but the Zustand store now expects Set<string> for both. All other call sites were properly updated to new Set(), but these two were missed.

This will cause a TypeScript compilation error, and if it somehow reaches runtime (e.g., via any casts), every call to .has() or .size on tracedColumns/tracedNodes throughout the codebase would crash.

Fix:

const onPaneClick = useCallback(() => {
  setTracedNodes(new Set());
  setTracedColumns(new Set());
  setSelectedColumn('');
  setActiveNode(undefined);
  setSelectedNode(undefined);
}, []);

⚠️ Bug: toggleEditMode doesn't clear local selectedEdge state

📄 openmetadata-ui/src/main/resources/ui/src/hooks/useLineageStore.ts:84 📄 openmetadata-ui/src/main/resources/ui/src/context/LineageProvider/LineageProvider.tsx:196

The old onLineageEditClick in LineageProvider explicitly called setSelectedEdge(undefined) (local state) to close the drawer when toggling edit mode. The new toggleEditMode in the Zustand store sets selectedEdge: undefined on the store's copy, but LineageProvider still manages selectedEdge as local React state (line 196: const [selectedEdge, setSelectedEdge] = useState<Edge>()). The store's setter and the provider's local setter are different.

This means when a user toggles edit mode while an edge drawer is open, the drawer won't close because the local selectedEdge is never cleared.

Fix: Either:

1. Move selectedEdge fully into the Zustand store and consume it from there in the provider, OR
2. Have the CustomControls component call a provider callback (like the old onLineageEditClick) that clears local state alongside toggling the store's edit mode.

💡 Edge Case: onNodesChange drops 'add' and 'reset' change types silently

📄 openmetadata-ui/src/main/resources/ui/src/hooks/useMapBasedNodesEdges.ts:175

The custom onNodesChange in useMapBasedNodesEdges.ts only handles 'remove', 'position', 'dimensions', and 'select' change types. ReactFlow's NodeChange union also includes 'add' and 'reset'. Similarly, onEdgesChange only handles 'remove' and 'select', missing 'add' and 'reset'.

In the current codebase, nodes/edges are added via dedicated addNodes/setNodes calls so 'add' changes likely never fire, and 'reset' is typically internal. However, if ReactFlow internals or future code paths emit these change types, they'll be silently dropped.

Consider adding a no-op comment for the unhandled cases or implementing them defensively:

case 'add':
  next.set(change.item.id, change.item);
  break;
case 'reset':
  return new Map(change.items?.map(n => [n.id, n]) ?? []);

✅ 8 resolved

✅ Bug: Empty Drawer rendered without open prop or children

📄 openmetadata-ui/src/main/resources/ui/src/context/LineageProvider/LineageProvider.tsx:2075
After the refactor, the first Drawer block (rendered when isEditMode && selectedEdge) became an empty self-closing <Drawer /> with no open prop and no children. The EdgeInfoDrawer that was previously rendered inside this Drawer was moved into the second Drawer block.

This means:

1. When isEditMode && selectedEdge is truthy, an empty Drawer mounts with open defaulting to false (MUI default) — it does nothing useful.
2. The intent of showing an edge info drawer in edit mode is broken since there's no content.

Suggested fix: Remove this entire first Drawer block since EdgeInfoDrawer is now rendered in the second unified Drawer. The second Drawer already handles selectedEdge rendering via the {selectedEdge && <EdgeInfoDrawer ... />} branch. However, note the second Drawer only renders when !isEditMode, so if edge info should be shown during edit mode, the condition on the second Drawer needs updating too.

✅ Performance: tracedColumns.includes() is O(n) per call in renderSentinel

📄 openmetadata-ui/src/main/resources/ui/src/hooks/useLineageStore.ts:26 📄 openmetadata-ui/src/main/resources/ui/src/components/Entity/EntityLineage/NodeChildren/NodeChildren.component.tsx:619
tracedColumns is stored as an array (string[]) in the Zustand store, and renderSentinel calls tracedColumns.includes(fullyQualifiedName ?? '') for every off-page column on every render. Similarly, tracedColumns.includes() is called in other hot paths like renderColumnsData.

Given that this PR already converted columnsHavingLineage from string[] to Set<string> for O(1) lookups, consider applying the same optimization to tracedColumns (and tracedNodes). This would be especially beneficial for lineage graphs with many columns where the traced columns list could grow.

Suggested fix: Change tracedColumns and tracedNodes from string[] to Set<string> in the store, and use .has() instead of .includes() in all lookup sites.

✅ Edge Case: columnInfoDropdown misses handleClickColumnInfoDropdown dep

📄 openmetadata-ui/src/main/resources/ui/src/components/Entity/EntityLineage/LineageNodeLabelV1.tsx:230
In the same columnInfoDropdown useMemo, the onClick handler references handleClickColumnInfoDropdown which is a useCallback whose identity can change (it depends on toggleColumnsList). Since handleClickColumnInfoDropdown is not in the dependency array, the button could reference a stale click handler after toggleColumnsList changes.

Suggested fix: Add handleClickColumnInfoDropdown to the dependency array:

}, [childrenCount, childrenHeading, isChildrenListExpanded, handleClickColumnInfoDropdown]);

✅ Bug: Commented-out updateNodeInternals may break node rendering

📄 openmetadata-ui/src/main/resources/ui/src/components/Entity/EntityLineage/NodeChildren/NodeChildren.component.tsx
The call to updateNodeInternals?.(node.id) has been commented out without explanation:

if (node.id) {
   //   updateNodeInternals?.(node.id);
}

In ReactFlow, updateNodeInternals is important to call when node dimensions or handles change - it tells ReactFlow to recalculate the node's internal state including handle positions. Commenting this out could cause issues where:

1. Handle positions don't update correctly when columns are selected/traced
2. Edges may not connect to the correct positions after column changes
3. Visual glitches when column selection state changes

If this was intentionally disabled for performance reasons, this should be documented with a comment explaining why, or the entire useEffect block should be removed if it's no longer needed.

Suggested fix: Either restore the call if it's needed:

if (node.id) {
-   //   updateNodeInternals?.(node.id);
+   updateNodeInternals?.(node.id);
}

Or remove the entire useEffect block if it's truly not needed anymore.

✅ Bug: Commented-out code left in EntityFooter component

📄 openmetadata-ui/src/main/resources/ui/src/components/Entity/EntityLineage/LineageNodeLabelV1.tsx
There is commented-out code // const { isEditMode } = useLineageProvider(); left in the EntityFooter component. This appears to be debug/refactoring leftover that should be removed to keep the codebase clean. While harmless, it creates confusion about whether this code is needed or was intentionally removed.

Suggested fix: Remove the commented-out line:

-    //   const { isEditMode } = useLineageProvider();

...and 3 more resolved from earlier reviews

Rules 🎸 1 action taken

Gitar Rules

🎸 Summary Enhancement: Updated PR description summary to include new useMapBasedNodesEdges hook and revised performance optimization details

_{2 rules not applicable. Show all rules by commenting gitar display:verbose.}

Options

Auto-apply is off → Gitar will not commit updates to this branch.
Display: compact → Showing less information.

Comment with these commands to change:

Auto-apply	Compact
gitar auto-apply:on	gitar display:verbose

_{Was this helpful? React with 👍 / 👎 | Gitar}