Skip to content

Commit bc91c39

Browse files
benfdkingbenfdking
committed
temp
1 parent 285afd8 commit bc91c39

6 files changed

Lines changed: 224 additions & 93 deletions

File tree

‎vscode/bus/src/callbacks.ts‎

Lines changed: 41 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,13 @@
1+
import { Result } from './result'
2+
3+
export type CallbackShape = Record<string, any>
4+
15
export type Callback = {
26
openFile: {
37
uri: string
48
}
5-
queryRequest: {
6-
requestId: string
7-
url: string
8-
method?: string
9-
params?: Record<string, string>
10-
body?: Record<string, any>
11-
}
12-
}
9+
rpcResponse: RPCResponse
10+
} & CallbackShape
1311

1412
/**
1513
* A tuple type representing a callback event with its associated payload.
@@ -23,22 +21,52 @@ export type CallbackEvent = {
2321
[K in keyof Callback]: { key: K; payload: Callback[K] }
2422
}[keyof Callback]
2523

24+
export type VSCodeCallbackShape = Record<string, any>
25+
2626
/**
2727
* A tuple type representing a VSCode event with its associated payload.
2828
*/
2929
export type VSCodeCallback = {
3030
changeFocusOnFile: {
3131
path: string
3232
}
33-
queryResponse: {
34-
requestId: string
35-
response: any
36-
}
3733
savedFile: {
3834
fileUri: string
3935
}
40-
}
36+
rpcRequest: RPCRequest
37+
} & VSCodeCallbackShape
4138

4239
export type VSCodeEvent = {
4340
[K in keyof VSCodeCallback]: { key: K; payload: VSCodeCallback[K] }
4441
}[keyof VSCodeCallback]
42+
43+
type RPCMethodsShape = Record<string, { params: any; result: any }>
44+
45+
export type RPCMethods = {
46+
get_active_file: {
47+
params: {}
48+
result: {
49+
fileUri?: string
50+
}
51+
}
52+
api_query: {
53+
params: {
54+
url: string
55+
method: string
56+
params: any
57+
body: any
58+
}
59+
result: any
60+
}
61+
} & RPCMethodsShape
62+
63+
export type RPCRequest = {
64+
requestId: string
65+
method: keyof RPCMethods
66+
params: RPCMethods[keyof RPCMethods]['params']
67+
}
68+
69+
export type RPCResponse = {
70+
requestId: string
71+
result: Result<RPCMethods[keyof RPCMethods]['result'], string>
72+
}

‎vscode/extension/src/webviews/lineagePanel.ts‎

