SQLMesh · benfdking · Jul 4, 2025 · Jul 4, 2025
diff --git a/vscode/react/src/components/graph/ModelColumns.tsx b/vscode/react/src/components/graph/ModelColumns.tsx
@@ -10,7 +10,7 @@ import {
   isNotNil,
   truncate,
 } from '@/utils/index'
-import { EnumSide, toID, type Side } from './types'
+import { toID, type PartialColumnHandleId, type Side } from './types'
 import { NoSymbolIcon } from '@heroicons/react/24/solid'
 import { ClockIcon, ExclamationCircleIcon } from '@heroicons/react/24/outline'
 import clsx from 'clsx'
@@ -38,7 +38,7 @@ export function ModelColumns({
   columns,
   disabled,
   className,
-  limit = 5,
+  limit,
   withHandles = false,
   withDescription = true,
   maxHeight = '50vh',
@@ -47,7 +47,7 @@ export function ModelColumns({
   columns: Column[]
   disabled?: boolean
   className?: string
-  limit?: number
+  limit: number
   withHandles?: boolean
   withDescription?: boolean
   maxHeight?: string
@@ -125,7 +125,7 @@ export function ModelColumns({
   }
 
   const isSelectManually = useCallback(
-    function isSelectManually(columnName: string): boolean {
+    function isSelectManually(columnName: ColumnName): boolean {
       if (isNil(manuallySelectedColumn)) return false
 
       const [selectedModel, selectedColumn] = manuallySelectedColumn
@@ -138,12 +138,10 @@ export function ModelColumns({
   )
 
   const removeEdges = useCallback(
-    function removeEdges(columnId: string): void {
+    function removeEdges(columnId: PartialColumnHandleId): void {
       const visited = new Set<string>()
 
-      removeActiveEdges(
-        walk(columnId, EnumSide.Left).concat(walk(columnId, EnumSide.Right)),
-      )
+      removeActiveEdges(walk(columnId, 'left').concat(walk(columnId, 'right')))
 
       if (connections.size === 0 && isNotNil(lineageCache)) {
         setLineage(lineageCache)
@@ -164,12 +162,12 @@ export function ModelColumns({
         return edges
           .map(edge =>
             [
-              side === EnumSide.Left
-                ? [toID(EnumSide.Left, id), toID(EnumSide.Right, edge)]
-                : [toID(EnumSide.Left, edge), toID(EnumSide.Right, id)],
+              side === 'left'
+                ? [toID('left', id), toID('right', edge)]
+                : [toID('left', edge), toID('right', id)],
             ].concat(walk(edge, side)),
           )
-          .flat() as Array<[string, string]>
+          .flat() as Array<[PartialColumnHandleId, PartialColumnHandleId]>
       }
     },
     [removeActiveEdges, connections],
@@ -324,8 +322,8 @@ function ModelColumn({
   withHandles = false,
   withDescription = true,
 }: {
-  id: string
-  nodeId: string
+  id: PartialColumnHandleId
+  nodeId: ModelEncodedFQN
   column: Column
   disabled?: boolean
   isActive?: boolean
@@ -337,7 +335,7 @@ function ModelColumn({
   updateColumnLineage: (
     lineage: ColumnLineageApiLineageModelNameColumnNameGet200,
   ) => void
-  removeEdges: (columnId: string) => void
+  removeEdges: (columnId: PartialColumnHandleId) => void
   selectManually?: React.Dispatch<
     React.SetStateAction<
       [ModelSQLMeshModel<InitialSQLMeshModel>, Column] | undefined
@@ -454,8 +452,8 @@ function ColumnHandles({
   children,
   className,
 }: {
-  nodeId: string
-  id: string
+  nodeId: ModelEncodedFQN
+  id: PartialColumnHandleId
   children: React.ReactNode
   className?: string
   hasLeft?: boolean
@@ -482,7 +480,7 @@ function ColumnHandles({
       {hasLeft && (
         <Handle
           type="target"
-          id={toID(EnumSide.Left, id)}
+          id={toID('left', id)}
           position={Position.Left}
           isConnectable={false}
           className="w-2 h-2 rounded-full"
@@ -492,7 +490,7 @@ function ColumnHandles({
       {hasRight && (
         <Handle
           type="source"
-          id={toID(EnumSide.Right, id)}
+          id={toID('right', id)}
           position={Position.Right}
           isConnectable={false}
           className="w-2 h-2 rounded-full"
@@ -510,19 +508,13 @@ function ColumnDisplay({
   disabled = false,
   withDescription = true,
 }: {
-  columnName: string
+  columnName: ColumnName
   columnType: string
   columnDescription?: ColumnDescription
   disabled?: boolean
   withDescription?: boolean
   className?: string
 }): JSX.Element {
-  let decodedColumnName = columnName
-
-  try {
-    decodedColumnName = decodeURI(columnName)
-  } catch {}
-
   return (
     <div
       className={clsx(
@@ -533,7 +525,7 @@ function ColumnDisplay({
     >
       <div className="w-full flex justify-between items-center">
         <span
-          title={decodedColumnName}
+          title={columnName}
           className={clsx('flex items-center', disabled && 'opacity-50')}
         >
           {disabled && (
@@ -542,7 +534,7 @@ function ColumnDisplay({
               className="w-3 h-3 mr-2"
             />
           )}
-          {truncate(decodedColumnName, 50, 20)}
+          {truncate(columnName, 50, 20)}
         </span>
         <span
           title={columnType}

diff --git a/vscode/react/src/components/graph/ModelNode.tsx b/vscode/react/src/components/graph/ModelNode.tsx
@@ -9,6 +9,8 @@ import { ModelNodeHeaderHandles } from './ModelNodeHeaderHandles'
 import { ModelColumns } from './ModelColumns'
 import { fromAPIColumn, type Column } from '@/domain/column'
 import { decode, type ModelEncodedFQN } from '@/domain/models'
+import { toKeys } from './types'
+import { MAX_VISIBLE_COLUMNS } from './constants'
 
 export const EnumLineageNodeModelType = {
   ...ModelType,
@@ -57,33 +59,24 @@ export default function ModelNode({
     const model = modelsArray.find((m: Model) => m.fqn === decodedId)
     const modelColumns = model?.columns?.map(fromAPIColumn) ?? []
 
-    Object.keys(lineage[decodedId]?.columns ?? {}).forEach((column: string) => {
-      const found = modelColumns.find(({ name }: any) => {
-        try {
-          return name === decodeURI(column)
-        } catch {
-          return name === column
-        }
-      })
-
+    toKeys(lineage[decodedId]?.columns ?? {}).forEach(column => {
+      const found = modelColumns.find(({ name }) => name === column)
       if (isNil(found)) {
         modelColumns.push(
           fromAPIColumn({ name: column, type: EnumColumnType.UNKNOWN }),
         )
       }
     })
-
-    modelColumns.forEach((column: any) => {
+    return modelColumns.map(column => {
       let columnType = column.type ?? EnumColumnType.UNKNOWN
-
       if (columnType.startsWith(EnumColumnType.STRUCT)) {
         columnType = EnumColumnType.STRUCT
       }
-
-      column.type = columnType
+      return {
+        ...column,
+        type: columnType,
+      }
     })
-
-    return modelColumns
   }, [id, models, lineage])
 
   const highlightedNodeModels = useMemo(
@@ -136,7 +129,7 @@ export default function ModelNode({
   // Ensure nodeData.type is a valid LineageNodeModelType
   const nodeType: LineageNodeModelType = Object.values(
     EnumLineageNodeModelType,
-  ).includes(nodeData.type as any)
+  ).includes(nodeData.type)
     ? (nodeData.type as LineageNodeModelType)
     : EnumLineageNodeModelType.unknown
 
@@ -222,12 +215,13 @@ export default function ModelNode({
             ? undefined
             : handleSelect
         }
-        count={hasHighlightedNodes ? undefined : columns.length}
+        numberOfColumns={columns.length}
       />
       {showColumns && (
         <ModelColumns
           className="nowheel rounded-b-lg bg-theme-lighter text-xs"
           nodeId={id}
+          limit={MAX_VISIBLE_COLUMNS}
           columns={columns}
           disabled={shouldDisableColumns}
           withHandles={true}

diff --git a/vscode/react/src/components/graph/ModelNodeHeaderHandles.tsx b/vscode/react/src/components/graph/ModelNodeHeaderHandles.tsx
@@ -3,10 +3,11 @@ import { Handle, Position } from 'reactflow'
 import 'reactflow/dist/base.css'
 import { getModelNodeTypeTitle } from './help'
 import { isNotNil, truncate } from '@/utils/index'
-import { EnumSide, toID } from './types'
+import { toID } from './types'
 import { ArrowRightCircleIcon } from '@heroicons/react/24/solid'
 import clsx from 'clsx'
 import { type LineageNodeModelType } from './ModelNode'
+import type { ModelEncodedFQN } from '@/domain/models'
 
 export function ModelNodeHeaderHandles({
   id,
@@ -17,16 +18,16 @@ export function ModelNodeHeaderHandles({
   isDraggable = false,
   label,
   type,
-  count,
+  numberOfColumns,
   handleClick,
   handleSelect,
 }: {
-  id: string
+  id: ModelEncodedFQN
   label: string
   type?: LineageNodeModelType
   hasLeft?: boolean
   hasRight?: boolean
-  count?: number
+  numberOfColumns?: number
   className?: string
   isSelected?: boolean
   isDraggable?: boolean
@@ -38,7 +39,7 @@ export function ModelNodeHeaderHandles({
       {hasLeft && (
         <Handle
           type="target"
-          id={toID(EnumSide.Left, id)}
+          id={toID('left', id)}
           position={Position.Left}
           isConnectable={false}
           className="-ml-2 border rounded-full overflow-hidden border-current"
@@ -87,17 +88,17 @@ export function ModelNodeHeaderHandles({
           >
             {truncate(decodeURI(label), 50, 20)}
           </span>
-          {isNotNil(count) && (
+          {isNotNil(numberOfColumns) && (
             <span className="flex justify-between ml-2 mr-1 px-2 rounded-full bg-neutral-10">
-              {count}
+              {numberOfColumns}
             </span>
           )}
         </span>
       </div>
       {hasRight && (
         <Handle
           type="source"
-          id={toID(EnumSide.Right, id)}
+          id={toID('right', id)}
           position={Position.Right}
           isConnectable={false}
           className="-mr-2 border rounded-full overflow-hidden border-current"

diff --git a/vscode/react/src/components/graph/constants.ts b/vscode/react/src/components/graph/constants.ts
@@ -0,0 +1,16 @@
+/**
+ * Space between nodes.
+ */
+export const NODE_BALANCE_SPACE = 64
+/**
+ * Height of a column line.
+ */
+export const COLUMN_LINE_HEIGHT = 24
+/**
+ * Assumed width of a character.
+ */
+export const CHAR_WIDTH = 8
+/**
+ * Maximum number of columns that can be visible in a node.
+ */
+export const MAX_VISIBLE_COLUMNS = 5