From f02e0dd665900f1c9f1fd0668a29a4729958988f Mon Sep 17 00:00:00 2001
From: Max Mykal <mykalmax@yahoo.com>
Date: Thu, 2 Oct 2025 11:41:47 -0700
Subject: [PATCH 1/2] add option to toggle node dragging

---
 .../src/components/Lineage/LineageLayout.tsx  | 60 ++++++++++++++++--
 .../Lineage/stories/Lineage.stories.tsx       | 63 ++++++++++++++-----
 .../Lineage/stories/ModelLineage.tsx          | 11 +++-
 3 files changed, 113 insertions(+), 21 deletions(-)

diff --git a/web/common/src/components/Lineage/LineageLayout.tsx b/web/common/src/components/Lineage/LineageLayout.tsx
index 411ace4e65..42800dfefa 100644
--- a/web/common/src/components/Lineage/LineageLayout.tsx
+++ b/web/common/src/components/Lineage/LineageLayout.tsx
@@ -2,7 +2,9 @@ import {
   Background,
   BackgroundVariant,
   Controls,
+  type EdgeChange,
   type EdgeTypes,
+  type NodeChange,
   type NodeTypes,
   ReactFlow,
   ReactFlowProvider,
@@ -12,6 +14,10 @@ import {
   getOutgoers,
   useReactFlow,
   useViewport,
+  useEdgesState,
+  useNodesState,
+  applyNodeChanges,
+  applyEdgeChanges,
 } from '@xyflow/react'
 
 import '@xyflow/react/dist/style.css'
@@ -55,6 +61,8 @@ export function LineageLayout<
   edgeTypes,
   className,
   controls,
+  nodesDraggable,
+  nodesConnectable,
   useLineage,
   onNodeClick,
   onNodeDoubleClick,
@@ -69,6 +77,8 @@ export function LineageLayout<
   nodeTypes?: NodeTypes
   edgeTypes?: EdgeTypes
   className?: string
+  nodesDraggable?: boolean
+  nodesConnectable?: boolean
   controls?:
     | React.ReactNode
     | (({ setCenter }: { setCenter: SetCenter }) => React.ReactNode)
@@ -86,6 +96,8 @@ export function LineageLayout<
       <LineageLayoutBase
         nodeTypes={nodeTypes}
         edgeTypes={edgeTypes}
+        nodesDraggable={nodesDraggable}
+        nodesConnectable={nodesConnectable}
         className={className}
         controls={controls}
         useLineage={useLineage}
@@ -107,6 +119,8 @@ function LineageLayoutBase<
   edgeTypes,
   className,
   controls,
+  nodesDraggable = false,
+  nodesConnectable = false,
   useLineage,
   onNodeClick,
   onNodeDoubleClick,
@@ -118,6 +132,8 @@ function LineageLayoutBase<
     TEdgeID,
     TPortID
   >
+  nodesDraggable?: boolean
+  nodesConnectable?: boolean
   nodeTypes?: NodeTypes
   edgeTypes?: EdgeTypes
   className?: string
@@ -140,8 +156,8 @@ function LineageLayoutBase<
     isBuildingLayout,
     currentNode,
     zoom,
-    nodes,
-    edges,
+    nodes: initialNodes,
+    edges: initialEdges,
     nodesMap,
     showOnlySelectedNodes,
     selectedNodeId,
@@ -152,6 +168,32 @@ function LineageLayoutBase<
     setSelectedEdges,
   } = useLineage()
 
