fix(vscode): allow undoing clicking of lineage columns (#4876)

Author: benfdking

vscode/react/src/components/graph/ModelColumns.tsx

@@ -28,7 +28,6 @@ import {
 import clsx from 'clsx'
 import {
   type ColumnDescription,
-  type Column,
   type ColumnLineageApiLineageModelNameColumnNameGet200,
   type LineageColumn,
   type LineageColumnSource,
@@ -47,6 +46,7 @@ import { useApiColumnLineage } from '@/api/index'
 import SourceList from '@/components/sourceList/SourceList'
 import type { Lineage } from '@/domain/lineage'
 import type { ModelName } from '@/domain/models'
+import type { Column } from '@/domain/column'
 
 export default function ModelColumns({
   nodeId,
@@ -207,7 +207,7 @@ export default function ModelColumns({
               id={toID(nodeId, column.name)}
               nodeId={nodeId}
               column={column}
-              disabled={true}
+              disabled={disabled}
               updateColumnLineage={updateColumnLineage}
               removeEdges={removeEdges}
               isActive={true}

vscode/react/src/components/graph/ModelNode.tsx

@@ -5,9 +5,9 @@ import { ModelType, type Model } from '@/api/client'
 import { useLineageFlow } from './context'
 import { type GraphNodeData } from './help'
 import { Position, type NodeProps } from 'reactflow'
-import { type Column } from '@/api/client'
 import ModelNodeHeaderHandles from './ModelNodeHeaderHandles'
 import ModelColumns from './ModelColumns'
+import { fromAPIColumn, type Column } from '@/domain/column'
 
 export const EnumLineageNodeModelType = {
   ...ModelType,
@@ -53,7 +53,7 @@ export default function ModelNode({
     const modelsArray = Object.values(models)
     const decodedId = decodeURIComponent(id)
     const model = modelsArray.find((m: Model) => m.fqn === decodedId)
-    const modelColumns = model?.columns ?? []
+    const modelColumns = model?.columns?.map(fromAPIColumn) ?? []
 
     Object.keys(lineage[decodedId]?.columns ?? {}).forEach((column: string) => {
       const found = modelColumns.find(({ name }: any) => {
@@ -65,7 +65,9 @@ export default function ModelNode({
       })
 
       if (isNil(found)) {
-        modelColumns.push({ name: column, type: EnumColumnType.UNKNOWN })
+        modelColumns.push(
+          fromAPIColumn({ name: column, type: EnumColumnType.UNKNOWN }),
+        )
       }
     })
 

vscode/react/src/components/graph/context.tsx

@@ -1,4 +1,4 @@
-import { type Column, type Model } from '@/api/client'
+import { type Model } from '@/api/client'
 import {
   createContext,
   useState,
@@ -11,6 +11,7 @@ import { EnumSide } from './types'
 import { type Node } from 'reactflow'
 import type { Lineage } from '@/domain/lineage'
 import type { ModelSQLMeshModel } from '@/domain/sqlmesh-model'
+import type { Column } from '@/domain/column'
 
 export interface Connections {
   left: string[]

vscode/react/src/domain/column.ts

@@ -0,0 +1,18 @@
+import { type Column as APIColumn } from '@/api/client'
+import { type Branded } from '@bus/brand'
+
+export type ColumnName = Branded<string, 'ColumnName'>
+
+export type Column = {
+  name: ColumnName
+  type: string
+  description?: string
+}
+
+export function fromAPIColumn(column: APIColumn): Column {
+  return {
+    name: column.name as ColumnName,
+    type: column.type,
+    description: column.description ?? undefined,
+  }
+}