feat(vscode): adding ability to update columns

[ci skip]

Author: benfdking

requirements.md

@@ -0,0 +1,42 @@
+Ability to click update columns on a model in external models
+
+# Requirements 
+
+* [ ] For each model in external models file
+* [ ] Below each model in the yaml file, provide an LSP code action to update columns
+* [ ] The code action should be able to update the columns in the model
+
+# Useful files
+
+ sqlmesh/lsp/main.py is the main file for the LSP server.
+ sqlmesh/cli/main.py and create_external_model provides similar functionality for creating external models but for the CLI and only does the whole shebang.
+
+# Implementation Plan
+
+1. **Create a new code action handler for external models YAML files**
+   - Add support for YAML files in the LSP context's `get_code_actions` method
+   - Detect when cursor is on or near a model definition in external_models.yaml
+
+2. **Add YAML file parsing to identify model positions**
+   - Parse the YAML file to identify model definitions and their line positions
+   - Track the name field and columns section for each model
+
+3. **Implement the "Update Columns" code action**
+   - When triggered, the action will:
+     - Extract the model name from the YAML
+     - Use the engine adapter to fetch current columns from the database
+     - Generate the updated columns section in YAML format
+     - Replace the existing columns section with the new one
+
+4. **Integration points**
+   - Extend `LSPContext.get_code_actions()` to handle YAML files
+   - Add YAML parsing logic to identify model boundaries
+   - Reuse existing column fetching logic from `schema_loader.py`
+   - Use the adapter's `columns()` method to get schema information
+
+5. **Code action workflow**
+   - User opens external_models.yaml in editor
+   - LSP server identifies model definitions in the file
+   - When cursor is on a model, show "Update columns" code action
+   - On activation, fetch latest schema and update the YAML
+ 

sqlmesh/lsp/commands.py

@@ -0,0 +1,6 @@
+from sqlmesh.utils.pydantic import PydanticModel
+
+EXTERNAL_MODEL_UPDATE_COLUMNS = 'sqlmesh.external_model_update_columns'
+
+class ExternalModelUpdateColumnsRequest(PydanticModel):
+    model_name: str

sqlmesh/lsp/context.py

@@ -3,12 +3,14 @@
 from sqlmesh.core.context import Context
 import typing as t
 
-from sqlmesh.core.model.definition import SqlModel
+from sqlmesh.core.model.definition import SqlModel, ExternalModel
 from sqlmesh.core.linter.definition import AnnotatedRuleViolation
+from sqlmesh.lsp.commands import EXTERNAL_MODEL_UPDATE_COLUMNS
 from sqlmesh.lsp.custom import ModelForRendering
 from sqlmesh.lsp.custom import AllModelsResponse, RenderModelEntry
 from sqlmesh.lsp.uri import URI
 from lsprotocol import types
+from pathlib import Path
 
 
 @dataclass
@@ -228,6 +230,40 @@ def get_code_actions(
                         edit=types.WorkspaceEdit(changes={params.text_document.uri: text_edits}),
                     )
                     code_actions.append(code_action)
+        return code_actions
+
+    def get_code_lenses(self, uri: URI ) -> t.Optional[t.List[types.CodeLens]]:
+        models_in_file = self.map.get(uri.to_path())
+        if isinstance(models_in_file, ModelTarget):
+            models = [self.context.get_model(model) for model in models_in_file.names]
+            if any(isinstance(model, ExternalModel) for model in models):
+                code_lenses = self._get_external_model_code_lenses(uri)
+                if code_lenses:
+                    return code_lenses
+
+        return None
+
+    def _get_external_model_code_lenses(self, uri: URI) -> t.List[types.CodeLens]:
+        from sqlmesh.lsp.reference import get_yaml_model_name_ranges
+
+        """Get code actions for external models YAML files."""
+        ranges = get_yaml_model_name_ranges(uri.to_path())
+
+        code_actions: t.List[types.CodeLens] = []
+        for name, range in ranges.items():
+            # Create a code action to update columns for external models
+            command = types.Command(
+                title="Update Columns",
+                command=EXTERNAL_MODEL_UPDATE_COLUMNS,
+                arguments=[
+                     name,
+                ],
+            )
+            code_action = types.CodeLens(
+                range=range,
+                command=command,
+            )
+            code_actions.append(code_action)
 
         return code_actions if code_actions else None
 

sqlmesh/lsp/main.py

@@ -23,6 +23,7 @@
     ApiResponseGetLineage,
     ApiResponseGetModels,
 )
