feat(vscode): add completion provider (#4300)

Author: benfdking

sqlmesh/lsp/completions.py

@@ -0,0 +1,77 @@
+from functools import lru_cache
+from sqlglot import Dialect, 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.
+    """
+    if file_uri is not None and context is not None:
+        models = context.map[file_uri]
+        if models:
+            model = models[0]
+            model_from_context = context.context.get_model(model)
+            if model_from_context is not None:
+                if model_from_context.dialect:
+                    return get_keywords_from_tokenizer(model_from_context.dialect)
+    if context is not None:
+        return get_keywords_from_tokenizer(context.context.default_dialect)
+    return get_keywords_from_tokenizer(None)
+
+
+@lru_cache()
+def get_keywords_from_tokenizer(dialect: t.Optional[str] = None) -> t.Set[str]:
+    """
+    Return a list of sql keywords for a given dialect. This is separate from
+    the direct use of Tokenizer.KEYWORDS.keys() because that returns a set of
+    keywords that are expanded, e.g. "ORDER BY" -> ["ORDER", "BY"].
+    """
+    tokenizer = Tokenizer
+    if dialect is not None:
+        try:
+            tokenizer = Dialect.get_or_raise(dialect).tokenizer_class
+        except Exception:
+            pass
+
+    expanded_keywords = set()
+    for keyword in tokenizer.KEYWORDS.keys():
+        parts = keyword.split(" ")
+        expanded_keywords.update(parts)
+    return 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)

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/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,19 @@
+export interface AllModelsMethod {
+  method: 'sqlmesh/all_models'
+  request: AllModelsRequest
+  response: AllModelsResponse
+}
+
+// @eslint-disable-next-line  @typescript-eslint/consistent-type-definition
+export type CustomLSPMethods = AllModelsMethod
+
+interface AllModelsRequest {
+  textDocument: {
+    uri: string
+  }
+}
+
+interface AllModelsResponse {
+  models: string[]
+  keywords: 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))
+    }
+  }
 }