progress: showing the lineage graph

[ci skip]

Author: benfdking

package-lock.json

@@ -60,7 +60,7 @@
     "@types/fs-extra": "^11.0.4",
     "@vscode/python-extension": "^1.0.5",
     "fs-extra": "^11.3.0",
-    "vscode-languageclient": "^9.0.1",
+    "vscode-languageclient": "^9.0.1"
   },
   "devDependencies": {
     "@types/mocha": "^10.0.10",

vscode/extension/package.json

@@ -1,5 +1,6 @@
 import { useApiModelLineage, useApiModels } from '@/api/index'
 import { useEffect, useMemo, useRef, useState } from 'react'
+import { type ModelSQLMeshModel } from '@/domain/sqlmesh-model'
 import { type HighlightedNodes, useLineageFlow } from './context'
 import { ChevronDownIcon } from '@heroicons/react/24/solid'
 import ReactFlow, {
@@ -18,6 +19,7 @@ import ReactFlow, {
 } from 'reactflow'
 import Loading from '@/components/loading/Loading'
 import Spinner from '@/components/logo/Spinner'
+import { createLineageWorker } from '@/workers'
 import { isArrayEmpty, isFalse, isNil, isNotNil } from '@/utils/index'
 import ListboxShow from '@/components/listbox/ListboxShow'
 import clsx from 'clsx'
@@ -42,7 +44,7 @@ export function ModelLineage({
   model,
   highlightedNodes,
 }: {
-  model: Model 
+  model: ModelSQLMeshModel
   highlightedNodes?: HighlightedNodes
 }): JSX.Element {
   const {
@@ -67,24 +69,32 @@ export function ModelLineage({
     isFetching: isFetchingModelLineage,
   } = useApiModelLineage(model.name)
   const { isFetching: isFetchingModels } = useApiModels()
-  const [isMergingModels, setIsMergingModels] = useState(false)
+
+  const [isMegringModels, setIsMergingModels] = useState(false)
   const [modelLineage, setModelLineage] =
     useState<ModelLineageApiLineageModelNameGet200 | undefined>(undefined)
 
-  console.log('model to lineage', model)
-  console.log('modelLineage', modelLineage)
-  console.log('models', models)
-  console.log('isFetchingModelLineage', isFetchingModelLineage)
+  useEffect(() => {
+    const lineageWorker = createLineageWorker()
 
+    lineageWorker.addEventListener('message', handleLineageWorkerMessage)
 
-  useEffect(() => {
     getModelLineage()
       .then(({ data }) => {
         setModelLineage(data)
 
         if (isNil(data)) return
 
         setIsMergingModels(true)
+
+        lineageWorker.postMessage({
+          topic: 'lineage',
+          payload: {
+            currentLineage: {},
+            newLineage: data,
+            mainNode: model.fqn,
+          },
+        })
       })
       .catch(error => {
         handleError?.(error)
@@ -101,6 +111,9 @@ export function ModelLineage({
     return () => {
       // void cancel?.()
 
+      lineageWorker.removeEventListener('message', handleLineageWorkerMessage)
+      lineageWorker.terminate()
+
       setLineage({})
       setNodeConnections({})
       setMainNode(undefined)
@@ -113,10 +126,10 @@ export function ModelLineage({
       modelName = encodeURI(modelName)
 
       if (
-        !models[modelName] &&
-        !unknownModels[modelName]
+        isFalse(modelName in models) &&
+        isFalse(modelName in unknownModels)
       ) {
-        unknownModels[modelName]
+        unknownModels.add(modelName)
       }
     })
 
@@ -127,13 +140,27 @@ export function ModelLineage({
     setHighlightedNodes(highlightedNodes ?? {})
   }, [highlightedNodes])
 
-  const isFetching =
-    isFetchingModelLineage || isFetchingModels || isMergingModels
-  console.log('isFetchingModelLineage', isFetchingModelLineage)
-  console.log('isFetchingModels', isFetchingModels)
-  console.log('isMegringModels', isMergingModels)
-  console.log('isFetching', isFetching)
+  function handleLineageWorkerMessage(e: MessageEvent): void {
+    if (e.data.topic === 'lineage') {
+      setIsMergingModels(false)
+      setNodeConnections(e.data.payload.nodesConnections)
+      setLineage(e.data.payload.lineage)
 
+      if (
+        Object.values(e.data.payload?.lineage ?? {}).length > WITH_COLUMNS_LIMIT
+      ) {
+        setWithColumns(false)
+      }
+    }
+
+    if (e.data.topic === 'error') {
+      handleError?.(e.data.error)
+      setIsMergingModels(false)
+    }
+  }
+
+  const isFetching =
+    isFetchingModelLineage || isFetchingModels || isMegringModels
 
   return (
     <div className="relative h-full w-full overflow-hidden">

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

@@ -46,7 +46,7 @@ export default function ModelNode({
   } = useLineageFlow()
 
   const columns: Column[] = useMemo(() => {
-    const model = models.get(id)
+    const model = models[id]
     const modelColumns = model?.columns ?? []
 
     Object.keys(lineage[id]?.columns ?? {}).forEach((column: string) => {

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

@@ -172,7 +172,7 @@ function getNodeMap({
   const modelNames = Object.keys(lineage)
 
   return modelNames.reduce((acc: Record<string, Node>, modelName: string) => {
-    const model = models.get(modelName)
+    const model = models[modelName]
     const node = createGraphNode(modelName, {
       label: model?.displayName ?? modelName,
       withColumns,
@@ -187,7 +187,7 @@ function getNodeMap({
           : EnumLineageNodeModelType.cte,
     })
     const columnsCount = withColumns
-      ? (models.get(modelName)?.columns?.length ?? 0)
+      ? (models[modelName]?.columns?.length ?? 0)
       : 0
 
     const maxWidth = Math.min(
@@ -216,7 +216,7 @@ function getNodeMap({
 
   function getNodeMaxWidth(label: string, hasColumns: boolean = false): number {
     const defaultWidth = label.length * CHAR_WIDTH
-    const columns = models.get(label)?.columns ?? []
+    const columns = models[label]?.columns ?? []
 
     return hasColumns
       ? Math.max(...columns.map(getColumnWidth), defaultWidth)

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

@@ -0,0 +1,6 @@
+import { type LineageColumn } from '@/api/client'
+
+export interface Lineage {
+  models: string[]
+  columns?: Record<string, LineageColumn>
+}

vscode/react/src/domain/lineage.ts

@@ -18,7 +18,7 @@ export function isNumber(value: unknown): value is number {
   return typeof value === 'number'
 }
 
-export function isPrimitive(value: unknown): value is Primitive {
+export function isPrimitive(value: unknown): value is string | number | boolean {
   return isString(value) || isNumber(value) || typeof value === 'boolean'
 }
 
@@ -130,7 +130,7 @@ export function toRatio(
   return (top / bottom) * multiplier
 }
 
-export function parseJSON<T>(value: string | null): Optional<T> {
+export function parseJSON<T>(value: string | null): T | undefined {
   if (isNil(value)) return undefined
 
   try {

vscode/react/src/utils/index.ts

@@ -0,0 +1,3 @@
+import createLineageWorker from './lineage.ts?worker&inline'
+
+export { createLineageWorker }

vscode/react/src/workers/index.ts

@@ -0,0 +1,126 @@
+import { isFalse, isNil, isStringEmptyOrNil, toID } from '@/utils/index'
+import { type Lineage } from '@/domain/lineage'
+
+export interface ConnectedNode {
+  id?: string
+  edges: ConnectedNode[]
+}
+
+const EnumDirection = {
+  Upstream: 'upstream',
+  Downstream: 'downstream',
+} as const
+
+type Direction = (typeof EnumDirection)[keyof typeof EnumDirection]
+
+const scope = self as any
+
+scope.onmessage = async (e: MessageEvent) => {
+  if (e.data.topic === 'lineage') {
+    try {
+      const { currentLineage, newLineage, mainNode } = e.data.payload
+      const lineage = await mergeLineageWithModels(currentLineage, newLineage)
+      const nodesConnections = await getNodesConnections(mainNode, lineage)
+
+      scope.postMessage({
+        topic: 'lineage',
+        payload: {
+          lineage,
+          nodesConnections,
+        },
+      })
+    } catch (error) {
+      scope.postMessage({
+        topic: 'error',
+        error,
+      })
+    }
+  }
+}
+
+async function mergeLineageWithModels(
+  currentLineage: Record<string, Lineage> = {},
+  data: Record<string, string[]> = {},
+): Promise<Record<string, Lineage>> {
+  return Object.entries(data).reduce(
+    (acc: Record<string, Lineage>, [key, models = []]) => {
+      key = encodeURI(key)
+
+      acc[key] = {
+        models: models.map(encodeURI),
+        columns: currentLineage?.[key]?.columns ?? undefined,
+      }
+
+      return acc
+    },
+    {},
+  )
+}
+
+async function getNodesConnections(
+  mainNode: string,
+  lineage: Record<string, Lineage> = {},
+): Promise<Record<string, ConnectedNode>> {
+  return new Promise((resolve, reject) => {
+    if (isNil(lineage) || isNil(mainNode)) return {}
+
+    const distances: Record<string, ConnectedNode> = {}
+
+    try {
+      getConnectedNodes(EnumDirection.Upstream, mainNode, lineage, distances)
+      getConnectedNodes(EnumDirection.Downstream, mainNode, lineage, distances)
+    } catch (error) {
+      reject(error)
+    }
+
+    resolve(distances)
+  })
+}
+
+function getConnectedNodes(
+  direction: Direction = EnumDirection.Downstream,
+  node: string,
+  lineage: Record<string, Lineage> = {},
+  result: Record<string, ConnectedNode> = {},
+): void {
+  const isDownstream = direction === EnumDirection.Downstream
+  let models: string[] = []
+
+  if (isDownstream) {
+    models = Object.keys(lineage).filter(key =>
+      lineage[key]!.models.includes(node),
+    )
+  } else {
+    models = lineage[node]?.models ?? []
+  }
+
+  if (isFalse(node in result)) {
+    result[node] = createConnectedNode()
+  }
+
+  for (const model of models) {
+    const connectedNode = isDownstream
+      ? createConnectedNode(node, model, [result[node]!])
+      : createConnectedNode(model, node, [result[node]!])
+
+    if (model in result) {
+      result[model]!.edges.push(connectedNode)
+    } else {
+      result[model] = connectedNode
+      getConnectedNodes(direction, model, lineage, result)
+    }
+  }
+}
+
+function createConnectedNode(
+  source?: string,
+  target?: string,
+  edges: ConnectedNode[] = [],
+): ConnectedNode {
+  const id = toID(source, target)
+
+  return {
+    id: isStringEmptyOrNil(id) ? undefined : id,
+    edges,
+  }
+}