some pr comments

Author: mykalmax

pnpm-lock.yaml

@@ -17,3 +17,26 @@ type Brand<B> = { [__brand]: B }
  * userId == userName -> compile error
  */
 export type Branded<T, B> = T & Brand<B>
+
+/**
+ * Constraint that only accepts branded string types
+ */
+export type BrandedString = string & Brand<string>
+
+/**
+ * BrandedRecord is a type that creates a branded Record type with strict key checking.
+ * This ensures that Record<BrandedKey1, V> is NOT assignable to Record<BrandedKey2, V>
+ *
+ * @example
+ * type ModelFQN = Branded<string, 'ModelFQN'>
+ * type ModelName = Branded<string, 'ModelName'>
+ *
+ * type FQNMap = BrandedRecord<ModelFQN, string>
+ * type NameMap = BrandedRecord<ModelName, string>
+ *
+ * const fqnMap: FQNMap = {}
+ * const nameMap: NameMap = fqnMap // TypeScript error!
+ */
+export type BrandedRecord<K extends BrandedString, V> = Record<K, V> & {
+  readonly __recordKeyBrand: K
+}

vscode/bus/src/brand.ts

@@ -19,16 +19,15 @@
     "@heroicons/react": "^2.2.0",
     "@radix-ui/react-select": "^2.2.5",
     "@tanstack/react-query": "^5.83.0",
-    "@tanstack/react-router": "^1.129.8",
+    "@tanstack/react-router": "^1.131.26",
     "@tanstack/react-router-devtools": "^1.131.26",
     "@tanstack/react-virtual": "^3.13.12",
-    "@tanstack/router-plugin": "^1.129.8",
+    "@tanstack/router-plugin": "^1.131.26",
     "@tobikodata/sqlmesh-common": "file:../../web/common",
     "apache-arrow": "^19.0.1",
     "clsx": "^2.1.1",
     "cronstrue": "3.3.0",
     "elkjs": "^0.8.2",
-    "lodash": "4.17.21",
     "lucide-react": "0.542.0",
     "orval": "^7.10.0",
     "react": "^18.3.1",
@@ -47,12 +46,11 @@
     "@storybook/react-vite": "^9.0.18",
     "@testing-library/dom": "^10.4.1",
     "@testing-library/react": "^16.3.0",
-    "@types/lodash": "4.17.20",
     "@types/react": "^18.3.23",
     "@types/react-dom": "^18.3.7",
     "@vitejs/plugin-react": "^4.7.0",
-    "@vitest/browser": "3.2.3",
-    "@vitest/coverage-v8": "3.2.3",
+    "@vitest/browser": "3.2.4",
+    "@vitest/coverage-v8": "3.2.4",
     "jsdom": "^26.1.0",
     "playwright": "^1.54.1",
     "storybook": "^9.0.18",

vscode/react/package.json

@@ -1,12 +1,9 @@
 import React from 'react'
 
