Skip to content

Commit 23083f9

Browse files
authored
fix(vscode): improve the tcloud stability (#4734)
1 parent 6e90e3e commit 23083f9

9 files changed

Lines changed: 606 additions & 64 deletions

File tree

sqlmesh/lsp/main.py

Lines changed: 8 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -294,32 +294,6 @@ def did_open(ls: LanguageServer, params: types.DidOpenTextDocumentParams) -> Non
294294
SQLMeshLanguageServer._diagnostics_to_lsp_diagnostics(diagnostics),
295295
)
296296

297-
@self.server.feature(types.TEXT_DOCUMENT_DID_CHANGE)
298-
def did_change(ls: LanguageServer, params: types.DidChangeTextDocumentParams) -> None:
299-
if self.lsp_context is None:
300-
current_path = Path.cwd()
301-
self._ensure_context_in_folder(current_path)
302-
if self.lsp_context is None:
303-
ls.log_trace("No context found after did_change")
304-
return
305-
306-
uri = URI(params.text_document.uri)
307-
context = self._context_get_or_load(uri)
308-
309-
models = context.map[uri.to_path()]
310-
if models is None or not isinstance(models, ModelTarget):
311-
return
312-
313-
# Get diagnostics from context (which handles caching)
314-
diagnostics = context.lint_model(uri)
315-
316-
# Only publish diagnostics if client doesn't support pull diagnostics
317-
if not self.client_supports_pull_diagnostics:
318-
ls.publish_diagnostics(
319-
params.text_document.uri,
320-
SQLMeshLanguageServer._diagnostics_to_lsp_diagnostics(diagnostics),
321-
)
322-
323297
@self.server.feature(types.TEXT_DOCUMENT_DID_SAVE)
324298
def did_save(ls: LanguageServer, params: types.DidSaveTextDocumentParams) -> None:
325299
uri = URI(params.text_document.uri)
@@ -753,15 +727,19 @@ def _ensure_context_for_document(
753727

754728
loaded = False
755729
# Ascend directories to look for config
756-
while path.parents and not loaded:
730+
current_dir = path.parent # Start from the file's parent directory
731+
while current_dir.parents and not loaded:
757732
for ext in ("py", "yml", "yaml"):
758-
config_path = path / f"config.{ext}"
733+
config_path = current_dir / f"config.{ext}"
759734
if config_path.exists():
760-
if self._create_lsp_context([path]):
735+
if self._create_lsp_context([current_dir]):
761736
loaded = True
762737
# Re-check context for the document now that it's loaded
763738
return self._ensure_context_for_document(document_uri)
764-
path = path.parent
739+
# Check if we've reached the filesystem root to prevent infinite loops
740+
if current_dir == current_dir.parent:
741+
break
742+
current_dir = current_dir.parent
765743

766744
# If still no context found, try the workspace folders
767745
if not loaded:

tooling/vscode/extensions.json

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
"dbaeumer.vscode-eslint",
66
"amodio.tsl-problem-matcher",
77
"ms-vscode.extension-test-runner",
8-
"ms-playwright.playwright"
8+
"ms-playwright.playwright",
9+
"esbenp.prettier-vscode"
910
]
1011
}

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

Lines changed: 51 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -147,51 +147,69 @@ export const installSqlmeshEnterprise = async (
147147
return ok(true)
148148
}
149149

