chore(vscode): adding typing to worker message bus (#4918)

Author: benfdking

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

@@ -17,7 +17,7 @@ import ReactFlow, {
 } from 'reactflow'
 import Loading from '@/components/loading/Loading'
 import Spinner from '@/components/logo/Spinner'
-import { createLineageWorker } from '@/workers/index'
+import { createLineageWorker } from '@/components/graph/workers/index'
 import { isArrayEmpty, isFalse, isNil, isNotNil } from '@/utils/index'
 import clsx from 'clsx'
 import ModelNode from './ModelNode'
@@ -35,7 +35,13 @@ import {
   type ModelLineage as ModelLineageType,
 } from '@/domain/lineage'
 import './Graph.css'
-import { toKeys } from './types'
+import {
+  toKeys,
+  type LineageWorkerMessage,
+  type LineageWorkerRequestMessage,
+  type LineageWorkerResponseMessage,
+  type LineageWorkerErrorMessage,
+} from './types'
 import { encode } from '@/domain/models'
 
 const WITH_COLUMNS_LIMIT = 30
@@ -95,14 +101,15 @@ export function ModelLineage({
 
         setIsMergingModels(true)
 
-        lineageWorker.postMessage({
+        const message: LineageWorkerRequestMessage = {
           topic: 'lineage',
           payload: {
             currentLineage: {},
             newLineage: data,
             mainNode: model.fqn,
           },
-        })
+        }
+        lineageWorker.postMessage(message)
       })
       .catch(error => {
         handleError?.(error)
@@ -146,21 +153,25 @@ export function ModelLineage({
     setHighlightedNodes(highlightedNodes ?? {})
   }, [highlightedNodes])
 
-  function handleLineageWorkerMessage(e: MessageEvent): void {
+  function handleLineageWorkerMessage(
+    e: MessageEvent<LineageWorkerMessage>,
+  ): void {
     if (e.data.topic === 'lineage') {
+      const message = e.data as LineageWorkerResponseMessage
       setIsMergingModels(false)
-      setNodeConnections(e.data.payload.nodesConnections)
-      setLineage(e.data.payload.lineage)
+      setNodeConnections(message.payload.nodesConnections)
+      setLineage(message.payload.lineage)
 
       if (
-        Object.values(e.data.payload?.lineage ?? {}).length > WITH_COLUMNS_LIMIT
+        Object.values(message.payload.lineage ?? {}).length > WITH_COLUMNS_LIMIT
       ) {
         setWithColumns(false)
       }
     }
 
     if (e.data.topic === 'error') {
-      handleError?.(e.data.error)
+      const message = e.data as LineageWorkerErrorMessage
+      handleError?.(message.error)
       setIsMergingModels(false)
     }
   }

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

@@ -14,7 +14,7 @@ import type { ModelEncodedFQN, ModelName } from '@/domain/models'
 import type { ColumnName } from '@/domain/column'
 import type { Model } from '@/api/client'
 import { toID, toKeys } from './types'
-import type { ConnectedNode } from '@/workers/lineage'
+import type { ConnectedNode } from '@/components/graph/types'
 
 export interface Connections {
   left: string[]

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

@@ -15,7 +15,7 @@ import {
   type LineageNodeModelType,
 } from './ModelNode'
 import type { Lineage } from '@/domain/lineage'
-import type { ConnectedNode } from '@/workers/lineage'
+import type { ConnectedNode } from '@/components/graph/types'
 import { encode, type ModelEncodedFQN, type ModelURI } from '@/domain/models'
 import type { Column, ColumnName } from '@/domain/column'
 import type { ModelSQLMeshModel } from '@/domain/sqlmesh-model'

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

@@ -1,9 +1,12 @@
 import type { ColumnName } from '@/domain/column'
 import type { ModelEncodedFQN } from '@/domain/models'
 import type { Branded } from '@bus/brand'
+import type { Lineage } from '@/domain/lineage'
 
 export type Side = 'left' | 'right'
 
+export type Direction = 'upstream' | 'downstream'
+
 export type NodeId = string
 
 export type EdgeId = string
@@ -55,3 +58,44 @@ export function toKeys<K extends string, V>(obj: Record<K, V>): K[] {
 }
 
 export type ModelLineage = Record<ModelEncodedFQN, ModelEncodedFQN[]>
+
+// Worker Message Types
+export interface ConnectedNode {
+  id?: string
+  edges: ConnectedNode[]
+}
+
+export interface LineageWorkerRequestPayload {
+  currentLineage: Record<string, Lineage>
+  newLineage: Record<string, string[]>
+  mainNode: string
+}
+
+export interface LineageWorkerResponsePayload {
+  lineage: Record<string, Lineage>
+  nodesConnections: Record<string, ConnectedNode>
+}
+
+export interface LineageWorkerErrorPayload {
+  error: Error
+}
+
+export interface LineageWorkerRequestMessage {
+  topic: 'lineage'
+  payload: LineageWorkerRequestPayload
+}
+
+export interface LineageWorkerResponseMessage {
+  topic: 'lineage'
+  payload: LineageWorkerResponsePayload
+}
+
+export interface LineageWorkerErrorMessage {
+  topic: 'error'
+  error: Error
+}
+
+export type LineageWorkerMessage =
+  | LineageWorkerRequestMessage
+  | LineageWorkerResponseMessage
+  | LineageWorkerErrorMessage

vscode/react/src/workers/index.ts …ct/src/components/graph/workers/index.tsvscode/react/src/workers/index.ts renamed to vscode/react/src/components/graph/workers/index.ts vscode/react/src/workers/index.ts renamed to vscode/react/src/components/graph/workers/index.ts

@@ -1,41 +1,46 @@
 import { isFalse, isNil } from '@/utils/index'
 import { type Lineage } from '@/domain/lineage'
 import type { ModelEncodedFQN } from '@/domain/models'
-import { toID, type NodeId } from '@/components/graph/types'
-
-export interface ConnectedNode {
-  id?: string
-  edges: ConnectedNode[]
+import {
+  toID,
+  type NodeId,
+  type LineageWorkerMessage,
+  type LineageWorkerRequestMessage,
+  type LineageWorkerResponseMessage,
+  type LineageWorkerErrorMessage,
+  type ConnectedNode,
+} from '@/components/graph/types'
+import type { Direction } from '../types'
+
+interface WorkerScope {
+  onmessage: ((e: MessageEvent<LineageWorkerMessage>) => void) | null
+  postMessage: (message: LineageWorkerMessage) => void
 }
 
-const EnumDirection = {
-  Upstream: 'upstream',
-  Downstream: 'downstream',
-} as const
-
-type Direction = (typeof EnumDirection)[keyof typeof EnumDirection]
-
-const scope = self as any
+const scope = self as unknown as WorkerScope
 
-scope.onmessage = async (e: MessageEvent) => {
+scope.onmessage = async (e: MessageEvent<LineageWorkerMessage>) => {
   if (e.data.topic === 'lineage') {
     try {
-      const { currentLineage, newLineage, mainNode } = e.data.payload
+      const message = e.data as LineageWorkerRequestMessage
+      const { currentLineage, newLineage, mainNode } = message.payload
       const lineage = await mergeLineageWithModels(currentLineage, newLineage)
       const nodesConnections = await getNodesConnections(mainNode, lineage)
 
-      scope.postMessage({
+      const responseMessage: LineageWorkerResponseMessage = {
         topic: 'lineage',
         payload: {
           lineage,
           nodesConnections,
         },
-      })
+      }
+      scope.postMessage(responseMessage)
     } catch (error) {
-      scope.postMessage({
+      const errorMessage: LineageWorkerErrorMessage = {
         topic: 'error',
-        error,
-      })
+        error: error as Error,
+      }
+      scope.postMessage(errorMessage)
     }
   }
 }
@@ -69,8 +74,8 @@ async function getNodesConnections(
     const distances: Record<string, ConnectedNode> = {}
 
     try {
-      getConnectedNodes(EnumDirection.Upstream, mainNode, lineage, distances)
-      getConnectedNodes(EnumDirection.Downstream, mainNode, lineage, distances)
+      getConnectedNodes('upstream', mainNode, lineage, distances)
+      getConnectedNodes('downstream', mainNode, lineage, distances)
     } catch (error) {
       reject(error)
     }
@@ -80,12 +85,12 @@ async function getNodesConnections(
 }
 
 function getConnectedNodes(
-  direction: Direction = EnumDirection.Downstream,
+  direction: Direction = 'downstream',
   node: string,
   lineage: Record<string, Lineage> = {},
   result: Record<string, ConnectedNode> = {},
 ): void {
-  const isDownstream = direction === EnumDirection.Downstream
+  const isDownstream = direction === 'downstream'
   let models: string[] = []
 
   if (isDownstream) {