Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
38 changes: 8 additions & 30 deletions sqlmesh/lsp/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -294,32 +294,6 @@ def did_open(ls: LanguageServer, params: types.DidOpenTextDocumentParams) -> Non
SQLMeshLanguageServer._diagnostics_to_lsp_diagnostics(diagnostics),
)

@self.server.feature(types.TEXT_DOCUMENT_DID_CHANGE)
def did_change(ls: LanguageServer, params: types.DidChangeTextDocumentParams) -> None:
if self.lsp_context is None:
current_path = Path.cwd()
self._ensure_context_in_folder(current_path)
if self.lsp_context is None:
ls.log_trace("No context found after did_change")
return

uri = URI(params.text_document.uri)
context = self._context_get_or_load(uri)

models = context.map[uri.to_path()]
if models is None or not isinstance(models, ModelTarget):
return

# Get diagnostics from context (which handles caching)
diagnostics = context.lint_model(uri)

# Only publish diagnostics if client doesn't support pull diagnostics
if not self.client_supports_pull_diagnostics:
ls.publish_diagnostics(
params.text_document.uri,
SQLMeshLanguageServer._diagnostics_to_lsp_diagnostics(diagnostics),
)

@self.server.feature(types.TEXT_DOCUMENT_DID_SAVE)
def did_save(ls: LanguageServer, params: types.DidSaveTextDocumentParams) -> None:
uri = URI(params.text_document.uri)
Expand Down Expand Up @@ -753,15 +727,19 @@ def _ensure_context_for_document(

loaded = False
# Ascend directories to look for config
while path.parents and not loaded:
current_dir = path.parent # Start from the file's parent directory
while current_dir.parents and not loaded:
Comment thread
benfdking marked this conversation as resolved.
for ext in ("py", "yml", "yaml"):
config_path = path / f"config.{ext}"
config_path = current_dir / f"config.{ext}"
if config_path.exists():
if self._create_lsp_context([path]):
if self._create_lsp_context([current_dir]):
loaded = True
# Re-check context for the document now that it's loaded
return self._ensure_context_for_document(document_uri)
path = path.parent
# Check if we've reached the filesystem root to prevent infinite loops
if current_dir == current_dir.parent:
break
current_dir = current_dir.parent

# If still no context found, try the workspace folders
if not loaded:
Expand Down
3 changes: 2 additions & 1 deletion tooling/vscode/extensions.json
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
"dbaeumer.vscode-eslint",
"amodio.tsl-problem-matcher",
"ms-vscode.extension-test-runner",
"ms-playwright.playwright"
"ms-playwright.playwright",
"esbenp.prettier-vscode"
]
}
84 changes: 51 additions & 33 deletions vscode/extension/src/utilities/sqlmesh/sqlmesh.ts
Original file line number Diff line number Diff line change
Expand Up @@ -147,51 +147,69 @@ export const installSqlmeshEnterprise = async (
return ok(true)
}

let installationLock: Promise<Result<boolean, ErrorType>> | undefined = undefined

