feat: add plan diff lsp endpoint and vscode command

Author: benfdking

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]

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:

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": {

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 })
+  }
+}

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(

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[]
+}