feat: add model autocomplete to vscode

[ci skip]

Author: benfdking

examples/sushi/models/latest_order.sql

@@ -11,5 +11,4 @@ MODEL (
 
 SELECT id, customer_id, start_ts, end_ts, event_date
 FROM sushi.orders
-ORDER BY event_date DESC LIMIT 1
-
+ORDER BY event_date DESC LIMIT 1

sqlmesh/lsp/custom.py

@@ -0,0 +1,21 @@
+from lsprotocol import types
+import typing as t
+from sqlmesh.utils.pydantic import PydanticModel
+
+ALL_MODELS_FEATURE = "sqlmesh/all_models"
+
+
+class AllModelsRequest(PydanticModel):
+    """
+    Request to get all the models that are in the current project.
+    """
+
+    textDocument: types.TextDocumentIdentifier
+
+
+class AllModelsResponse(PydanticModel):
+    """
+    Response to get all the models that are in the current project.
+    """
+
+    models: t.List[str]

sqlmesh/lsp/main.py

@@ -12,6 +12,7 @@
 from sqlmesh.core.context import Context
 from sqlmesh.core.linter.definition import AnnotatedRuleViolation
 from sqlmesh.lsp.context import LSPContext
+from sqlmesh.lsp.custom import ALL_MODELS_FEATURE, AllModelsRequest, AllModelsResponse
 from sqlmesh.lsp.reference import get_model_definitions_for_a_path
 
 
@@ -38,6 +39,12 @@ def __init__(
     def _register_features(self) -> None:
         """Register LSP features on the internal LanguageServer instance."""
 
+        @self.server.feature(ALL_MODELS_FEATURE)
+        def all_models(ls: LanguageServer, params: AllModelsRequest) -> AllModelsResponse:
+            context = self._context_get_or_load(params.textDocument.uri)
+            models = context.context.models
+            return AllModelsResponse(models=[model.name for model in models.values()])
+
         @self.server.feature(types.TEXT_DOCUMENT_DID_OPEN)
         def did_open(ls: LanguageServer, params: types.DidOpenTextDocumentParams) -> None:
             context = self._context_get_or_load(params.text_document.uri)
@@ -130,43 +137,6 @@ def formatting(
                 ls.show_message(f"Error formatting SQL: {e}", types.MessageType.Error)
                 return []
 
-        @self.server.feature(types.TEXT_DOCUMENT_DEFINITION)
-        def goto_definition(
-            ls: LanguageServer, params: types.DefinitionParams
-        ) -> t.List[types.LocationLink]:
-            """Jump to an object's definition."""
-            try:
-                self._ensure_context_for_document(params.text_document.uri)
-                document = ls.workspace.get_document(params.text_document.uri)
-                if self.lsp_context is None:
-                    raise RuntimeError(f"No context found for document: {document.path}")
-
-                references = get_model_definitions_for_a_path(
-                    self.lsp_context, params.text_document.uri
-                )
-                if len(references) == 0:
-                    return []
-
-                return [
-                    types.LocationLink(
-                        target_uri=reference.uri,
-                        target_selection_range=types.Range(
-                            start=types.Position(line=0, character=0),
-                            end=types.Position(line=0, character=0),
-                        ),
-                        target_range=types.Range(
-                            start=types.Position(line=0, character=0),
-                            end=types.Position(line=0, character=0),
-                        ),
-                        origin_selection_range=reference.range,
-                    )
-                    for reference in references
-                ]
-
-            except Exception as e:
-                ls.show_message(f"Error formatting SQL: {e}", types.MessageType.Error)
-                return []
-
         @self.server.feature(types.TEXT_DOCUMENT_DEFINITION)
         def goto_definition(
             ls: LanguageServer, params: types.DefinitionParams

vscode/extension/package.json

@@ -57,6 +57,11 @@
         "command": "sqlmesh.signout",
         "title": "Sign out from Tobiko Cloud",
         "description": "SQLMesh"
+      },
+      {
+        "command": "sqlmesh.allModels",
+        "title": "Get all models",
+        "description": "SQLMesh"
       }
     ]
   },

vscode/extension/src/completion/completion.ts

@@ -0,0 +1,30 @@
+import * as vscode from 'vscode'
+import { LSPClient } from '../lsp/lsp'
+import { isErr } from '../utilities/functional/result'
+
+export const selector: vscode.DocumentSelector = {
+  pattern: '**/*.sql',
+}
+
+export const completionProvider = (
+  lsp: LSPClient,
+): vscode.CompletionItemProvider => {
+  return {
+    async provideCompletionItems(document) {
+      const result = await lsp.call_custom_method('sqlmesh/all_models', {
+        textDocument: {
+          uri: document.uri.fsPath,
+        },
+      })
+      const modelCompletions = isErr(result)
+        ? []
+        : result.value.models.map((model): vscode.CompletionItem => {
+            return new vscode.CompletionItem(
+              model,
+              vscode.CompletionItemKind.Module,
+            )
+          })
+      return new vscode.CompletionList(modelCompletions)
+    },
+  }
+}

vscode/extension/src/extension.ts

@@ -18,6 +18,8 @@ import {
   handleSqlmeshLspNotFoundError,
   handleSqlmeshLspDependenciesMissingError,
 } from './utilities/errors'
+import { completionProvider } from './completion/completion'
+import { selector } from './completion/completion'
 
 let lspClient: LSPClient | undefined
 
@@ -82,6 +84,13 @@ export async function activate(context: vscode.ExtensionContext) {
     context.subscriptions.push(lspClient)
   }
 
+  context.subscriptions.push(
+    vscode.languages.registerCompletionItemProvider(
+      selector,
+      completionProvider(lspClient),
+    ),
+  )
+
   const restart = async () => {
     if (lspClient) {
       traceVerbose('Restarting LSP client')

vscode/extension/src/lsp/custom.ts

@@ -0,0 +1,32 @@
+export interface AllModelsMethod {
+  method: 'sqlmesh/all_models'
+  request: AllModelsRequest
+  response: AllModelsResponse
+}
+
+export interface TestMethod {
+  method: 'sqlmesh/test'
+  request: TestRequest
+  response: TestResponse
+}
+
+// @eslint-disable-next-line  @typescript-eslint/consistent-type-definition
+export type CustomLSPMethods = AllModelsMethod | TestMethod
+
+interface AllModelsRequest {
+  textDocument: {
+    uri: string
+  }
+}
+
+interface AllModelsResponse {
+  models: string[]
+}
+
+interface TestRequest {
+  foo: string
+}
+
+interface TestResponse {
+  bar: string
+}

vscode/extension/src/lsp/lsp.ts

@@ -10,6 +10,7 @@ import { err, isErr, ok, Result } from '../utilities/functional/result'
 import { getWorkspaceFolders } from '../utilities/common/vscodeapi'
 import { traceError } from '../utilities/common/log'
 import { ErrorType } from '../utilities/errors'
+import { CustomLSPMethods } from './custom'
 
 let outputChannel: OutputChannel | undefined
 
@@ -98,4 +99,23 @@ export class LSPClient implements Disposable {
   public async dispose() {
     await this.stop()
   }
+
+  public async call_custom_method<
+    Method extends CustomLSPMethods['method'],
+    Request extends Extract<CustomLSPMethods, { method: Method }>['request'],
+    Response extends Extract<CustomLSPMethods, { method: Method }>['response'],
+  >(method: Method, request: Request): Promise<Result<Response, string>> {
+    if (!this.client) {
+      return err('lsp client not ready')
+    }
+    try {
+      const result = await this.client.sendRequest<Response>(method, request)
+      return ok(result)
+    } catch (error) {
+      traceError(
+        `lsp '${method}' request ${JSON.stringify(request)} failed: ${JSON.stringify(error)}`,
+      )
+      return err(JSON.stringify(error))
+    }
+  }
 }