Skip to content

Commit 41064d2

Browse files
benfdkingbenfdking
committed
feat: make vscode extension work with enterprise
1 parent 1865bd9 commit 41064d2

2 files changed

Lines changed: 179 additions & 59 deletions

File tree

‎vscode/extension/src/lsp/lsp.ts‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ import { ServerOptions, LanguageClientOptions, LanguageClient, TransportKind } f
33
import { sqlmesh_lsp_exec } from "../utilities/sqlmesh/sqlmesh"
44
import { err, isErr, ok, Result } from "../utilities/functional/result"
55
import { getWorkspaceFolders } from "../utilities/common/vscodeapi"
6+
import { traceError } from "../utilities/common/log"
67

78
let outputChannel: OutputChannel | undefined
89

@@ -20,10 +21,12 @@ export class LSPClient implements Disposable {
2021

2122
const sqlmesh = await sqlmesh_lsp_exec()
2223
if (isErr(sqlmesh)) {
24+
traceError(`Failed to get sqlmesh_lsp_exec: ${sqlmesh.error}`)
2325
return sqlmesh
2426
}
2527
const workspaceFolders = getWorkspaceFolders()
2628
if (workspaceFolders.length !== 1) {
29+
traceError(`Invalid number of workspace folders: ${workspaceFolders.length}`)
2730
return err("Invalid number of workspace folders")
2831
}
2932

@@ -36,13 +39,15 @@ export class LSPClient implements Disposable {
3639
options: {
3740
cwd: workspacePath,
3841
},
42+
args: sqlmesh.value.args,
3943
},
4044
debug: {
4145
command: sqlmesh.value.bin,
4246
transport: TransportKind.stdio,
4347
options: {
4448
cwd: workspacePath,
45-
}
49+
},
50+
args: sqlmesh.value.args,
4651
}
4752
}
4853
let clientOptions: LanguageClientOptions = {

‎vscode/extension/src/utilities/sqlmesh/sqlmesh.ts‎

Lines changed: 173 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -1,82 +1,197 @@
11
import path from "path"
22
import { traceLog, traceVerbose } from "../common/log"
33
import { getInterpreterDetails } from "../common/python"
4-
import { Result, ok } from "../functional/result"
4+
import { Result, err, isErr, ok } from "../functional/result"
55
import { getProjectRoot } from "../common/utilities"
6+
import { execFile } from "child_process"
7+
import { promisify } from "util"
8+
69

710
export type sqlmesh_exec = {
8-
workspacePath: string;
9-
bin: string;
10-
env: Record<string, string | undefined>;
11+
workspacePath: string;
12+
bin: string;
13+
env: Record<string, string | undefined>;
14+
args: string[];
15+
};
16+
17+
/**
18+
* Check if tcloud is installed in the current Python environment.
19+
*
20+
* @returns A Result indicating whether tcloud is installed.
21+
*/
22+
export const is_tcloud_installed = async (): Promise<
23+
Result<boolean, string>
24+
> => {
25+
const interpreterDetails = await getInterpreterDetails()
26+
if (!interpreterDetails.path) {
27+
return err("No Python interpreter found")
28+
}
29+
traceVerbose(
30+
`Using interpreter from Python extension: ${interpreterDetails.path.join(
31+
" "
32+
)}`
33+
)
34+
35+
const pythonPath = interpreterDetails.path[0]
36+
const checkScript = `
37+
import sys
38+
if sys.version_info >= (3, 12):
39+
from importlib import metadata
40+
else:
41+
import importlib_metadata as metadata
42+
43+
try:
44+
metadata.version('tcloud')
45+
print("true")
46+
except metadata.PackageNotFoundError:
47+
print("false")
48+
`
49+
try {
50+
const execFileAsync = promisify(execFile)
51+
traceVerbose(`Checking tcloud installation with script: ${checkScript}`)
52+
const { stdout } = await execFileAsync(pythonPath, ["-c", checkScript])
53+
traceVerbose(`tcloud installation check result: ${stdout.trim()}`)
54+
return ok(stdout.trim() === "true")
55+
} catch (error) {
56+
return err(`Failed to check tcloud installation: ${error}`)
57+
}
58+
}
59+
60+
/**
61+
* Get the tcloud executable for the current Python environment.
62+
*
63+
* @returns The tcloud executable for the current Python environment.
64+
*/
65+
export const get_tcloud_bin = async (): Promise<Result<string, string>> => {
66+
const interpreterDetails = await getInterpreterDetails()
67+
if (!interpreterDetails.path) {
68+
return err("No Python interpreter found")
69+
}
70+
const pythonPath = interpreterDetails.path[0]
71+
const binPath = path.join(path.dirname(pythonPath), "tcloud")
72+
return ok(binPath)
1173
}
1274

1375
/**
1476
* Get the sqlmesh executable for the current workspace.
15-
*
77+
*
1678
* @returns The sqlmesh executable for the current workspace.
1779
*/
1880
export const sqlmesh_exec = async (): Promise<Result<sqlmesh_exec, string>> => {
19-
const projectRoot = await getProjectRoot()
20-
const workspacePath = projectRoot.uri.fsPath
21-
const interpreterDetails = await getInterpreterDetails()
22-
traceLog(`Interpreter details: ${JSON.stringify(interpreterDetails)}`)
23-
if (interpreterDetails.path) {
24-
traceVerbose(`Using interpreter from Python extension: ${interpreterDetails.path.join(' ')}`)
81+
const projectRoot = await getProjectRoot()
82+
const workspacePath = projectRoot.uri.fsPath
83+
const interpreterDetails = await getInterpreterDetails()
84+
traceLog(`Interpreter details: ${JSON.stringify(interpreterDetails)}`)
85+
if (interpreterDetails.path) {
86+
traceVerbose(
87+
`Using interpreter from Python extension: ${interpreterDetails.path.join(
88+
" "
89+
)}`
90+
)
91+
}
92+
if (interpreterDetails.isVirtualEnvironment) {
93+
traceLog("Using virtual environment")
94+
const tcloudInstalled = await is_tcloud_installed()
95+
if (isErr(tcloudInstalled)) {
96+
return tcloudInstalled
2597
}
26-
if (interpreterDetails.isVirtualEnvironment) {
27-
traceLog('Using virtual environment')
28-
const binPath = path.join(interpreterDetails.binPath!, 'sqlmesh')
29-
traceLog(`Bin path: ${binPath}`)
30-
return ok({
31-
bin: binPath,
32-
workspacePath,
33-
env: {
34-
PYTHONPATH: interpreterDetails.path?.[0],
35-
VIRTUAL_ENV: path.dirname(interpreterDetails.binPath!),
36-
PATH: path.join(path.dirname(interpreterDetails.binPath!), 'bin')
37-
}
38-
})
39-
} else {
40-
return ok({
41-
bin: 'sqlmesh',
42-
workspacePath,
43-
env: {},
44-
})
98+
if (tcloudInstalled.value) {
99+
const tcloudBin = await get_tcloud_bin()
100+
if (isErr(tcloudBin)) {
101+
return tcloudBin
102+
}
103+
return ok({
104+
bin: `${tcloudBin.value} sqlmesh`,
105+
workspacePath,
106+
env: {
107+
PYTHONPATH: interpreterDetails.path?.[0],
108+
VIRTUAL_ENV: path.dirname(interpreterDetails.binPath!),
109+
PATH: interpreterDetails.binPath!,
110+
},
111+
args:[],
112+
})
45113
}
114+
const binPath = path.join(interpreterDetails.binPath!, "sqlmesh")
115+
traceLog(`Bin path: ${binPath}`)
116+
return ok({
117+
bin: binPath,
118+
workspacePath,
119+
env: {
120+
PYTHONPATH: interpreterDetails.path?.[0],
121+
VIRTUAL_ENV: path.dirname(interpreterDetails.binPath!),
122+
PATH: interpreterDetails.binPath!,
123+
},
124+
args: [],
125+
})
126+
} else {
127+
return ok({
128+
bin: "sqlmesh",
129+
workspacePath,
130+
env: {},
131+
args: [],
132+
})
133+
}
46134
}
47135

48136
/**
49137
* Get the sqlmesh_lsp executable for the current workspace.
50-
*
138+
*
51139
* @returns The sqlmesh_lsp executable for the current workspace.
52140
*/
53-
export const sqlmesh_lsp_exec = async (): Promise<Result<sqlmesh_exec, string>> => {
54-
const projectRoot = await getProjectRoot()
55-
const workspacePath = projectRoot.uri.fsPath
56-
const interpreterDetails = await getInterpreterDetails()
57-
traceLog(`Interpreter details: ${JSON.stringify(interpreterDetails)}`)
58-
if (interpreterDetails.path) {
59-
traceVerbose(`Using interpreter from Python extension: ${interpreterDetails.path.join(' ')}`)
141+
export const sqlmesh_lsp_exec = async (): Promise<
142+
Result<sqlmesh_exec, string>
143+
> => {
144+
const projectRoot = await getProjectRoot()
145+
const workspacePath = projectRoot.uri.fsPath
146+
const interpreterDetails = await getInterpreterDetails()
147+
traceLog(`Interpreter details: ${JSON.stringify(interpreterDetails)}`)
148+
if (interpreterDetails.path) {
149+
traceVerbose(
150+
`Using interpreter from Python extension: ${interpreterDetails.path.join(
151+
" "
152+
)}`
153+
)
154+
}
155+
if (interpreterDetails.isVirtualEnvironment) {
156+
traceLog("Using virtual environment")
157+
const tcloudInstalled = await is_tcloud_installed()
158+
if (isErr(tcloudInstalled)) {
159+
return tcloudInstalled
60160
}
61-
if (interpreterDetails.isVirtualEnvironment) {
62-
traceLog('Using virtual environment')
63-
const binPath = path.join(interpreterDetails.binPath!, 'sqlmesh_lsp')
64-
traceLog(`Bin path: ${binPath}`)
65-
return ok({
66-
bin: binPath,
67-
workspacePath,
68-
env: {
69-
PYTHONPATH: interpreterDetails.path?.[0],
70-
VIRTUAL_ENV: path.dirname(interpreterDetails.binPath!),
71-
PATH: path.join(path.dirname(interpreterDetails.binPath!), 'bin')
72-
}
73-
})
74-
} else {
75-
return ok({
76-
bin: 'sqlmesh_lsp',
77-
workspacePath,
78-
env: {},
79-
})
161+
if (tcloudInstalled.value) {
162+
traceLog("Tcloud installed, installing sqlmesh")
163+
const tcloudBin = await get_tcloud_bin()
164+
if (isErr(tcloudBin)) {
165+
return tcloudBin
166+
}
167+
const execFileAsync = promisify(execFile)
168+
await execFileAsync(tcloudBin.value, ["install_sqlmesh"], {
169+
cwd: workspacePath,
170+
env: {
171+
PYTHONPATH: interpreterDetails.path?.[0],
172+
VIRTUAL_ENV: path.dirname(interpreterDetails.binPath!),
173+
PATH: interpreterDetails.binPath!,
174+
},
175+
})
80176
}
81-
177+
const binPath = path.join(interpreterDetails.binPath!, "sqlmesh_lsp")
178+
traceLog(`Bin path: ${binPath}`)
179+
return ok({
180+
bin: binPath,
181+
workspacePath,
182+
env: {
183+
PYTHONPATH: interpreterDetails.path?.[0],
184+
VIRTUAL_ENV: path.dirname(interpreterDetails.binPath!),
185+
PATH: path.join(path.dirname(interpreterDetails.binPath!), "bin"),
186+
},
187+
args: [],
188+
})
189+
} else {
190+
return ok({
191+
bin: "sqlmesh_lsp",
192+
workspacePath,
193+
env: {},
194+
args: [],
195+
})
196+
}
82197
}

0 commit comments

Comments
 (0)