+  const [nodes, setNodes] = React.useState(initialNodes)
+  const [edges, setEdges] = React.useState(initialEdges)
+
+  const onNodesChange = React.useCallback(
+    (changes: NodeChange<LineageNode<TNodeData, TNodeID>>[]) => {
+      setNodes(
+        applyNodeChanges<LineageNode<TNodeData, TNodeID>>(changes, nodes),
+      )
+    },
+    [nodes, setNodes],
+  )
+
+  const onEdgesChange = React.useCallback(
+    (
+      changes: EdgeChange<LineageEdge<TEdgeData, TNodeID, TEdgeID, TPortID>>[],
+    ) => {
+      setEdges(
+        applyEdgeChanges<LineageEdge<TEdgeData, TNodeID, TEdgeID, TPortID>>(
+          changes,
+          edges,
+        ),
+      )
+    },
+    [edges, setEdges],
+  )
+
   const updateZoom = React.useMemo(() => debounce(setZoom, 200), [setZoom])
 
   const zoomToCurrentNode = React.useCallback(
@@ -221,6 +263,14 @@ function LineageLayoutBase<
     [nodes, edges],
   )
 
+  React.useEffect(() => {
+    setNodes(initialNodes)
+  }, [initialNodes])
+
+  React.useEffect(() => {
+    setEdges(initialEdges)
+  }, [initialEdges])
+
   React.useEffect(() => {
     if (selectedNodeId == null) {
       setShowOnlySelectedNodes(false)
@@ -332,8 +382,10 @@ function LineageLayoutBase<
         edges={edges}
         nodeTypes={nodeTypes}
         edgeTypes={edgeTypes}
-        nodesDraggable={false}
-        nodesConnectable={false}
+        onNodesChange={onNodesChange}
+        onEdgesChange={onEdgesChange}
+        nodesDraggable={nodesDraggable}
+        nodesConnectable={nodesConnectable}
         zoomOnDoubleClick={false}
         panOnScroll={true}
         zoomOnScroll={true}
diff --git a/web/common/src/components/Lineage/stories/Lineage.stories.tsx b/web/common/src/components/Lineage/stories/Lineage.stories.tsx
index 4ad8ca9f8b..6e16bed61e 100644
--- a/web/common/src/components/Lineage/stories/Lineage.stories.tsx
+++ b/web/common/src/components/Lineage/stories/Lineage.stories.tsx
@@ -17,33 +17,56 @@ export const LineageModel = () => {
     >
       <style>{`
         :root {
+          --color-metadata-label: rgba(100, 100, 100, 1);
+          --color-metadata-value: rgba(10, 10, 10, 1);
 
-          --color-lineage-control-button-tooltip-background: rgba(150, 150, 150, 1);
-          --color-lineage-control-button-tooltip-foreground: rgba(255, 255, 255, 1);
+          --color-tooltip-background: rgba(150, 150, 150, 1);
+          --color-tooltip-foreground: rgba(255, 255, 255, 1);
 
-          --color-lineage-divider: rgba(0, 0, 0, 0.1);
-          --color-lineage-background: rgba(255, 255, 255, 1);
-          --color-lineage-border: rgba(0, 0, 0, 0.1);
+          --color-filterable-list-counter-background: rgba(200, 0, 0, 1);
+          --color-filterable-list-counter-foreground: rgba(255, 255, 255, 1);
+          --color-filterable-list-input-background: rgba(250, 250, 250, 1);
+          --color-filterable-list-input-foreground: rgba(0, 0, 0, 1);
+          --color-filterable-list-input-placeholder: rgba(100, 100, 100, 1);
+          --color-filterable-list-input-border: rgba(100, 100, 100, 1);
 
           --color-lineage-control-background: rgba(250, 250, 250, 1);
           --color-lineage-control-background-hover: rgba(245, 245, 245, 1);
           --color-lineage-control-icon-background: rgba(0, 0, 0, 1);
           --color-lineage-control-icon-foreground: rgba(255, 255, 255, 1);
+          --color-lineage-control-button-tooltip-background: rgba(150, 150, 150, 1);
+          --color-lineage-control-button-tooltip-foreground: rgba(255, 255, 255, 1);
 
-          --color-lineage-grid-dot: rgba(0, 0, 0, 1);
+          --color-model-name-grayscale-link-underline: rgba(125, 125, 125, 1);
+          --color-model-name-grayscale-link-underline-hover: rgba(125, 125, 125, 1);
+          --color-model-name-link-underline: rgba(0, 0, 0, 1);
+          --color-model-name-link-underline-hover: rgba(0, 0, 0, 1);
+          --color-model-name-catalog: rgba(0, 0, 0, 1);
+          --color-model-name-schema: rgba(0, 0, 0, 1);
+          --color-model-name-model: rgba(0, 0, 0, 1);
+          --color-model-name-grayscale-catalog: rgba(100, 100, 100, 1);
+          --color-model-name-grayscale-schema: rgba(50, 50, 50, 1);
+          --color-model-name-grayscale-model: rgba(10, 10, 10, 1);
+          --color-model-name-copy-icon: rgba(100, 100, 100, 1);
+          --color-model-name-copy-icon-hover: rgba(125, 125, 125, 1);
+
+          --color-lineage-background: rgba(255, 255, 255, 1);
+          --color-lineage-divider: rgba(0, 0, 0, 0.1);
+          --color-lineage-border: rgba(0, 0, 0, 0.1);
+
+          --color-lineage-grid-dot: rgba(0, 0, 0, 0.1);
 
-          --color-lineage-node-appendix-background: transparent;
-          
           --color-lineage-node-background: rgba(255, 255, 255, 1);
           --color-lineage-node-foreground: rgba(0, 0, 0, 0.75);
+          --color-lineage-node-selected-border: rgba(0, 120, 120, 0.5);
           --color-lineage-node-border: rgba(0, 0, 0, 0.1);
           --color-lineage-node-border-hover: rgba(0, 0, 0, 0.2);
 
-          --color-lineage-node-selected-border: rgba(0, 120, 120, 0.5);
-
-          --color-lineage-node-badge-background: rgba(240, 240, 240, 1);
+          --color-lineage-node-badge-background: rgba(200, 200, 200, 1);
           --color-lineage-node-badge-foreground: rgba(0, 0, 0, 1);
 
+          --color-lineage-node-appendix-background: transparent;
+
           --color-lineage-node-type-background-sql: rgba(0, 0, 120, 1);
           --color-lineage-node-type-foreground-sql: rgba(0, 0, 120, 1);
           --color-lineage-node-type-border-sql: rgba(0, 0, 120, 1);
@@ -52,7 +75,18 @@ export const LineageModel = () => {
           --color-lineage-node-type-foreground-python: rgba(120, 0, 120, 1);
           --color-lineage-node-type-border-python: rgba(120, 0, 120, 1);
 
+          --color-lineage-node-type-background-source: rgba(120, 120, 0, 1);
+          --color-lineage-node-type-foreground-source: rgba(120, 120, 0, 1);
+          --color-lineage-node-type-border-source: rgba(120, 120, 0, 1);
+
+          --color-lineage-node-type-background-cte-subquery: rgba(120, 120, 120, 1);
+          --color-lineage-node-type-foreground-cte-subquery: rgba(120, 120, 120, 1);
+          --color-lineage-node-type-border-cte-subquery: rgba(120, 120, 120, 1);
+
           --color-lineage-node-type-handle-icon-background: rgba(255, 255, 255, 1);
+          --color-lineage-node-type-handle-icon-foreground: rgba(0, 0, 0, 1);
+
+          --color-lineage-edge: rgba(0, 0, 0, 0.1);
 
           --color-lineage-node-port-background: rgba(70, 0, 0, 0.05);
           --color-lineage-node-port-handle-source: rgba(70, 0, 0, 1);
@@ -61,15 +95,12 @@ export const LineageModel = () => {
           --color-lineage-node-port-edge-target: rgba(130, 0, 0, 1);
 
           --color-lineage-model-column-error-background: rgba(255, 0, 0, 1);
-          --color-lineage-model-column-source-background: rgba(200, 0, 0, 1);
-          --color-lineage-model-column-expression-background: rgba(100, 0, 0, 1);
+          --color-lineage-model-column-source-background: rgba(0, 200, 200, 1);
+          --color-lineage-model-column-expression-background: rgba(0, 10, 100, 1);
           --color-lineage-model-column-error-icon: rgba(255, 0, 0, 1);
           --color-lineage-model-column-active: rgba(70, 0, 0, 0.1);
           --color-lineage-model-column-icon: rgba(0, 0, 0, 1);
           --color-lineage-model-column-icon-active: rgba(0, 0, 0, 1);
-
-          --color-filterable-list-counter-background: rgba(200, 0, 0, 1);
-          --color-filterable-list-counter-foreground: rgba(200, 200, 200, 1);
         }
       `}</style>
       <ModelLineage
diff --git a/web/common/src/components/Lineage/stories/ModelLineage.tsx b/web/common/src/components/Lineage/stories/ModelLineage.tsx
index d704b6a209..2902350919 100644
--- a/web/common/src/components/Lineage/stories/ModelLineage.tsx
+++ b/web/common/src/components/Lineage/stories/ModelLineage.tsx
@@ -1,5 +1,5 @@
 import { debounce } from 'lodash'
-import { Focus, Rows2, Rows3 } from 'lucide-react'
+import { Focus, LockOpen, Rows2, Rows3, Lock } from 'lucide-react'
 import React from 'react'
 
 import { type ColumnLevelLineageAdjacencyList } from '../LineageColumnLevel/ColumnLevelLineageContext'
@@ -74,6 +74,7 @@ export const ModelLineage = ({
 }) => {
   const [zoom, setZoom] = React.useState(ZOOM_THRESHOLD)
   const [isBuildingLayout, setIsBuildingLayout] = React.useState(false)
+  const [nodesDraggable, setNodesDraggable] = React.useState(false)
   const [edges, setEdges] = React.useState<
     LineageEdge<EdgeData, ModelNodeId, ModelEdgeId, ModelColumnID>[]
   >([])
@@ -388,6 +389,7 @@ export const ModelLineage = ({
         nodeTypes={nodeTypes}
         edgeTypes={edgeTypes}
         className={className}
+        nodesDraggable={nodesDraggable}
         controls={
           <>
             <LineageControlButton
@@ -408,6 +410,13 @@ export const ModelLineage = ({
             >
               <LineageControlIcon Icon={Focus} />
             </LineageControlButton>
+            <LineageControlButton
+              text={nodesDraggable ? 'Lock nodes' : 'Unlock nodes'}
+              onClick={() => setNodesDraggable(prev => !prev)}
+              disabled={isBuildingLayout}
+            >
+              <LineageControlIcon Icon={nodesDraggable ? Lock : LockOpen} />
+            </LineageControlButton>
           </>
         }
       />

From 3a710e251830b8dca9386b3f903d07db58a7deac Mon Sep 17 00:00:00 2001
From: Max Mykal <mykalmax@yahoo.com>
Date: Thu, 2 Oct 2025 11:56:16 -0700
Subject: [PATCH 2/2] clean up

---
 web/common/src/components/Lineage/LineageLayout.tsx | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/web/common/src/components/Lineage/LineageLayout.tsx b/web/common/src/components/Lineage/LineageLayout.tsx
index 42800dfefa..4b2e06b0b3 100644
--- a/web/common/src/components/Lineage/LineageLayout.tsx
+++ b/web/common/src/components/Lineage/LineageLayout.tsx
@@ -14,8 +14,6 @@ import {
   getOutgoers,
   useReactFlow,
   useViewport,
-  useEdgesState,
-  useNodesState,
   applyNodeChanges,
   applyEdgeChanges,
 } from '@xyflow/react'
@@ -340,8 +338,6 @@ function LineageLayoutBase<
   React.useEffect(() => {
     if (currentNode?.id) {
       setSelectedNodeId(currentNode.id)
-    } else if (selectedNodeId) {
-      // setSelectedNodeId(selectedNodeId);
     } else {
       const node = nodes.length > 0 ? nodes[nodes.length - 1] : null