feat(vscode): add basic column lineage (#4476)

Author: benfdking

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

@@ -28,8 +28,12 @@ export default function ModelNode({
   data,
   sourcePosition,
   targetPosition,
-}: NodeProps): JSX.Element {
-  const nodeData: GraphNodeData = data ?? {}
+}: NodeProps<GraphNodeData>): JSX.Element {
+  const nodeData: GraphNodeData = data ?? {
+    label: '',
+    type: EnumLineageNodeModelType.unknown,
+    withColumns: false,
+  }
   const {
     // connections,
     models,
@@ -128,11 +132,18 @@ export default function ModelNode({
   const isMainNode = mainNode === id
   const isHighlightedNode = highlightedNodeModels.includes(id)
   const isSelected = selectedNodes.has(id)
-  const isModelSQL = nodeData.type === EnumLineageNodeModelType.sql
-  const isCTE = nodeData.type === EnumLineageNodeModelType.cte
-  const isModelExternal = nodeData.type === EnumLineageNodeModelType.external
-  const isModelSeed = nodeData.type === EnumLineageNodeModelType.seed
-  const isModelUnknown = nodeData.type === EnumLineageNodeModelType.unknown
+  // Ensure nodeData.type is a valid LineageNodeModelType
+  const nodeType: LineageNodeModelType = Object.values(
+    EnumLineageNodeModelType,
+  ).includes(nodeData.type as any)
+    ? (nodeData.type as LineageNodeModelType)
+    : EnumLineageNodeModelType.unknown
+
+  const isModelSQL = nodeType === EnumLineageNodeModelType.sql
+  const isCTE = nodeType === EnumLineageNodeModelType.cte
+  const isModelExternal = nodeType === EnumLineageNodeModelType.external
+  const isModelSeed = nodeType === EnumLineageNodeModelType.seed
+  const isModelUnknown = nodeType === EnumLineageNodeModelType.unknown
   const showColumns = isArrayNotEmpty(columns) && isFalse(hasHighlightedNodes)
   const isActiveNode =
     selectedNodes.size > 0 || activeNodes.size > 0 || withConnected
@@ -188,7 +199,7 @@ export default function ModelNode({
     >
       <ModelNodeHeaderHandles
         id={id}
-        // type={nodeData.type}
+        type={nodeType}
         label={nodeData.label}
         isSelected={isSelected}
         isDraggable={true}

vscode/react/src/components/graph/help.ts

@@ -22,6 +22,8 @@ export interface GraphNodeData {
   label: string
   type: LineageNodeModelType
   withColumns: boolean
+  width?: number
+  height?: number
   [key: string]: any
 }
 
@@ -170,20 +172,25 @@ function getNodeMap({
   const modelNames = Object.keys(lineage)
 
   return modelNames.reduce((acc: Record<string, Node>, modelName: string) => {
-    const model = models[modelName]
+    const decodedModelName = modelName.includes('%')
+      ? decodeURI(modelName)
+      : modelName
+    const model = Object.values(models).find(m => m.fqn === decodedModelName)
+    console.log('model', model)
+    const nodeType: LineageNodeModelType = isNotNil(model)
+      ? (model.type as LineageNodeModelType)
+      : // If model name present in lineage but not in global models
+        // it means either this is a CTE or model is UNKNOWN
+        // CTEs only have connections between columns
+        // where UNKNOWN model has connection only from another model
+        unknownModels.has(modelName)
+        ? EnumLineageNodeModelType.unknown
+        : EnumLineageNodeModelType.cte
+
     const node = createGraphNode(modelName, {
-      // @ts-ignore
-      label: model?.displayName ?? modelName,
+      label: model?.name ?? modelName,
       withColumns,
-      type: isNotNil(model)
-        ? (model.type as LineageNodeModelType)
-        : // If model name present in lineage but not in global models
-          // it means either this is a CTE or model is UNKNOWN
-          // CTEs only have connections between columns
-          // where UNKNOWN model has connection only from another model
-          unknownModels.has(modelName)
-          ? EnumLineageNodeModelType.unknown
-          : EnumLineageNodeModelType.cte,
+      type: nodeType,
     })
     const columnsCount = withColumns
       ? (models[modelName]?.columns?.length ?? 0)
@@ -267,7 +274,7 @@ function createGraphNode(
   data: GraphNodeData,
   position: XYPosition = { x: 0, y: 0 },
   hidden: boolean = false,
-): Node {
+): Node<GraphNodeData> {
   return {
     id,
     dragHandle: '.drag-handle',

vscode/react/src/domain/sqlmesh-model.ts

@@ -10,8 +10,7 @@ import {
 } from '@/api/client'
 import { isArrayNotEmpty } from '@/utils/index'
 import { ModelInitial } from './initial'
-
-type Lineage = any
+import type { Lineage } from './lineage'
 
 export interface InitialSQLMeshModel extends Model {
   lineage?: Record<string, Lineage>