-import { debounce } from 'lodash'
 import { Focus, Rows2, Rows3 } from 'lucide-react'
 import {
   FactoryEdgeWithGradient,
   EdgeWithGradient,
-  type LineageAdjacencyList,
-  type LineageDetails,
   ZOOM_THRESHOLD,
   type LineageEdge,
   type LineageNodesMap,
@@ -34,20 +31,21 @@ import {
 } from '@tobikodata/sqlmesh-common/lineage'
 
 import {
+  type BrandedLineageAdjacencyList,
+  type BrandedLineageDetails,
   type ColumnName,
   type EdgeData,
   type ModelColumnID,
   type ModelEdgeId,
   type ModelLineageNodeDetails,
   type ModelNodeId,
   type NodeData,
-  type NodeType,
   ModelLineageContext,
 } from './ModelLineageContext'
 import { ModelNode } from './ModelNode'
 import { useModelLineage } from './ModelLineageContext'
-import type { ModelName } from '@/domain/models'
-import { getNodeTypeColorVar } from './help'
+import type { ModelFQN } from '@/domain/models'
+import { NODE_TYPE_COLOR_VAR } from './help'
 
 const nodeTypes = {
   node: ModelNode,
@@ -64,9 +62,9 @@ export const ModelLineage = ({
   className,
   onNodeClick,
 }: {
-  adjacencyList: LineageAdjacencyList<ModelName>
-  lineageDetails: LineageDetails<ModelName, ModelLineageNodeDetails>
-  selectedModelName?: ModelName
+  adjacencyList: BrandedLineageAdjacencyList<ModelFQN>
+  lineageDetails: BrandedLineageDetails<ModelFQN, ModelLineageNodeDetails>
+  selectedModelName?: ModelFQN
   className?: string
   onNodeClick?: (
     event: React.MouseEvent<Element, MouseEvent>,
@@ -98,7 +96,7 @@ export const ModelLineage = ({
 
   const [showColumns, setShowColumns] = React.useState(false)
   const [columnLevelLineage, setColumnLevelLineage] = React.useState<
-    Map<ModelColumnID, ColumnLevelLineageAdjacencyList<ModelName, ColumnName>>
+    Map<ModelColumnID, ColumnLevelLineageAdjacencyList<ModelFQN, ColumnName>>
   >(new Map())
   const [fetchingColumns, setFetchingColumns] = React.useState<
     Set<ModelColumnID>
@@ -108,17 +106,17 @@ export const ModelLineage = ({
     adjacencyListColumnLevel,
     selectedColumns,
     adjacencyListKeysColumnLevel,
-  } = useColumnLevelLineage<ModelName, ColumnName, ModelColumnID>(
+  } = useColumnLevelLineage<ModelFQN, ColumnName, ModelColumnID>(
     columnLevelLineage,
   )
 
   const adjacencyListKeys = React.useMemo(() => {
-    let keys: ModelName[] = []
+    let keys: ModelFQN[] = []
 
     if (adjacencyListKeysColumnLevel.length > 0) {
       keys = adjacencyListKeysColumnLevel
     } else {
-      keys = Object.keys(adjacencyList) as ModelName[]
+      keys = Object.keys(adjacencyList) as ModelFQN[]
     }
 
     return keys
@@ -131,7 +129,7 @@ export const ModelLineage = ({
       const node = createNode('node', nodeId, {
         name: detail.name,
         displayName: detail.display_name,
-        model_type: detail.model_type as NodeType,
+        model_type: detail.model_type,
         identifier: detail.identifier,
         kind: detail.kind,
         cron: detail.cron,
@@ -177,7 +175,7 @@ export const ModelLineage = ({
 
   const transformedNodesMap = React.useMemo(() => {
     return getTransformedNodes<
-      ModelName,
+      ModelFQN,
       ModelLineageNodeDetails,
       NodeData,
       ModelNodeId
@@ -201,19 +199,15 @@ export const ModelLineage = ({
         data.startColor = 'var(--color-lineage-node-port-edge-source)'
       } else {
         if (sourceNode?.data?.model_type) {
-          data.startColor = getNodeTypeColorVar(
-            sourceNode.data.model_type as NodeType,
-          )
+          data.startColor = NODE_TYPE_COLOR_VAR[sourceNode.data.model_type]
         }
       }
 
       if (targetHandleId) {
         data.endColor = 'var(--color-lineage-node-port-edge-target)'
       } else {
         if (targetNode?.data?.model_type) {
-          data.endColor = getNodeTypeColorVar(
-            targetNode.data.model_type as NodeType,
-          )
+          data.endColor = NODE_TYPE_COLOR_VAR[targetNode.data.model_type]
         }
       }
 
@@ -237,7 +231,7 @@ export const ModelLineage = ({
   const edgesColumnLevel = React.useMemo(
     () =>
       getEdgesFromColumnLineage<
-        ModelName,
+        ModelFQN,
         ColumnName,
         EdgeData,
         ModelEdgeId,
@@ -254,34 +248,32 @@ export const ModelLineage = ({
     return edgesColumnLevel.length > 0
       ? edgesColumnLevel
       : getTransformedModelEdgesTargetSources<
-          ModelName,
+          ModelFQN,
           EdgeData,
           ModelNodeId,
           ModelEdgeId,
           ModelColumnID
         >(adjacencyListKeys, adjacencyList, transformEdge)
   }, [adjacencyListKeys, adjacencyList, transformEdge, edgesColumnLevel])
 
-  const calculateLayout = React.useMemo(() => {
-    return debounce(
-      (
-        eds: LineageEdge<EdgeData, ModelNodeId, ModelEdgeId, ModelColumnID>[],
-        nds: LineageNodesMap<NodeData>,
-      ) => {
-        const { edges, nodesMap } = buildLayout<
-          NodeData,
-          EdgeData,
-          ModelNodeId,
-          ModelEdgeId,
-          ModelColumnID
-        >({ edges: eds, nodesMap: nds })
-        setEdges(edges)
-        setNodesMap(nodesMap)
-        setIsBuildingLayout(false)
-      },
-      200,
-    )
-  }, [])
+  const calculateLayout = React.useCallback(
+    (
+      eds: LineageEdge<EdgeData, ModelNodeId, ModelEdgeId, ModelColumnID>[],
+      nds: LineageNodesMap<NodeData>,
+    ) => {
+      const { edges, nodesMap } = buildLayout<
+        NodeData,
+        EdgeData,
+        ModelNodeId,
+        ModelEdgeId,
+        ModelColumnID
+      >({ edges: eds, nodesMap: nds })
+      setEdges(edges)
+      setNodesMap(nodesMap)
+      setIsBuildingLayout(false)
+    },
+    [edges, nodesMap, setEdges, setNodesMap, setIsBuildingLayout],
+  )
 
   const nodes = React.useMemo(() => {
     return Object.values(nodesMap)
@@ -317,27 +309,25 @@ export const ModelLineage = ({
     } else {
       calculateLayout(transformedEdges, transformedNodesMap)
     }
-  }, [
-    calculateLayout,
-    showOnlySelectedNodes,
-    transformedEdges,
-    transformedNodesMap,
-  ])
+  }, [showOnlySelectedNodes, transformedEdges, transformedNodesMap])
 
+  // currentNodeId is passed from the parent component
+  // we it change we need to reset the selectedNodeId
   React.useEffect(() => {
     if (currentNodeId) {
       setSelectedNodeId(currentNodeId)
     }
   }, [currentNodeId])
 
+  // When the selectedColumns is empty it measn we dont have any selected columns
+  // so we need to set the selectedNodeId back to the currentNode.id
+  // where currentNode derived from currentNodeId if present in nodesMap
+  // if the currentNode is null we need to set the selectedNodeId to null
   React.useEffect(() => {
-    if (
-      (selectedColumns.size === 0 && currentNode?.id !== selectedNodeId) ||
-      selectedNodeId == null
-    ) {
+    if (selectedColumns.size === 0 && selectedNodeId != currentNode?.id) {
       setSelectedNodeId(currentNode?.id || null)
     }
-  }, [selectedColumns, currentNode?.id, selectedNodeId])
+  }, [selectedColumns, currentNode?.id])
 
   function toggleColumns() {
     setShowColumns(prev => !prev)

vscode/react/src/pages/ModelLineage.tsx

@@ -1,31 +1,52 @@
-import type { ModelName } from '@/domain/models'
-import type { Branded } from '@tobikodata/sqlmesh-common'
+import type { ModelType } from '@/api/client'
+import type { ModelFQN, ModelName } from '@/domain/models'
+import type { Branded, BrandedString } from '@bus/brand'
 import {
   type Column,
-  type ColumnLevelLineageAdjacencyList,
   type ColumnLevelLineageContextValue,
   type LineageContextValue,
   type PathType,
   getInitial as getLineageContextInitial,
   getColumnLevelLineageContextInitial,
   createLineageContext,
+  type LineageAdjacencyList,
+  type LineageDetails,
+  type ColumnLevelLineageAdjacencyList,
 } from '@tobikodata/sqlmesh-common/lineage'
 
+export type BrandedLineageAdjacencyList<K extends BrandedString> =
+  LineageAdjacencyList<K> & {
+    readonly __adjacencyListKeyBrand: K
+  }
+
+export type BrandedLineageDetails<K extends BrandedString, V> = LineageDetails<
+  K,
+  V
+> & {
+  readonly __lineageDetailsKeyBrand: K
+}
+
+export type BrandedColumnLevelLineageAdjacencyList<
+  K extends BrandedString,
+  V extends BrandedString,
+> = ColumnLevelLineageAdjacencyList<K, V> & {
+  readonly __columnLevelLineageAdjacencyListKeyBrand: K
+  readonly __columnLevelLineageAdjacencyListColumnKeyBrand: V
+}
+
 export type ColumnName = Branded<string, 'ColumnName'>
 export type ModelColumnID = Branded<string, 'ModelColumnID'>
 export type ModelNodeId = Branded<string, 'ModelNodeId'>
 export type ModelEdgeId = Branded<string, 'ModelEdgeId'>
 export type ModelColumn = Column & {
   id: ModelColumnID
   name: ColumnName
-  columnLineageData?: ColumnLevelLineageAdjacencyList<ModelName, ColumnName>
 }
 
-export type NodeType = 'sql' | 'python'
 export type ModelLineageNodeDetails = {
-  name: ModelName
-  display_name: string
-  model_type: string
+  name: ModelFQN
+  display_name: ModelName
+  model_type: ModelType
   identifier?: string | null
   version?: string | null
   dialect?: string | null
@@ -37,9 +58,9 @@ export type ModelLineageNodeDetails = {
 }
 
 export type NodeData = {
-  name: ModelName
-  displayName: string
-  model_type: NodeType
+  name: ModelFQN
+  displayName: ModelName
+  model_type: ModelType
   identifier?: string | null
   version?: string | null
   kind?: string | null
@@ -58,7 +79,7 @@ export type EdgeData = {
 }
 
 export type ModelLineageContextValue = ColumnLevelLineageContextValue<
-  ModelName,
+  ModelFQN,
   ColumnName,
   ModelColumnID
 > &
@@ -72,11 +93,7 @@ export type ModelLineageContextValue = ColumnLevelLineageContextValue<
 
 export const initial = {
   ...getLineageContextInitial<ModelNodeId, ModelEdgeId>(),
-  ...getColumnLevelLineageContextInitial<
-    ModelName,
-    ColumnName,
-    ModelColumnID
-  >(),
+  ...getColumnLevelLineageContextInitial<ModelFQN, ColumnName, ModelColumnID>(),
 }
 
 export const { Provider, useLineage } = createLineageContext<