From c8b0b7a3115a487283a96c7ea1692a52f1918a44 Mon Sep 17 00:00:00 2001 From: Ben <9087625+benfdking@users.noreply.github.com> Date: Wed, 11 Jun 2025 21:53:29 +0100 Subject: [PATCH] feat: add plan diff lsp endpoint and vscode command --- sqlmesh/lsp/custom.py | 27 ++++++++++ sqlmesh/lsp/main.py | 40 ++++++++++++++ vscode/extension/package.json | 5 ++ vscode/extension/src/commands/planDiff.ts | 66 +++++++++++++++++++++++ vscode/extension/src/extension.ts | 8 +++ vscode/extension/src/lsp/custom.ts | 34 ++++++++++++ 6 files changed, 180 insertions(+) create mode 100644 vscode/extension/src/commands/planDiff.ts diff --git a/sqlmesh/lsp/custom.py b/sqlmesh/lsp/custom.py index 618b4a44bc..995112ab83 100644 --- a/sqlmesh/lsp/custom.py +++ b/sqlmesh/lsp/custom.py @@ -143,3 +143,30 @@ class FormatProjectResponse(CustomMethodResponseBaseClass): """ pass + + +LIST_ENVIRONMENTS_FEATURE = "sqlmesh/list_environments" + + +class ListEnvironmentsRequest(CustomMethodRequestBaseClass): + pass + + +class ListEnvironmentsResponse(CustomMethodResponseBaseClass): + environments: t.List[str] + + +PLAN_DIFF_FEATURE = "sqlmesh/plan_diff" + + +class PlanDiffRequest(CustomMethodRequestBaseClass): + environment: str + + +class PlanDiffEntry(PydanticModel): + name: str + diff: str + + +class PlanDiffResponse(CustomMethodResponseBaseClass): + diffs: t.List[PlanDiffEntry] diff --git a/sqlmesh/lsp/main.py b/sqlmesh/lsp/main.py index f9bfe46114..7bc8aecdf2 100755 --- a/sqlmesh/lsp/main.py +++ b/sqlmesh/lsp/main.py @@ -31,6 +31,8 @@ RENDER_MODEL_FEATURE, SUPPORTED_METHODS_FEATURE, FORMAT_PROJECT_FEATURE, + LIST_ENVIRONMENTS_FEATURE, + PLAN_DIFF_FEATURE, AllModelsRequest, AllModelsResponse, AllModelsForRenderRequest, @@ -43,6 +45,11 @@ FormatProjectRequest, FormatProjectResponse, CustomMethod, + ListEnvironmentsRequest, + ListEnvironmentsResponse, + PlanDiffRequest, + PlanDiffResponse, + PlanDiffEntry, ) from sqlmesh.lsp.hints import get_hints from sqlmesh.lsp.reference import ( @@ -89,6 +96,8 @@ def __init__( API_FEATURE: self._custom_api, SUPPORTED_METHODS_FEATURE: self._custom_supported_methods, FORMAT_PROJECT_FEATURE: self._custom_format_project, + LIST_ENVIRONMENTS_FEATURE: self._custom_list_environments, + PLAN_DIFF_FEATURE: self._custom_plan_diff, } # Register LSP features (e.g., formatting, hover, etc.) @@ -211,6 +220,37 @@ def _custom_api( raise NotImplementedError(f"API request not implemented: {request.url}") + def _custom_list_environments( + self, ls: LanguageServer, params: ListEnvironmentsRequest + ) -> ListEnvironmentsResponse: + if self.lsp_context is None: + current_path = Path.cwd() + self._ensure_context_in_folder(current_path) + if self.lsp_context is None: + raise RuntimeError("No context found") + + envs = self.lsp_context.context._new_state_sync().get_environments_summary() + return ListEnvironmentsResponse(environments=[e.name for e in envs]) + + def _custom_plan_diff(self, ls: LanguageServer, params: PlanDiffRequest) -> PlanDiffResponse: + if self.lsp_context is None: + current_path = Path.cwd() + self._ensure_context_in_folder(current_path) + if self.lsp_context is None: + raise RuntimeError("No context found") + + plan = self.lsp_context.context.plan_builder( + params.environment, + skip_tests=True, + skip_backfill=True, + ).build() + + diffs = [ + PlanDiffEntry(name=name, diff=plan.context_diff.text_diff(name)) + for name in plan.context_diff.modified_snapshots + ] + return PlanDiffResponse(diffs=diffs) + def _custom_supported_methods( self, ls: LanguageServer, params: SupportedMethodsRequest ) -> SupportedMethodsResponse: diff --git a/vscode/extension/package.json b/vscode/extension/package.json index 8e36c459c1..431e3f7c9d 100644 --- a/vscode/extension/package.json +++ b/vscode/extension/package.json @@ -93,6 +93,11 @@ "title": "Render Model", "description": "Render the model in the current editor", "icon": "$(open-preview)" + }, + { + "command": "sqlmesh.plan", + "title": "Show Plan Diff", + "description": "Create a plan and show the diff" } ], "menus": { diff --git a/vscode/extension/src/commands/planDiff.ts b/vscode/extension/src/commands/planDiff.ts new file mode 100644 index 0000000000..a1ecb98630 --- /dev/null +++ b/vscode/extension/src/commands/planDiff.ts @@ -0,0 +1,66 @@ +import * as vscode from 'vscode' +import { LSPClient } from '../lsp/lsp' +import { isErr } from '@bus/result' + +export function planDiff(lspClient?: LSPClient) { + return async () => { + if (!lspClient) { + vscode.window.showErrorMessage('LSP client not available') + return + } + + const envResult = await lspClient.call_custom_method( + 'sqlmesh/list_environments', + {}, + ) + + if (isErr(envResult)) { + vscode.window.showErrorMessage( + `Failed to list environments: ${envResult.error.message}`, + ) + return + } + + const env = await vscode.window.showQuickPick(envResult.value.environments, { + placeHolder: 'Select environment to plan', + }) + + if (!env) { + return + } + + const diffResult = await lspClient.call_custom_method('sqlmesh/plan_diff', { + environment: env, + }) + + if (isErr(diffResult)) { + vscode.window.showErrorMessage( + `Failed to get plan diff: ${diffResult.error.message}`, + ) + return + } + + if (!diffResult.value.diffs.length) { + vscode.window.showInformationMessage('No changes detected') + return + } + + let selected = diffResult.value.diffs[0] + if (diffResult.value.diffs.length > 1) { + const pick = await vscode.window.showQuickPick( + diffResult.value.diffs.map(d => ({ label: d.name })), + { placeHolder: 'Select a model diff to view' }, + ) + if (!pick) { + return + } + selected = diffResult.value.diffs.find(d => d.name === pick.label)! + } + + const doc = await vscode.workspace.openTextDocument({ + content: selected.diff, + language: 'diff', + }) + await vscode.window.showTextDocument(doc, { preview: true }) + } +} diff --git a/vscode/extension/src/extension.ts b/vscode/extension/src/extension.ts index 8749c61fb2..0e593a4179 100644 --- a/vscode/extension/src/extension.ts +++ b/vscode/extension/src/extension.ts @@ -13,6 +13,7 @@ import { signOut } from './commands/signout' import { signIn } from './commands/signin' import { signInSpecifyFlow } from './commands/signinSpecifyFlow' import { renderModel, reRenderModelForSourceFile } from './commands/renderModel' +import { planDiff } from './commands/planDiff' import { isErr } from '@bus/result' import { handleError } from './utilities/errors' import { selector, completionProvider } from './completion/completion' @@ -83,6 +84,13 @@ export async function activate(context: vscode.ExtensionContext) { ), ) + context.subscriptions.push( + vscode.commands.registerCommand( + 'sqlmesh.plan', + planDiff(lspClient), + ), + ) + // Register the webview const lineagePanel = new LineagePanel(context.extensionUri, lspClient) context.subscriptions.push( diff --git a/vscode/extension/src/lsp/custom.ts b/vscode/extension/src/lsp/custom.ts index be11419b79..50cf1ff515 100644 --- a/vscode/extension/src/lsp/custom.ts +++ b/vscode/extension/src/lsp/custom.ts @@ -33,6 +33,8 @@ export type CustomLSPMethods = | AllModelsForRenderMethod | SupportedMethodsMethod | FormatProjectMethod + | ListEnvironmentsMethod + | PlanDiffMethod interface AllModelsRequest { textDocument: { @@ -106,3 +108,35 @@ interface FormatProjectRequest {} // eslint-disable-next-line @typescript-eslint/no-empty-object-type interface FormatProjectResponse {} + +export interface ListEnvironmentsMethod { + method: 'sqlmesh/list_environments' + request: ListEnvironmentsRequest + response: ListEnvironmentsResponse +} + +// eslint-disable-next-line @typescript-eslint/no-empty-interface +interface ListEnvironmentsRequest {} + +interface ListEnvironmentsResponse { + environments: string[] +} + +export interface PlanDiffMethod { + method: 'sqlmesh/plan_diff' + request: PlanDiffRequest + response: PlanDiffResponse +} + +interface PlanDiffRequest { + environment: string +} + +interface PlanDiffEntry { + name: string + diff: string +} + +interface PlanDiffResponse { + diffs: PlanDiffEntry[] +}