temp [ci skip]

Author: benfdking

vscode/bus/src/callbacks.ts

@@ -1,6 +1,6 @@
 export type Callback = {
   openFile: {
-    path: string
+    uri: string
   }
   queryRequest: {
     requestId: string

vscode/extension/src/webviews/lineagePanel.ts

@@ -1,4 +1,4 @@
-import { CallbackEvent } from '@bus/callbacks'
+import { Callback, CallbackEvent } from '@bus/callbacks'
 import {
   Disposable,
   Uri,
@@ -57,42 +57,24 @@ export class LineagePanel implements WebviewViewProvider, Disposable {
     // Set content options for external URL access
     // Set up message listener for events from the iframe
     webviewView.webview.onDidReceiveMessage(
-      async message => {
-        if (!message) {
+      async request => {
+        console.log('received message', request)
+        if (!request) {
           return
         }
-        if (!message.key) {
+        if (!request.key) {
           return
         }
+        const message: CallbackEvent = request
         switch (message.key) {
-          case 'vscode_callback': {
-            const payload: CallbackEvent = message.payload
-            // Handle callback events from the iframe
-            switch (payload.key) {
-              case 'openFile': {
-                const workspaceFolders = getWorkspaceFolders()
-                if (workspaceFolders.length != 1) {
-                  throw new Error('Only one workspace folder is supported')
-                }
-                const workspaceFolder = workspaceFolders[0]
-                const fullPath = Uri.joinPath(
-                  workspaceFolder.uri,
-                  payload.payload.path,
-                )
-                const document = await workspace.openTextDocument(fullPath)
-                await window.showTextDocument(document)
-                break
-              }
-              case 'queryRequest': {
-                throw new Error('queryRequest is not supported')
-              }
-              default: {
-                console.error(
-                  'Unhandled message type under vscode_callback: ',
-                  message,
-                )
-              }
+          case 'openFile': {
+            const workspaceFolders = getWorkspaceFolders()
+            if (workspaceFolders.length != 1) {
+              throw new Error('Only one workspace folder is supported')
             }
+            const fullPath = Uri.parse(message.payload.uri)
+            const document = await workspace.openTextDocument(fullPath)
+            await window.showTextDocument(document)
             break
           }
           case 'queryRequest': {
@@ -102,7 +84,7 @@ export class LineagePanel implements WebviewViewProvider, Disposable {
               'sqlmesh/api',
               payload,
             )
-             webviewView.webview.postMessage({
+            webviewView.webview.postMessage({
               key: 'query_response',
               payload: response,
               requestId,

vscode/react/src/domain/model.ts

@@ -3,6 +3,7 @@ interface Model {
   fqn: string
   path: string
   dialect: string
+  full_path: string
   // type: ModelType
   // columns: Column[]
   // default_catalog?: ModelDefaultCatalog

vscode/react/src/hooks/vscode.ts

@@ -7,4 +7,3 @@ import { sendVSCodeMessage } from '@/utils/vscodeapi'
  * so we need to send messages to the parent window
  */
 export const useVSCode = () => sendVSCodeMessage
-

vscode/react/src/pages/lineage.tsx

@@ -8,19 +8,24 @@ import { useApiModels } from '@/api'
 import LineageFlowProvider from '@/components/graph/context'
 import { ModelLineage } from '@/components/graph/ModelLineage'
 import { useVSCode } from '@/hooks/vscode'
-import React from 'react'
+import React, { useEffect, useState } from 'react'
 import { ModelSQLMeshModel } from '@/domain/sqlmesh-model'
+import { useEventBus } from '@/hooks/eventBus'
+import type { VSCodeEvent } from '@bus/callbacks'
+import { setTokenSourceMapRange } from 'typescript'
 
 export function LineagePage() {
+  const { emit } = useEventBus()
+
   // Handle messages from VSCode extension
   React.useEffect(() => {
     const handleMessage = (event: MessageEvent) => {
       // Ensure the message is from VSCode
       if (event.data && event.data.key === 'vscode_send') {
-        const payload = event.data.payload
+        const payload: VSCodeEvent = event.data.payload
         switch (payload.key) {
           case 'changeFocusOnFile':
-            console.log('Changing focus to file:', payload.payload.path)
+            console.log('emitted focus to file:', payload.payload.path)
             emit('changeFocusedFile', { fileUri: payload.payload.path })
             break
           default:
@@ -57,8 +62,16 @@ export function LineagePage() {
 }
 
 function Lineage() {
+  const { on } = useEventBus()
+  const [selectedModelSet, setSelectedModelSet] = useState<string | undefined>(undefined)
   const { data: models, isLoading: isLoadingModels } = useApiModels()
-  if (isLoadingModels || models === undefined) {
+  React.useEffect(() => {
+    if (selectedModelSet === undefined && models && models.length > 0) {
+      setSelectedModelSet(models[0].name)
+    }
+  }, [models, selectedModelSet])
+
+  if (isLoadingModels || models === undefined || selectedModelSet === undefined) {
     return <div>Loading models...</div>
   }
   if (!Array.isArray(models)) {
@@ -71,7 +84,15 @@ function Lineage() {
     },
     {} as Record<string, Model>,
   )
-  const selectedModel = Object.keys(modelsRecord)[0]
+  const selectedModel = selectedModelSet
+  on('changeFocusedFile', (fileUri) => {
+    console.log('on changeFocusedFile', fileUri)
+    const full_path = fileUri.fileUri.startsWith("file://") ? fileUri.fileUri.substring(7) : fileUri.fileUri
+    const model = Object.values(modelsRecord).find((m: Model) => m.full_path === full_path)
+    if (model) {
+      setSelectedModelSet(model.name)
+    }
+  })
 
   return (
     <LineageComponentFromWeb
@@ -96,7 +117,7 @@ export function LineageComponentFromWeb({
     if (!model) {
       throw new Error('Model not found')
     }
-    vscode('openFile', { path: model.path })
+    vscode('openFile', { uri: "file://" + model.full_path })
   }
 
   function handleError(error: any): void {

vscode/react/src/utils/vscodeapi.ts

@@ -14,8 +14,10 @@ export const sendVSCodeMessage = <K extends keyof Callback> (
     payload: payload,
   }
   if (window !== window.parent) {
+    console.log('sending message to parent', eventPayload)
     window.parent.postMessage(eventPayload, '*')
   } else {
+    console.log('sending message to self', eventPayload)
     getVSCodeAPI().postMessage(eventPayload)
   }
 }