150+
let installationLock: Promise<Result<boolean, ErrorType>> | undefined = undefined
151+
150152
/**
151153
* Checks if sqlmesh enterprise is installed and updated. If not, it will install it.
152154
* This will also create a progress message in vscode in order to inform the user that sqlmesh enterprise is being installed.
155+
* Uses a lock mechanism to prevent parallel executions.
153156
*
154157
* @returns A Result indicating whether sqlmesh enterprise was installed in the call.
155158
*/
156159
export const ensureSqlmeshEnterpriseInstalled = async (): Promise<
157160
Result<boolean, ErrorType>
158161
> => {
159-
traceInfo('Ensuring sqlmesh enterprise is installed')
160-
const isInstalled = await isSqlmeshEnterpriseInstalled()
161-
if (isErr(isInstalled)) {
162-
return isInstalled
163-
}
164-
if (isInstalled.value) {
165-
traceInfo('Sqlmesh enterprise is installed')
166-
return ok(false)
162+
// If there's an ongoing installation, wait for it to complete
163+
if (installationLock) {
164+
return installationLock
167165
}
168-
traceInfo('Sqlmesh enterprise is not installed, installing...')
169-
const abortController = new AbortController()
170-
const installResult = await window.withProgress(
171-
{
172-
location: ProgressLocation.Notification,
173-
title: 'Installing sqlmesh enterprise...',
174-
cancellable: true,
175-
},
176-
async (progress, token) => {
177-
// Connect the cancellation token to our abort controller
178-
token.onCancellationRequested(() => {
179-
abortController.abort()
180-
traceInfo('Sqlmesh enterprise installation cancelled')
181-
window.showInformationMessage('Installation cancelled')
182-
})
183-
progress.report({ message: 'Installing sqlmesh enterprise...' })
184-
const result = await installSqlmeshEnterprise(abortController)
185-
if (isErr(result)) {
186-
return result
166+
167+
// Create a new lock
168+
installationLock = (async () => {
169+
try {
170+
traceInfo('Ensuring sqlmesh enterprise is installed')
171+
const isInstalled = await isSqlmeshEnterpriseInstalled()
172+
if (isErr(isInstalled)) {
173+
return isInstalled
174+
}
175+
if (isInstalled.value) {
176+
traceInfo('Sqlmesh enterprise is installed')
177+
return ok(false)
178+
}
179+
traceInfo('Sqlmesh enterprise is not installed, installing...')
180+
const abortController = new AbortController()
181+
const installResult = await window.withProgress(
182+
{
183+
location: ProgressLocation.Notification,
184+
title: 'SQLMesh',
185+
cancellable: true,
186+
},
187+
async (progress, token) => {
188+
// Connect the cancellation token to our abort controller
189+
token.onCancellationRequested(() => {
190+
abortController.abort()
191+
traceInfo('Sqlmesh enterprise installation cancelled')
192+
window.showInformationMessage('Installation cancelled')
193+
})
194+
progress.report({ message: 'Installing enterprise python package...' })
195+
const result = await installSqlmeshEnterprise(abortController)
196+
if (isErr(result)) {
197+
return result
198+
}
199+
return ok(true)
200+
},
201+
)
202+
if (isErr(installResult)) {
203+
return installResult
187204
}
188205
return ok(true)
189-
},
190-
)
191-
if (isErr(installResult)) {
192-
return installResult
193-
}
194-
return ok(true)
206+
} finally {
207+
// Clear the lock when done
208+
installationLock = undefined
209+
}
210+
})()
211+
212+
return installationLock
195213
}
196214

197215
/**
Lines changed: 215 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,215 @@
1+
import { expect, test } from '@playwright/test'
2+
import path from 'path'
3+
import fs from 'fs-extra'
4+
import os from 'os'
5+
import { exec } from 'child_process'
6+
import { promisify } from 'util'
7+
import { REPO_ROOT, startVSCode, SUSHI_SOURCE_PATH } from './utils'
8+
9+
const execAsync = promisify(exec)
10+
11+
/**
12+
* Helper function to create and set up a Python virtual environment
13+
*/
14+
async function setupPythonEnvironment(envDir: string): Promise<string> {
15+
// Create virtual environment
16+
const pythonCmd = process.platform === 'win32' ? 'python' : 'python3'
17+
const { stderr } = await execAsync(`${pythonCmd} -m venv "${envDir}"`)
18+
if (stderr && !stderr.includes('WARNING')) {
19+
throw new Error(`Failed to create venv: ${stderr}`)
20+
}
21+
22+
// Get paths
23+
const isWindows = process.platform === 'win32'
24+
const binDir = path.join(envDir, isWindows ? 'Scripts' : 'bin')
25+
const pythonPath = path.join(binDir, isWindows ? 'python.exe' : 'python')
26+
const pipPath = path.join(binDir, isWindows ? 'pip.exe' : 'pip')
27+
28+
// Install the mock tcloud package
29+
const mockTcloudPath = path.join(__dirname, 'tcloud')
30+
const { stderr: pipErr1 } = await execAsync(
31+
`"${pipPath}" install -e "${mockTcloudPath}"`,
32+
)
33+
if (pipErr1 && !pipErr1.includes('WARNING') && !pipErr1.includes('notice')) {
34+
throw new Error(`Failed to install mock tcloud: ${pipErr1}`)
35+
}
36+
37+
// Install sqlmesh from the local repository with LSP support
38+
const sqlmeshRepoPath = path.join(__dirname, '..', '..', '..') // Navigate to repo root from tests dir
39+
const { stderr: pipErr2 } = await execAsync(
40+
`"${pipPath}" install -e "${sqlmeshRepoPath}[lsp,bigquery]" "${REPO_ROOT}/examples/custom_materializations"`,
41+
)
42+
if (pipErr2 && !pipErr2.includes('WARNING') && !pipErr2.includes('notice')) {
43+
throw new Error(`Failed to install sqlmesh: ${pipErr2}`)
44+
}
45+
46+
return pythonPath
47+
}
48+
49+
/**
50+
* Helper function to set up a pre-authenticated tcloud state
51+
*/
52+
async function setupAuthenticatedState(tempDir: string): Promise<void> {
53+
const authStateFile = path.join(tempDir, '.tcloud_auth_state.json')
54+
const authState = {
55+
is_logged_in: true,
56+
id_token: {
57+
iss: 'https://mock.tobikodata.com',
58+
aud: 'mock-audience',
59+
sub: 'user-123',
60+
scope: 'openid email profile',
61+
iat: Math.floor(Date.now() / 1000),
62+
exp: Math.floor(Date.now() / 1000) + 3600, // Valid for 1 hour
63+
email: 'test@example.com',
64+
name: 'Test User',
65+
},
66+
}
67+
await fs.writeJson(authStateFile, authState)
68+
}
69+
70+
test.describe('Tcloud', () => {
71+
test('not signed in, shows sign in window', async ({}, testInfo) => {
72+
testInfo.setTimeout(120_000) // 2 minutes for venv creation and package installation
73+
const tempDir = await fs.mkdtemp(
74+
path.join(os.tmpdir(), 'vscode-test-tcloud-'),
75+
)
76+
const pythonEnvDir = path.join(tempDir, '.venv')
77+
78+
try {
79+
// Copy sushi project
80+
await fs.copy(SUSHI_SOURCE_PATH, tempDir)
81+
82+
// Create a tcloud.yaml to mark this as a tcloud project
83+
const tcloudConfig = {
84+
url: 'https://mock.tobikodata.com',
85+
org: 'test-org',
86+
project: 'test-project',
87+
}
88+
await fs.writeFile(
89+
path.join(tempDir, 'tcloud.yaml'),
90+
`url: ${tcloudConfig.url}\norg: ${tcloudConfig.org}\nproject: ${tcloudConfig.project}\n`,
91+
)
92+
93+
// Set up Python environment with mock tcloud and sqlmesh
94+
const pythonPath = await setupPythonEnvironment(pythonEnvDir)
95+
96+
// Configure VS Code settings to use our Python environment
97+
const settings = {
98+
'python.defaultInterpreterPath': pythonPath,
99+
'sqlmesh.environmentPath': pythonEnvDir,
100+
}
101+
await fs.ensureDir(path.join(tempDir, '.vscode'))
102+
await fs.writeJson(
103+
path.join(tempDir, '.vscode', 'settings.json'),
104+
settings,
105+
{ spaces: 2 },
106+
)
107+
108+
// Start VS Code
109+
const { window, close } = await startVSCode(tempDir)
110+
111+
// Open a SQL file to trigger SQLMesh activation
112+
// Wait for the models folder to be visible
113+
await window.waitForSelector('text=models')
114+
115+
// Click on the models folder
116+
await window
117+
.getByRole('treeitem', { name: 'models', exact: true })
118+
.locator('a')
119+
.click()
120+
121+
// Open the top_waiters model
122+
await window
123+
.getByRole('treeitem', { name: 'customers.sql', exact: true })
124+
.locator('a')
125+
.click()
126+
127+
// Wait for the file to open
128+
await window.waitForTimeout(2000)
129+
130+
await window.waitForSelector(
131+
'text=Please sign in to Tobiko Cloud to use SQLMesh',
132+
)
133+
134+
// Close VS Code
135+
await close()
136+
} finally {
137+
// Clean up
138+
await fs.remove(tempDir)
139+
}
140+
})
141+
142+
test('signed in and not installed shows installation window and can see output', async ({}, testInfo) => {
143+
testInfo.setTimeout(120_000) // 2 minutes for venv creation and package installation
144+
const tempDir = await fs.mkdtemp(
145+
path.join(os.tmpdir(), 'vscode-test-tcloud-'),
146+
)
147+
const pythonEnvDir = path.join(tempDir, '.venv')
148+
149+
try {
150+
// Copy sushi project
151+
await fs.copy(SUSHI_SOURCE_PATH, tempDir)
152+
153+
// Create a tcloud.yaml to mark this as a tcloud project
154+
const tcloudConfig = {
155+
url: 'https://mock.tobikodata.com',
156+
org: 'test-org',
157+
project: 'test-project',
158+
}
159+
await fs.writeFile(
160+
path.join(tempDir, 'tcloud.yaml'),
161+
`url: ${tcloudConfig.url}\norg: ${tcloudConfig.org}\nproject: ${tcloudConfig.project}\n`,
162+
)
163+
164+
// Write mock ".tcloud_auth_state.json" file
165+
await setupAuthenticatedState(tempDir)
166+
167+
// Set up Python environment with mock tcloud and sqlmesh
168+
const pythonPath = await setupPythonEnvironment(pythonEnvDir)
169+
170+
// Configure VS Code settings to use our Python environment
171+
const settings = {
172+
'python.defaultInterpreterPath': pythonPath,
173+
'sqlmesh.environmentPath': pythonEnvDir,
174+
}
175+
await fs.ensureDir(path.join(tempDir, '.vscode'))
176+
await fs.writeJson(
177+
path.join(tempDir, '.vscode', 'settings.json'),
178+
settings,
179+
{ spaces: 2 },
180+
)
181+
182+
// Start VS Code
183+
const { window, close } = await startVSCode(tempDir)
184+
185+
// Open a SQL file to trigger SQLMesh activation
186+
// Wait for the models folder to be visible
187+
await window.waitForSelector('text=models')
188+
189+
// Click on the models folder
190+
await window
191+
.getByRole('treeitem', { name: 'models', exact: true })
192+
.locator('a')
193+
.click()
194+
195+
// Open the top_waiters model
196+
await window
197+
.getByRole('treeitem', { name: 'customers.sql', exact: true })
198+
.locator('a')
199+
.click()
200+
201+
await window.waitForSelector('text=Installing enterprise python package')
202+
expect(
203+
await window.locator('text=Installing enterprise python package'),
204+
).toHaveCount(2)
205+
206+
await window.waitForSelector('text=Loaded SQLMesh context')
207+
208+
// Close VS Code
209+
await close()
210+
} finally {
211+
// Clean up
212+
await fs.remove(tempDir)
213+
}
214+
})
215+
})

0 commit comments

Comments
 (0)