/**
* Checks if sqlmesh enterprise is installed and updated. If not, it will install it.
* This will also create a progress message in vscode in order to inform the user that sqlmesh enterprise is being installed.
* Uses a lock mechanism to prevent parallel executions.
*
* @returns A Result indicating whether sqlmesh enterprise was installed in the call.
*/
export const ensureSqlmeshEnterpriseInstalled = async (): Promise<
Result<boolean, ErrorType>
> => {
traceInfo('Ensuring sqlmesh enterprise is installed')
const isInstalled = await isSqlmeshEnterpriseInstalled()
if (isErr(isInstalled)) {
return isInstalled
}
if (isInstalled.value) {
traceInfo('Sqlmesh enterprise is installed')
return ok(false)
// If there's an ongoing installation, wait for it to complete
if (installationLock) {
return installationLock
}
traceInfo('Sqlmesh enterprise is not installed, installing...')
const abortController = new AbortController()
const installResult = await window.withProgress(
{
location: ProgressLocation.Notification,
title: 'Installing sqlmesh enterprise...',
cancellable: true,
},
async (progress, token) => {
// Connect the cancellation token to our abort controller
token.onCancellationRequested(() => {
abortController.abort()
traceInfo('Sqlmesh enterprise installation cancelled')
window.showInformationMessage('Installation cancelled')
})
progress.report({ message: 'Installing sqlmesh enterprise...' })
const result = await installSqlmeshEnterprise(abortController)
if (isErr(result)) {
return result

// Create a new lock
installationLock = (async () => {
try {
traceInfo('Ensuring sqlmesh enterprise is installed')
const isInstalled = await isSqlmeshEnterpriseInstalled()
if (isErr(isInstalled)) {
return isInstalled
}
if (isInstalled.value) {
traceInfo('Sqlmesh enterprise is installed')
return ok(false)
}
traceInfo('Sqlmesh enterprise is not installed, installing...')
const abortController = new AbortController()
const installResult = await window.withProgress(
{
location: ProgressLocation.Notification,
title: 'SQLMesh',
cancellable: true,
},
async (progress, token) => {
// Connect the cancellation token to our abort controller
token.onCancellationRequested(() => {
abortController.abort()
traceInfo('Sqlmesh enterprise installation cancelled')
window.showInformationMessage('Installation cancelled')
})
progress.report({ message: 'Installing enterprise python package...' })
const result = await installSqlmeshEnterprise(abortController)
if (isErr(result)) {
return result
}
return ok(true)
},
)
if (isErr(installResult)) {
return installResult
}
return ok(true)
},
)
if (isErr(installResult)) {
return installResult
}
return ok(true)
} finally {
// Clear the lock when done
installationLock = undefined
}
})()

return installationLock
}

/**
Expand Down
215 changes: 215 additions & 0 deletions vscode/extension/tests/tcloud.spec.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,215 @@
import { expect, test } from '@playwright/test'
import path from 'path'
import fs from 'fs-extra'
import os from 'os'
import { exec } from 'child_process'
import { promisify } from 'util'
import { REPO_ROOT, startVSCode, SUSHI_SOURCE_PATH } from './utils'

const execAsync = promisify(exec)

/**
* Helper function to create and set up a Python virtual environment
*/
async function setupPythonEnvironment(envDir: string): Promise<string> {
// Create virtual environment
const pythonCmd = process.platform === 'win32' ? 'python' : 'python3'
const { stderr } = await execAsync(`${pythonCmd} -m venv "${envDir}"`)
if (stderr && !stderr.includes('WARNING')) {
throw new Error(`Failed to create venv: ${stderr}`)
}

// Get paths
const isWindows = process.platform === 'win32'
const binDir = path.join(envDir, isWindows ? 'Scripts' : 'bin')
const pythonPath = path.join(binDir, isWindows ? 'python.exe' : 'python')
const pipPath = path.join(binDir, isWindows ? 'pip.exe' : 'pip')

// Install the mock tcloud package
const mockTcloudPath = path.join(__dirname, 'tcloud')
const { stderr: pipErr1 } = await execAsync(
`"${pipPath}" install -e "${mockTcloudPath}"`,
)
if (pipErr1 && !pipErr1.includes('WARNING') && !pipErr1.includes('notice')) {
throw new Error(`Failed to install mock tcloud: ${pipErr1}`)
}

// Install sqlmesh from the local repository with LSP support
const sqlmeshRepoPath = path.join(__dirname, '..', '..', '..') // Navigate to repo root from tests dir
const { stderr: pipErr2 } = await execAsync(
`"${pipPath}" install -e "${sqlmeshRepoPath}[lsp,bigquery]" "${REPO_ROOT}/examples/custom_materializations"`,
)
if (pipErr2 && !pipErr2.includes('WARNING') && !pipErr2.includes('notice')) {
throw new Error(`Failed to install sqlmesh: ${pipErr2}`)
}

return pythonPath
}

/**
* Helper function to set up a pre-authenticated tcloud state
*/
async function setupAuthenticatedState(tempDir: string): Promise<void> {
const authStateFile = path.join(tempDir, '.tcloud_auth_state.json')
const authState = {
is_logged_in: true,
id_token: {
iss: 'https://mock.tobikodata.com',
aud: 'mock-audience',
sub: 'user-123',
scope: 'openid email profile',
iat: Math.floor(Date.now() / 1000),
exp: Math.floor(Date.now() / 1000) + 3600, // Valid for 1 hour
email: 'test@example.com',
name: 'Test User',
},
}
await fs.writeJson(authStateFile, authState)
}

test.describe('Tcloud', () => {
test('not signed in, shows sign in window', async ({}, testInfo) => {
testInfo.setTimeout(120_000) // 2 minutes for venv creation and package installation
const tempDir = await fs.mkdtemp(
path.join(os.tmpdir(), 'vscode-test-tcloud-'),
)
const pythonEnvDir = path.join(tempDir, '.venv')

try {
// Copy sushi project
await fs.copy(SUSHI_SOURCE_PATH, tempDir)

// Create a tcloud.yaml to mark this as a tcloud project
const tcloudConfig = {
url: 'https://mock.tobikodata.com',
org: 'test-org',
project: 'test-project',
}
await fs.writeFile(
path.join(tempDir, 'tcloud.yaml'),
`url: ${tcloudConfig.url}\norg: ${tcloudConfig.org}\nproject: ${tcloudConfig.project}\n`,
)

// Set up Python environment with mock tcloud and sqlmesh
const pythonPath = await setupPythonEnvironment(pythonEnvDir)

// Configure VS Code settings to use our Python environment
const settings = {
'python.defaultInterpreterPath': pythonPath,
'sqlmesh.environmentPath': pythonEnvDir,
}
await fs.ensureDir(path.join(tempDir, '.vscode'))
await fs.writeJson(
path.join(tempDir, '.vscode', 'settings.json'),
settings,
{ spaces: 2 },
)

// Start VS Code
const { window, close } = await startVSCode(tempDir)

// Open a SQL file to trigger SQLMesh activation
// Wait for the models folder to be visible
await window.waitForSelector('text=models')

// Click on the models folder
await window
.getByRole('treeitem', { name: 'models', exact: true })
.locator('a')
.click()

// Open the top_waiters model
await window
.getByRole('treeitem', { name: 'customers.sql', exact: true })
.locator('a')
.click()

// Wait for the file to open
await window.waitForTimeout(2000)

await window.waitForSelector(
'text=Please sign in to Tobiko Cloud to use SQLMesh',
)

// Close VS Code
await close()
} finally {
// Clean up
await fs.remove(tempDir)
}
})

test('signed in and not installed shows installation window and can see output', async ({}, testInfo) => {
testInfo.setTimeout(120_000) // 2 minutes for venv creation and package installation
const tempDir = await fs.mkdtemp(
path.join(os.tmpdir(), 'vscode-test-tcloud-'),
)
const pythonEnvDir = path.join(tempDir, '.venv')

try {
// Copy sushi project
await fs.copy(SUSHI_SOURCE_PATH, tempDir)

// Create a tcloud.yaml to mark this as a tcloud project
const tcloudConfig = {
url: 'https://mock.tobikodata.com',
org: 'test-org',
project: 'test-project',
}
await fs.writeFile(
path.join(tempDir, 'tcloud.yaml'),
`url: ${tcloudConfig.url}\norg: ${tcloudConfig.org}\nproject: ${tcloudConfig.project}\n`,
)

// Write mock ".tcloud_auth_state.json" file
await setupAuthenticatedState(tempDir)

// Set up Python environment with mock tcloud and sqlmesh
const pythonPath = await setupPythonEnvironment(pythonEnvDir)

// Configure VS Code settings to use our Python environment
const settings = {
'python.defaultInterpreterPath': pythonPath,
'sqlmesh.environmentPath': pythonEnvDir,
}
await fs.ensureDir(path.join(tempDir, '.vscode'))
await fs.writeJson(
path.join(tempDir, '.vscode', 'settings.json'),
settings,
{ spaces: 2 },
)

// Start VS Code
const { window, close } = await startVSCode(tempDir)

// Open a SQL file to trigger SQLMesh activation
// Wait for the models folder to be visible
await window.waitForSelector('text=models')

// Click on the models folder
await window
.getByRole('treeitem', { name: 'models', exact: true })
.locator('a')
.click()

// Open the top_waiters model
await window
.getByRole('treeitem', { name: 'customers.sql', exact: true })
.locator('a')
.click()

await window.waitForSelector('text=Installing enterprise python package')
expect(
await window.locator('text=Installing enterprise python package'),
).toHaveCount(2)

await window.waitForSelector('text=Loaded SQLMesh context')

// Close VS Code
await close()
} finally {
// Clean up
await fs.remove(tempDir)
}
})
})
Loading