feat: add model autocomplete to vscode

- includes dialect keywords
- includes models

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/completions.py

@@ -0,0 +1,58 @@
+from functools import lru_cache
+from sqlglot import Tokenizer
+from sqlmesh.lsp.custom import AllModelsResponse
+import typing as t
+from sqlmesh.lsp.context import LSPContext
+
+
+def get_sql_completions(context: t.Optional[LSPContext], file_uri: str) -> AllModelsResponse:
+    """
+    Return a list of completions for a given file.
+    """
+    return AllModelsResponse(
+        models=list(get_models(context, file_uri)),
+        keywords=list(get_keywords(context, file_uri)),
+    )
+
+
+def get_models(context: t.Optional[LSPContext], file_uri: t.Optional[str]) -> t.Set[str]:
+    """
+    Return a list of models for a given file.
+
+    If there is no context, return an empty list.
+    If there is a context, return a list of all models bar the ones the file itself defines.
+    """
+    if context is None:
+        return set()
+    all_models = set(model for models in context.map.values() for model in models)
+    if file_uri is not None:
+        models_file_refers_to = context.map[file_uri]
+        for model in models_file_refers_to:
+            all_models.discard(model)
+    return all_models
+
+
+def get_keywords(context: t.Optional[LSPContext], file_uri: t.Optional[str]) -> t.Set[str]:
+    """
+    Return a list of sql keywords for a given file.
+    If no context is provided, return ANSI SQL keywords.
+
+    If a context is provided but no file_uri is provided, returns the keywords
+    for the default dialect of the context.
+
+    If both a context and a file_uri are provided, returns the keywords
+    for the dialect of the model that the file belongs to.
+    """
+    return get_keywords_from_tokenizer()
+
+
+@lru_cache(maxsize=1)
+def get_keywords_from_tokenizer() -> t.Set[str]:
+    keywords = list(Tokenizer.KEYWORDS.keys())
+    expanded_keywords = []
+
+    for keyword in keywords:
+        parts = keyword.split(" ")
+        expanded_keywords.extend(parts)
+
+    return set(expanded_keywords)

sqlmesh/lsp/custom.py

@@ -0,0 +1,22 @@
+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]
+    keywords: t.List[str]

sqlmesh/lsp/main.py

@@ -11,7 +11,9 @@
 from sqlmesh._version import __version__
 from sqlmesh.core.context import Context
 from sqlmesh.core.linter.definition import AnnotatedRuleViolation
+from sqlmesh.lsp.completions import get_sql_completions
 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 +40,14 @@ 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:
+            try:
+                context = self._context_get_or_load(params.textDocument.uri)
+                return get_sql_completions(context, params.textDocument.uri)
+            except Exception as e:
+                return get_sql_completions(None, params.textDocument.uri)
+
         @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 +140,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

tests/lsp/test_completions.py

@@ -0,0 +1,44 @@
+import pytest
+from sqlglot import Tokenizer
+from sqlmesh.core.context import Context
+from sqlmesh.lsp.completions import get_keywords_from_tokenizer, get_sql_completions
+from sqlmesh.lsp.context import LSPContext
+
+
+TOKENIZER_KEYWORDS = set(Tokenizer.KEYWORDS.keys())
+
+
+@pytest.mark.fast
+def test_get_keywords_from_tokenizer():
+    assert len(get_keywords_from_tokenizer()) > len(TOKENIZER_KEYWORDS)
+
+
+@pytest.mark.fast
+def test_get_sql_completions_no_context():
+    completions = get_sql_completions(None, None)
+    assert len(completions.keywords) > len(TOKENIZER_KEYWORDS)
+    assert len(completions.models) == 0
+
+
+@pytest.mark.fast
+def test_get_sql_completions_with_context_no_file_uri():
+    context = Context(paths=["examples/sushi"])
+    lsp_context = LSPContext(context)
+
+    completions = get_sql_completions(lsp_context, None)
+    assert len(completions.keywords) > len(TOKENIZER_KEYWORDS)
+    assert "sushi.active_customers" in completions.models
+    assert "sushi.customers" in completions.models
+
+
+@pytest.mark.fast
+def test_get_sql_completions_with_context_and_file_uri():
+    context = Context(paths=["examples/sushi"])
+    lsp_context = LSPContext(context)
+
+    file_uri = next(
+        key for key in lsp_context.map.keys() if key.endswith("models/active_customers.sql")
+    )
+    completions = get_sql_completions(lsp_context, file_uri)
+    assert len(completions.keywords) > len(TOKENIZER_KEYWORDS)
+    assert "sushi.active_customers" not in completions.models

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,36 @@
+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,
+        },
+      })
+      if (isErr(result)) {
+        return []
+      }
+      const modelCompletions = result.value.models.map(
+        model =>
+          new vscode.CompletionItem(model, vscode.CompletionItemKind.Reference),
+      )
+      const keywordCompletions = result.value.keywords.map(
+        keyword =>
+          new vscode.CompletionItem(keyword, vscode.CompletionItemKind.Keyword),
+      )
+      return new vscode.CompletionList([
+        ...modelCompletions,
+        ...keywordCompletions,
+      ])
+    },
+  }
+}

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,33 @@
+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[]
+  keywords: 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))
+    }
+  }
 }