feat: add model autocomplete to vscode

[ci skip]

Author: benfdking

sqlmesh/lsp/custom.py

@@ -0,0 +1,18 @@
+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])
+
         @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)

sqlmesh/lsp/reference.py

@@ -15,10 +15,7 @@ class Reference(PydanticModel):
     range: Range
     uri: str
 
-
-def get_model_definitions_for_a_path(
-    lint_context: LSPContext, document_uri: str
-) -> t.List[Reference]:
+def get_model_definitions_for_a_path(lint_context: LSPContext, document_uri: str) -> t.List[Reference]:
     """
     Get the model references for a given path.
 
@@ -36,41 +33,38 @@ def get_model_definitions_for_a_path(
     # Ensure the path is a sql model
     if not document_uri.endswith(".sql"):
         return []
-
-    # Get the model
+    
+    # Get the model 
     models = lint_context.map[document_uri]
-    if models is None:
+    if models is None: 
         return []
     if len(models) == 0:
         return []
     model_name = models[0]
     model = lint_context.context.get_model(model_or_snapshot=model_name, raise_if_missing=False)
-    if model is None:
+    if model is None: 
         return []
     if not isinstance(model, SqlModel):
         return []
-
+    
     # Find all possible references
     tables = list(model.query.find_all(exp.Table))
     if len(tables) == 0:
         return []
-
+    
     references = []
     for table in tables:
         depends_on = model.depends_on
 
         # Normalize the table reference
         reference_name = table.this.this if table.db is None else f"{table.db}.{table.this.this}"
-        normalized_reference_name = normalize_model_name(
-            reference_name, default_catalog=lint_context.context.default_catalog
-        )
+        normalized_reference_name = normalize_model_name(reference_name, default_catalog=lint_context.context.default_catalog)
         if normalized_reference_name not in depends_on:
             continue
 
+
         # Get the referenced model uri
-        referenced_model = lint_context.context.get_model(
-            model_or_snapshot=normalized_reference_name, raise_if_missing=False
-        )
+        referenced_model = lint_context.context.get_model(model_or_snapshot=normalized_reference_name, raise_if_missing=False)
         if referenced_model is None:
             continue
         # Get the model uri
@@ -79,96 +73,42 @@ def get_model_definitions_for_a_path(
             continue
         # Fully qualify the path in case
         path = Path.resolve(Path(referenced_model_path))
-        referenced_model_uri = f"file://{path}"
-        read_file = open(path, "r").readlines()
-
+        referenced_model_path = f"file://{path}"
+        # Get the path to the file containing the reference
+        # file_path = document_uri.removeprefix("file://")
+        # read_file = open(file_path, "r").readlines() # Reading the file here is not needed for range calculation
+        
         # Extract metadata for positioning
-        table_meta = TokenPositionDetails.from_meta(table.this.meta)
-        table_range = _range_from_token_position_details(table_meta, read_file)
-        start_pos = table_range.start
-        end_pos = table_range.end
-
+        this_meta = table.this.meta
+        this_start_line_0 = this_meta['line'] - 1  # Convert to 0-indexed
+        this_start_col_0 = this_meta['col'] - 1   # Convert to 0-indexed
+        
+        # End position: Based on the start of 'this' token + its length
+        end_char_0 = this_start_col_0 + len(table.this.this)
+        end_pos = Position(line=this_start_line_0, character=end_char_0)
+        
+        # Start position: Initially set to the start of 'this' token
+        start_pos = Position(line=this_start_line_0, character=this_start_col_0)
+        
         # If there's a database qualifier, adjust the start position
-        db = table.args.get("db")
+        db = table.args.get('db')
         if db is not None:
-            db_meta = TokenPositionDetails.from_meta(db.meta)
-            db_range = _range_from_token_position_details(db_meta, read_file)
-            start_pos = db_range.start
-
+            db_meta = db.meta
+            db_start_line_0 = db_meta['line'] - 1  # Convert to 0-indexed
+            db_start_col_0 = db_meta['col'] - 1   # Convert to 0-indexed
+            start_pos = Position(line=db_start_line_0, character=db_start_col_0)
+        
         # If there's a catalog qualifier, adjust the start position further
-        catalog = table.args.get("catalog")
+        catalog = table.args.get('catalog')
         if catalog is not None:
-            catalog_meta = TokenPositionDetails.from_meta(catalog.meta)
-            catalog_range = _range_from_token_position_details(catalog_meta, read_file)
-            start_pos = catalog_range.start
+            catalog_meta = catalog.meta
+            catalog_start_line_0 = catalog_meta['line'] - 1 # Convert to 0-indexed
+            catalog_start_col_0 = catalog_meta['col'] - 1  # Convert to 0-indexed
+            start_pos = Position(line=catalog_start_line_0, character=catalog_start_col_0)
 
-        references.append(
-            Reference(uri=referenced_model_uri, range=Range(start=start_pos, end=end_pos))
-        )
+        references.append(Reference(
+            uri=referenced_model_path,
+            range=Range(start=start_pos, end=end_pos)
+        ))
 
     return references
-
-
-class TokenPositionDetails(PydanticModel):
-    """
-    Details about a token's position in the source code.
-
-    Attributes:
-        line (int): The line that the token ends on.
-        col (int): The column that the token ends on.
-        start (int): The start index of the token.
-        end (int): The ending index of the token.
-    """
-
-    line: int
-    col: int
-    start: int
-    end: int
-
-    @staticmethod
-    def from_meta(meta: t.Dict[str, int]) -> "TokenPositionDetails":
-        return TokenPositionDetails(
-            line=meta["line"],
-            col=meta["col"],
-            start=meta["start"],
-            end=meta["end"],
-        )
-
-
-def _range_from_token_position_details(
-    token_position_details: TokenPositionDetails, read_file: t.List[str]
-) -> Range:
-    """
-    Convert a TokenPositionDetails object to a Range object.
-
-    :param token_position_details: Details about a token's position
-    :param read_file: List of lines from the file
-    :return: A Range object representing the token's position
-    """
-    # Convert from 1-indexed to 0-indexed for line and column
-    end_line_0 = token_position_details.line - 1
-    end_col_0 = token_position_details.col
-
-    # Find the start line and column by counting backwards from the end position
-    start_pos = token_position_details.start
-    end_pos = token_position_details.end
-
-    # Initialize with the end position
-    start_line_0 = end_line_0
-    start_col_0 = end_col_0 - (end_pos - start_pos + 1)
-
-    # If start_col_0 is negative, we need to go back to previous lines
-    while start_col_0 < 0 and start_line_0 > 0:
-        start_line_0 -= 1
-        start_col_0 += len(read_file[start_line_0])
-        # Account for newline character
-        if start_col_0 >= 0:
-            break
-        start_col_0 += 1  # For the newline character
-
-    # Ensure we don't have negative values
-    start_col_0 = max(0, start_col_0)
-    return Range(
-        start=Position(line=start_line_0, character=start_col_0),
-        end=Position(line=end_line_0, character=end_col_0),
-    )

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,17 @@
+const selector: vscode.DocumentSelector = {
+    pattern: '**/my-special-file.json'   // ← only this file (you can use any glob)
+  };
+
+  const provider: vscode.CompletionItemProvider = {
+    provideCompletionItems(document, position) {
+      const items: vscode.CompletionItem[] = [];
+
+      const item = new vscode.CompletionItem('mySuggestion', vscode.CompletionItemKind.Keyword);
+      item.detail = 'Inserted by my extension';
+      item.insertText = 'mySuggestion';
+      items.push(item);
+
+      return items;
+    }
+  };
+

vscode/extension/src/extension.ts

@@ -82,6 +82,10 @@ export async function activate(context: vscode.ExtensionContext) {
     context.subscriptions.push(lspClient)
   }
 
+  ctx.subscriptions.push(
+    vscode.languages.registerCompletionItemProvider(selector, provider, ...['.', ':'])
+  );
+
   const restart = async () => {
     if (lspClient) {
       traceVerbose('Restarting LSP client')

vscode/extension/src/lsp/custom.ts

@@ -0,0 +1,31 @@
+export type AllModelsMethod = {
+  method: 'sqlmesh/all_models';
+  request: AllModelsRequest;
+  response: AllModelsResponse;
+};
+
+export type TestMethod = {
+  method: 'sqlmesh/test';
+  request: TestRequest;
+  response: TestResponse;
+};
+
+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
 
@@ -27,7 +28,9 @@ export class LSPClient implements Disposable {
 
     const sqlmesh = await sqlmesh_lsp_exec()
     if (isErr(sqlmesh)) {
-      traceError(`Failed to get sqlmesh_lsp_exec, ${JSON.stringify(sqlmesh.error)}`)
+      traceError(
+        `Failed to get sqlmesh_lsp_exec, ${JSON.stringify(sqlmesh.error)}`,
+      )
       return sqlmesh
     }
     const workspaceFolders = getWorkspaceFolders()
@@ -96,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))
+    }
+  }
 }