diff --git a/sqlmesh/lsp/main.py b/sqlmesh/lsp/main.py index f9bfe46114..641b570c84 100755 --- a/sqlmesh/lsp/main.py +++ b/sqlmesh/lsp/main.py @@ -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) @@ -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: 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: diff --git a/tooling/vscode/extensions.json b/tooling/vscode/extensions.json index 0271570408..b703cc6e84 100644 --- a/tooling/vscode/extensions.json +++ b/tooling/vscode/extensions.json @@ -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" ] } diff --git a/vscode/extension/src/utilities/sqlmesh/sqlmesh.ts b/vscode/extension/src/utilities/sqlmesh/sqlmesh.ts index 95cc94d38e..45d9cfbd4c 100644 --- a/vscode/extension/src/utilities/sqlmesh/sqlmesh.ts +++ b/vscode/extension/src/utilities/sqlmesh/sqlmesh.ts @@ -147,51 +147,69 @@ export const installSqlmeshEnterprise = async ( return ok(true) } +let installationLock: Promise> | 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 > => { - 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 } /** diff --git a/vscode/extension/tests/tcloud.spec.ts b/vscode/extension/tests/tcloud.spec.ts new file mode 100644 index 0000000000..3a99288546 --- /dev/null +++ b/vscode/extension/tests/tcloud.spec.ts @@ -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 { + // 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 { + 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) + } + }) +}) diff --git a/vscode/extension/tests/tcloud/README.md b/vscode/extension/tests/tcloud/README.md new file mode 100644 index 0000000000..c3c723be37 --- /dev/null +++ b/vscode/extension/tests/tcloud/README.md @@ -0,0 +1,53 @@ +# Mock tcloud CLI for Testing + +This directory contains a mock implementation of the tcloud CLI for testing the VSCode extension. + +## Implemented Commands + +The mock implements only the commands used by the VSCode extension: + +### Authentication Commands +- `tcloud auth vscode status` - Returns authentication status +- `tcloud auth vscode login-url` - Returns mock OAuth login URL +- `tcloud auth vscode start-server ` - Simulates OAuth callback +- `tcloud auth vscode device` - Returns mock device flow info +- `tcloud auth vscode poll_device ` - Simulates device flow success +- `tcloud auth logout` - Clears authentication state + +### SQLMesh Commands +- `tcloud is_sqlmesh_installed` - Checks installation status +- `tcloud install_sqlmesh` - Marks SQLMesh as installed +- `tcloud sqlmesh ` - Echoes sqlmesh commands + +## State Management + +The mock maintains state in two files: +- `.tcloud_auth_state.json` - Authentication state (logged in/out, ID token) +- `.sqlmesh_installed` - SQLMesh installation marker + +## Usage in Tests + +To use this mock in tests: + +1. Ensure the mock is in PATH or reference it directly +2. The mock will simulate successful authentication flows +3. State persists between calls for realistic testing + +## Example + +```bash +# Check auth status +./tcloud auth vscode status +# Output: {"is_logged_in": false, "id_token": null} + +# Simulate login +./tcloud auth vscode login-url +# Output: {"url": "https://mock-auth.example.com/auth?client_id=mock&redirect_uri=http://localhost:7890", "verifier_code": "mock_verifier_12345"} + +./tcloud auth vscode start-server mock_verifier_12345 +# Output: Mock server started successfully + +# Check status again +./tcloud auth vscode status +# Output: {"is_logged_in": true, "id_token": {"email": "test@example.com", "name": "Test User", "exp": 1736790123}} +``` \ No newline at end of file diff --git a/vscode/extension/tests/tcloud/mock_tcloud/__init__.py b/vscode/extension/tests/tcloud/mock_tcloud/__init__.py new file mode 100644 index 0000000000..98ad152e0e --- /dev/null +++ b/vscode/extension/tests/tcloud/mock_tcloud/__init__.py @@ -0,0 +1 @@ +# Mock tcloud package \ No newline at end of file diff --git a/vscode/extension/tests/tcloud/mock_tcloud/cli.py b/vscode/extension/tests/tcloud/mock_tcloud/cli.py new file mode 100755 index 0000000000..b9b80ab8f1 --- /dev/null +++ b/vscode/extension/tests/tcloud/mock_tcloud/cli.py @@ -0,0 +1,257 @@ +""" +Mock tcloud CLI for testing VSCode extension. +Implements only the commands used by the extension. +""" + +import json +import os +import subprocess +import sys +import time +from pathlib import Path + +import click + +def get_auth_state_file(): + """Get the path to the auth state file in the current working directory""" + return Path.cwd() / ".tcloud_auth_state.json" + + +def load_auth_state(): + """Load authentication state from file""" + auth_file = get_auth_state_file() + if auth_file.exists(): + with open(auth_file, "r") as f: + return json.load(f) + return {"is_logged_in": False, "id_token": None} + + +def save_auth_state(state): + """Save authentication state to file""" + auth_file = get_auth_state_file() + with open(auth_file, "w") as f: + json.dump(state, f) + + +@click.group(no_args_is_help=True) +@click.option( + "--project", + type=str, + help="The name of the project.", +) +@click.pass_context +def cli(ctx: click.Context, project: str) -> None: + """Mock Tobiko Cloud CLI""" + ctx.ensure_object(dict) + ctx.obj["project"] = project + + +@cli.command("is_sqlmesh_installed", hidden=True) +@click.pass_context +def is_sqlmesh_installed(ctx: click.Context) -> None: + """Check if SQLMesh Enterprise is installed""" + # For testing, we'll track installation state in a file in the current bin directory + # This matches where the test expects it to be + bin_dir = Path(sys.executable).parent + install_state_file = bin_dir / ".sqlmesh_installed" + is_installed = install_state_file.exists() + + print( + json.dumps( + { + "is_installed": is_installed, + } + ) + ) + + +@cli.command("install_sqlmesh") +@click.pass_context +def install_sqlmesh(ctx: click.Context) -> None: + """Install the correct version of SQLMesh Enterprise""" + + # For 3 seconds output to stdout + for i in range(3): + print(f"Installing SQLMesh Enterprise logs {i + 1}/3", flush=True) + time.sleep(1) + + # Simulate installation by creating a marker file in the bin directory + bin_dir = Path(sys.executable).parent + install_state_file = bin_dir / ".sqlmesh_installed" + install_state_file.touch() + + print("Mock SQLMesh Enterprise installed successfully") + + +@cli.command("sqlmesh") +@click.argument("args", nargs=-1) +@click.pass_context +def sqlmesh(ctx: click.Context, args) -> None: + """Run SQLMesh Enterprise commands""" + # Pass through to the real sqlmesh command + + # Get the path to sqlmesh in the same environment as this script + bin_dir = os.path.dirname(sys.executable) + sqlmesh_path = os.path.join(bin_dir, "sqlmesh") + + if not os.path.exists(sqlmesh_path): + # Try with .exe extension on Windows + sqlmesh_path = os.path.join(bin_dir, "sqlmesh.exe") + + if not os.path.exists(sqlmesh_path): + # Fall back to using sqlmesh from PATH + sqlmesh_path = "sqlmesh" + + # Execute the real sqlmesh with the provided arguments + result = subprocess.run([sqlmesh_path] + list(args), capture_output=False) + sys.exit(result.returncode) + + +@click.group() +def auth() -> None: + """ + Tobiko Cloud Authentication + """ + + +@auth.command() +def logout() -> None: + """Logout of any current session""" + save_auth_state({"is_logged_in": False, "id_token": None}) + print("Logged out successfully") + + +### Methods for VSCode +@auth.group(hidden=True) +def vscode() -> None: + """Commands for VSCode integration""" + pass + + +@vscode.command("login-url") +def login_url() -> None: + """ + Login to Tobiko Cloud. + + This returns a JSON object with the following fields: + - url: The URL to login open + """ + # Return mock OAuth URL and verifier + print( + json.dumps( + { + "url": "https://mock-auth.example.com/auth?client_id=mock&redirect_uri=http://localhost:7890", + "verifier_code": "mock_verifier_12345", + } + ) + ) + + +@vscode.command("start-server") +@click.argument("code_verifier", type=str, required=True) +def start_server(code_verifier: str) -> None: + """ + Start the server to catch the redirect from the browser. + """ + # Simulate successful authentication after a short delay + time.sleep(0.5) + + # Update auth state to logged in + save_auth_state( + { + "is_logged_in": True, + "id_token": { + "iss": "https://mock.tobikodata.com", + "aud": "mock-audience", + "sub": "user-123", + "scope": "openid email profile", + "iat": int(time.time()), + "exp": int(time.time()) + 3600, # Token expires in 1 hour + "email": "test@example.com", + "name": "Test User", + }, + } + ) + + # The real command would start a server, but for testing we just simulate success + print("Mock server started successfully") + + +@vscode.command("status") +def vscode_status() -> None: + """ + Auth status for logged in + """ + state = load_auth_state() + print( + json.dumps( + {"is_logged_in": state["is_logged_in"], "id_token": state["id_token"]} + ) + ) + + +@vscode.command("device") +def vscode_device() -> None: + """ + Initiate device flow for VSCode integration + """ + print( + json.dumps( + { + "device_code": "MOCK-DEVICE-CODE", + "user_code": "ABCD-1234", + "verification_uri": "https://mock-auth.example.com/device", + "verification_uri_complete": "https://mock-auth.example.com/device?user_code=ABCD-1234", + "expires_in": 600, + "interval": 5, + } + ) + ) + + +@vscode.command("poll_device") +@click.argument("device_code", type=str, required=True) +@click.option( + "-i", + "--interval", + type=int, + default=5, + help="The interval between polling attempts in seconds", +) +@click.option( + "-t", + "--timeout", + type=int, + default=300, + help="The timeout for the device flow in seconds", +) +def vscode_poll_device(device_code: str, interval: int, timeout: int) -> None: + """ + Poll the device flow for VSCode integration + """ + # For testing, we'll just succeed immediately + save_auth_state( + { + "is_logged_in": True, + "id_token": { + "iss": "https://mock.tobikodata.com", + "aud": "mock-audience", + "sub": "device-user-123", + "scope": "openid email profile", + "iat": int(time.time()), + "exp": int(time.time()) + 3600, + "email": "device@example.com", + "name": "Device User", + }, + } + ) + + print(json.dumps({"success": True})) + + +# Add auth group to main CLI +cli.add_command(auth) + + +if __name__ == "__main__": + cli() diff --git a/vscode/extension/tests/tcloud/pyproject.toml b/vscode/extension/tests/tcloud/pyproject.toml new file mode 100644 index 0000000000..505af5c783 --- /dev/null +++ b/vscode/extension/tests/tcloud/pyproject.toml @@ -0,0 +1,18 @@ +[build-system] +requires = ["setuptools>=61.0", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "mock-tcloud" +version = "0.1.0" +description = "Mock tcloud CLI for testing VSCode extension" +requires-python = ">=3.8" +dependencies = [ + "click>=8.0", +] + +[project.scripts] +tcloud = "mock_tcloud.cli:cli" + +[tool.setuptools] +packages = ["mock_tcloud"] \ No newline at end of file diff --git a/vscode/extension/tests/utils.ts b/vscode/extension/tests/utils.ts index 1a2c55e1e6..ff68c7291f 100644 --- a/vscode/extension/tests/utils.ts +++ b/vscode/extension/tests/utils.ts @@ -18,6 +18,7 @@ export const SUSHI_SOURCE_PATH = path.join( 'examples', 'sushi', ) +export const REPO_ROOT = path.join(__dirname, '..', '..', '..') /** * Launch VS Code and return the window and a function to close the app.