+from sqlmesh.lsp.commands import ExternalModelUpdateColumnsRequest, EXTERNAL_MODEL_UPDATE_COLUMNS
 from sqlmesh.lsp.completions import get_sql_completions
 from sqlmesh.lsp.context import (
     LSPContext,
@@ -63,6 +64,7 @@
 from web.server.api.endpoints.models import get_models
 from typing import Union
 from dataclasses import dataclass, field
+from sqlmesh.utils import yaml
 
 
 @dataclass
@@ -313,6 +315,25 @@ def function_call(ls: LanguageServer, params: t.Any) -> t.Dict[str, t.Any]:
 
             self.server.feature(name)(create_function_call(method))
 
+        @self.server.command(EXTERNAL_MODEL_UPDATE_COLUMNS)
+        def execute_command(ls: LanguageServer, raw: t.Any) -> None:
+            try:
+                if not isinstance(raw, list) or not len(list) == 1 or not isinstance(raw[0], str):
+                    raise ValueError("Invalid command parameters")
+
+                request = ExternalModelUpdateColumnsRequest.model_validate(raw[0])
+                context = self._context_get_or_load()
+                if not isinstance(context, LSPContext):
+                    raise ValueError("Context is not loaded or invalid")
+                context.context.update_external_model_columns(
+                    model_name=request.model_name,
+                )
+
+                ls.show_message(f"Executing command to update external model columns {raw}", types.MessageType.Info)
+            except Exception as e:
+                ls.show_message(f"Error executing command: {e}", types.MessageType.Error)
+                return None
+
         @self.server.feature(types.INITIALIZE)
         def initialize(ls: LanguageServer, params: types.InitializeParams) -> None:
             """Initialize the server when the client connects."""
@@ -688,6 +709,17 @@ def code_action(
                 ls.log_trace(f"Error getting code actions: {e}")
                 return None
 
+        @self.server.feature(types.TEXT_DOCUMENT_CODE_LENS)
+        def code_lens(ls: LanguageServer, params: types.CodeLensParams) -> t.List[types.CodeLens]:
+            try:
+                uri = URI(params.text_document.uri)
+                context = self._context_get_or_load(uri)
+                code_lenses = context.get_code_lenses(uri)
+                return code_lenses
+            except Exception as e:
+                ls.log_trace(f"Error getting code lenses: {e}")
+                return []
+
         @self.server.feature(
             types.TEXT_DOCUMENT_COMPLETION,
             types.CompletionOptions(trigger_characters=["@"]),  # advertise "@" for macros
@@ -902,6 +934,68 @@ def _uri_to_path(uri: str) -> Path:
         """Convert a URI to a path."""
         return URI(uri).to_path()
 
+    def _update_external_model_columns(self, ls: LanguageServer, arguments: t.List[t.Any]) -> t.Any:
+        """Update the columns for an external model in the YAML file."""
+        if len(arguments) != 3:
+            ls.show_message("Invalid arguments for update columns command", types.MessageType.Error)
+            return None
+
+        uri_str, model_name, model_idx = arguments
+        uri = URI(uri_str)
+        path = uri.to_path()
+
+        try:
+            # Get the context
+            context = self._context_get_or_load(uri)
+
+            # Read the YAML file
+            models = yaml.load(path)
+            if not isinstance(models, list) or model_idx >= len(models):
+                ls.show_message("Invalid model index", types.MessageType.Error)
+                return None
+
+            model = models[model_idx]
+            if model.get("name") != model_name:
+                ls.show_message("Model name mismatch", types.MessageType.Error)
+                return None
+
+            # Get the adapter and fetch columns
+            adapter = context.context.engine_adapter
+
+            # Get columns for the model
+            try:
+                columns = adapter.columns(model_name, include_pseudo_columns=True)
+            except Exception as e:
+                ls.show_message(
+                    f"Unable to fetch columns for {model_name}: {e}", types.MessageType.Error
+                )
+                return None
+
+            if not columns:
+                ls.show_message(f"No columns found for {model_name}", types.MessageType.Warning)
+                return None
+
+            # Update the model's columns
+            dialect = context.context.config.model_defaults.dialect
+            model["columns"] = {
+                col_name: dtype.sql(dialect=dialect) for col_name, dtype in columns.items()
+            }
+
+            # Write back to the file
+            with open(path, "w", encoding="utf-8") as f:
+                yaml.dump(models, f)
+
+            ls.show_message(f"Updated columns for {model_name}", types.MessageType.Info)
+
+            # Reload the context to pick up the changes
+            self._reload_context_and_publish_diagnostics(ls, uri, uri_str)
+
+        except Exception as e:
+            ls.show_message(f"Error updating columns: {e}", types.MessageType.Error)
+            ls.log_trace(f"Error updating columns: {e}")
+
+        return None
+
     def start(self) -> None:
         """Start the server with I/O transport."""
         logging.basicConfig(level=logging.DEBUG)

sqlmesh/lsp/reference.py

@@ -877,6 +877,34 @@ def _process_column_references(
     return references
 
 
+def get_yaml_model_name_ranges(path: Path) -> t.Dict[str, Range]:
+    """
+    Get a mapping of model names to their ranges in an external models YAML file.
+
+    Args:
+        path: The path to the YAML file.
+    Returns:
+        A dictionary mapping model names to their ranges in the YAML file.
+    """
+    yaml = YAML()
+    model_ranges: t.Dict[str, Range] = {}
+    with path.open("r", encoding="utf-8") as f:
+        data = yaml.load(f)
+
+    if not isinstance(data, list):
+        return model_ranges
+
+    for item in data:
+        if isinstance(item, dict) and "name" in item:
+            # Get size of block by taking the earliest line/col in the items block and the last line/col of the block
+            position_data = item.lc.data["name"]  # type: ignore
+            start = Position(line=position_data[2], character=position_data[3])
+            end = Position(line=position_data[2], character=position_data[3] + len(item["name"]))
+            model_ranges[item["name"]] = Range(start=start, end=end)
+
+    return model_ranges
+
+
 def _get_yaml_model_range(path: Path, model_name: str) -> t.Optional[Range]:
     """
     Find the range of a specific model block in a YAML file.
@@ -895,11 +923,8 @@ def _get_yaml_model_range(path: Path, model_name: str) -> t.Optional[Range]:
     if not isinstance(data, list):
         return None
 
-    for item in data:
-        if isinstance(item, dict) and item.get("name") == model_name:
-            # Get size of block by taking the earliest line/col in the items block and the last line/col of the block
-            position_data = item.lc.data["name"]  # type: ignore
-            start = Position(line=position_data[2], character=position_data[3])
-            end = Position(line=position_data[2], character=position_data[3] + len(item["name"]))
-            return Range(start=start, end=end)
+    # Get all model ranges in the YAML file
+    model_ranges = get_yaml_model_name_ranges(path)
+    if model_name in model_ranges:
+        return model_ranges[model_name]
     return None