Lines changed: 40 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import { CallbackEvent } from '@bus/callbacks'
1+
import { CallbackEvent, RPCRequest } from '@bus/callbacks'
22
import {
33
Disposable,
44
Uri,
@@ -65,6 +65,7 @@ export class LineagePanel implements WebviewViewProvider, Disposable {
6565
// Set up message listener for events from the iframe
6666
webviewView.webview.onDidReceiveMessage(
6767
async request => {
68+
console.log('request', request)
6869
if (!request) {
6970
return
7071
}
@@ -83,18 +84,44 @@ export class LineagePanel implements WebviewViewProvider, Disposable {
8384
await window.showTextDocument(document)
8485
break
8586
}
86-
case 'queryRequest': {
87-
const payload = message.payload
88-
const requestId = message.payload.requestId
89-
const response = await this.lsp.call_custom_method(
90-
'sqlmesh/api',
91-
payload as any,
92-
)
93-
webviewView.webview.postMessage({
94-
key: 'query_response',
95-
payload: response,
96-
requestId,
97-
})
87+
case 'rpcRequest': {
88+
console.log('rpcRequest', message.payload)
89+
const payload: RPCRequest = message.payload
90+
const requestId = payload.requestId
91+
switch (payload.method) {
92+
case 'api_query': {
93+
const response = await this.lsp.call_custom_method(
94+
'sqlmesh/api',
95+
payload.params,
96+
)
97+
const responseCallback: CallbackEvent = {
98+
key: 'rpcResponse',
99+
payload: {
100+
requestId,
101+
result: response,
102+
},
103+
}
104+
webviewView.webview.postMessage(responseCallback)
105+
break
106+
}
107+
case 'get_active_file': {
108+
const active_file = window.activeTextEditor?.document.uri.fsPath
109+
const responseCallback: CallbackEvent = {
110+
key: 'rpcResponse',
111+
payload: {
112+
requestId,
113+
result: {
114+
fileUri: active_file,
115+
},
116+
},
117+
}
118+
webviewView.webview.postMessage(responseCallback)
119+
break
120+
}
121+
default: {
122+
throw new Error(`Unhandled RPC method: ${payload.method}`)
123+
}
124+
}
98125
break
99126
}
100127
default:

‎vscode/react/src/api/instance.ts‎

Lines changed: 13 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import { sendVSCodeMessage } from '@/utils/vscodeapi'
1+
import { callRpc } from '@/utils/rpc'
22
import { isErr } from '@bus/result'
33

44
declare global {
@@ -38,43 +38,16 @@ export async function fetchAPI<T = any, B extends object = any>(
3838
config: FetchOptions<B>,
3939
_options?: Partial<FetchOptionsWithSignal>,
4040
): Promise<T & ResponseWithDetail> {
41-
// Generate a unique ID for this request
42-
// Create a promise that will resolve when we get a response with matching ID
43-
return new Promise((resolve, reject) => {
44-
const requestId = `query_${Date.now()}_${Math.random().toString(36).substr(2, 9)}`
45-
const messageHandler = (event: MessageEvent) => {
46-
if (
47-
event.data &&
48-
event.data.key === 'query_response' &&
49-
event.data.requestId === requestId
50-
) {
51-
// Remove the listener once we get our response
52-
window.removeEventListener('message', messageHandler)
53-
54-
const payload = event.data.payload
55-
if (isErr(payload)) {
56-
reject(new Error(payload.error as string))
57-
} else {
58-
resolve(payload.value.data)
59-
}
60-
}
61-
}
62-
63-
// Add the listener
64-
window.addEventListener('message', messageHandler)
65-
66-
sendVSCodeMessage('queryRequest', {
67-
requestId,
68-
url: config.url,
69-
params: config.params as any,
70-
body: config.data,
71-
method: config.method,
72-
})
73-
74-
// Set a timeout to prevent hanging promises
75-
setTimeout(() => {
76-
window.removeEventListener('message', messageHandler)
77-
reject(new Error('Query request timed out'))
78-
}, 30000) // 30 second timeout
79-
})
41+
const request = {
42+
url: config.url,
43+
method: config.method,
44+
params: config.params,
45+
body: config.data,
46+
}
47+
const result = await callRpc('api_query', request)
48+
if (isErr(result)) {
49+
throw new Error(result.error)
50+
}
51+
const response = result.value.data
52+
return response
8053
}

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -78,7 +78,7 @@ export function ModelLineage({
7878
useApiModelLineage(model.name)
7979
const { isFetching: isFetchingModels } = useApiModels()
8080

81-
const [isMegringModels, setIsMergingModels] = useState(false)
81+
const [isMergingModels, setIsMergingModels] = useState(false)
8282
const [modelLineage, setModelLineage] = useState<
8383
ModelLineageApiLineageModelNameGet200 | undefined
8484
>(undefined)
@@ -164,7 +164,7 @@ export function ModelLineage({
164164
}
165165

166166
const isFetching =
167-
isFetchingModelLineage || isFetchingModels || isMegringModels
167+
isFetchingModelLineage || isFetchingModels || isMergingModels
168168

169169
return (
170170
<div className="relative h-full w-full overflow-hidden">

‎vscode/react/src/pages/lineage.tsx‎

Lines changed: 74 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,8 @@ import { useEventBus } from '@/hooks/eventBus'
1515
import type { VSCodeEvent } from '@bus/callbacks'
1616
import { URI } from 'vscode-uri'
1717
import type { Model } from '@/api/client'
18+
import { useRpc } from '@/utils/rpc'
19+
import { isErr } from '@bus/result'
1820

1921
export function LineagePage() {
2022
const { emit } = useEventBus()
@@ -66,49 +68,96 @@ export function LineagePage() {
6668
}
6769

6870
function Lineage() {
69-
const [selectedModelSet, setSelectedModelSet] = useState<string | undefined>(
71+
const [selectedModel, setSelectedModel] = useState<string | undefined>(
7072
undefined,
7173
)
7274
const { on } = useEventBus()
7375
const queryClient = useQueryClient()
7476

7577
const { data: models, isLoading: isLoadingModels } = useApiModels()
78+
const rpc = useRpc()
79+
7680
React.useEffect(() => {
77-
if (selectedModelSet === undefined && models && Array.isArray(models)) {
78-
setSelectedModelSet(models[0].name)
81+
const fetchFirstModelIfNotSet = async (
82+
models: Model[],
83+
): Promise<string | undefined> => {
84+
console.log('fetchFirstModelIfNotSet', models)
85+
if (!Array.isArray(models)) {
86+
return undefined
87+
}
88+
const activeFile = await rpc('get_active_file', {})
89+
console.log('activeFileResponse', activeFile)
90+
// @ts-ignore
91+
if (!activeFile.fileUri) {
92+
return models[0].name
93+
}
94+
// @ts-ignore
95+
const fileUri: string = activeFile.fileUri
96+
const filePath = URI.parse(fileUri).fsPath
97+
const model = models.find((m: Model) => m.full_path === filePath)
98+
if (model) {
99+
return model.name
100+
}
101+
return undefined
102+
}
103+
if (selectedModel === undefined && Array.isArray(models)) {
104+
fetchFirstModelIfNotSet(models).then(modelName => {
105+
if (modelName && selectedModel === undefined) {
106+
setSelectedModel(modelName)
107+
}
108+
})
109+
}
110+
}, [models, selectedModel])
111+
112+
const modelsRecord =
113+
Array.isArray(models) &&
114+
models.reduce(
115+
(acc, model) => {
116+
acc[model.name] = model
117+
return acc
118+
},
119+
{} as Record<string, Model>,
120+
)
121+
122+
React.useEffect(() => {
123+
const handleChangeFocusedFile = (fileUri: { fileUri: string }) => {
124+
const full_path = URI.parse(fileUri.fileUri).fsPath
125+
const model = Object.values(modelsRecord).find(
126+
m => m.full_path === full_path,
127+
)
128+
if (model) {
129+
setSelectedModel(model.name)
130+
}
79131
}
80-
}, [models, selectedModelSet])
132+
133+
const handleSavedFile = () => {
134+
queryClient.invalidateQueries()
135+
}
136+
137+
const offChangeFocusedFile = on(
138+
'changeFocusedFile',
139+
handleChangeFocusedFile,
140+
)
141+
const offSavedFile = on('savedFile', handleSavedFile)
142+
143+
// If your event bus returns an "off" function, call it on cleanup
144+
return () => {
145+
if (offChangeFocusedFile) offChangeFocusedFile()
146+
if (offSavedFile) offSavedFile()
147+
}
148+
}, [on, queryClient, modelsRecord])
81149

82150
if (
83151
isLoadingModels ||
84152
models === undefined ||
85-
selectedModelSet === undefined
153+
modelsRecord === false ||
154+
selectedModel === undefined
86155
) {
87156
return <div>Loading models...</div>
88157
}
89158
if (!Array.isArray(models)) {
90159
return <div>Error: Models data is not in the expected format</div>
91160
}
92-
const modelsRecord = models.reduce(
93-
(acc, model) => {
94-
acc[model.name] = model
95-
return acc
96-
},
97-
{} as Record<string, Model>,
98-
)
99-
const selectedModel = selectedModelSet
100-
on('changeFocusedFile', fileUri => {
101-
const full_path = URI.parse(fileUri.fileUri).fsPath
102-
const model = Object.values(modelsRecord).find(
103-
m => m.full_path === full_path,
104-
)
105-
if (model) {
106-
setSelectedModelSet(model.name)
107-
}
108-
})
109-
on('savedFile', () => {
110-
queryClient.invalidateQueries()
111-
})
112161

113162
return (
114163
<LineageComponentFromWeb

0 commit comments

Comments
 (0)