feat(vscode): adding ability to update columns

Author: benfdking

pnpm-lock.yaml

@@ -57,28 +57,17 @@ def create_external_models_file(
         external_model_fqns -= existing_model_fqns
 
     with ThreadPoolExecutor(max_workers=max_workers) as pool:
-
-        def _get_columns(table: str) -> t.Optional[t.Dict[str, t.Any]]:
-            try:
-                return adapter.columns(table, include_pseudo_columns=True)
-            except Exception as e:
-                msg = f"Unable to get schema for '{table}': '{e}'."
-                if strict:
-                    raise SQLMeshError(msg) from e
-                get_console().log_warning(msg)
-                return None
-
         gateway_part = {"gateway": gateway} if gateway else {}
 
         schemas = [
             {
                 "name": exp.to_table(table).sql(dialect=dialect),
-                "columns": {c: dtype.sql(dialect=dialect) for c, dtype in columns.items()},
+                "columns": columns,
                 **gateway_part,
             }
             for table, columns in sorted(
                 pool.map(
-                    lambda table: (table, _get_columns(table)),
+                    lambda table: (table, get_columns(adapter, dialect, table, strict)),
                     external_model_fqns,
                 )
             )
@@ -94,3 +83,20 @@ def _get_columns(table: str) -> t.Optional[t.Dict[str, t.Any]]:
 
         with open(path, "w", encoding="utf-8") as file:
             yaml.dump(entries_to_keep + schemas, file)
+
+
+def get_columns(
+    adapter: EngineAdapter, dialect: DialectType, table: str, strict: bool
+) -> t.Optional[t.Dict[str, t.Any]]:
+    """
+    Return the column and their types in a dictionary
+    """
+    try:
+        columns = adapter.columns(table, include_pseudo_columns=True)
+        return {c: dtype.sql(dialect=dialect) for c, dtype in columns.items()}
+    except Exception as e:
+        msg = f"Unable to get schema for '{table}': '{e}'."
+        if strict:
+            raise SQLMeshError(msg) from e
+        get_console().log_warning(msg)
+        return None

sqlmesh/core/schema_loader.py

@@ -0,0 +1 @@
+EXTERNAL_MODEL_UPDATE_COLUMNS = "sqlmesh.external_model_update_columns"

sqlmesh/lsp/commands.py

@@ -1,16 +1,21 @@
 from dataclasses import dataclass
 from pathlib import Path
+from pygls.server import LanguageServer
 from sqlmesh.core.context import Context
 import typing as t
-
 from sqlmesh.core.linter.rule import Range
-from sqlmesh.core.model.definition import SqlModel
+from sqlmesh.core.model.definition import SqlModel, ExternalModel
 from sqlmesh.core.linter.definition import AnnotatedRuleViolation
+from sqlmesh.core.schema_loader import get_columns
+from sqlmesh.lsp.commands import EXTERNAL_MODEL_UPDATE_COLUMNS
 from sqlmesh.lsp.custom import ModelForRendering, TestEntry, RunTestResponse
 from sqlmesh.lsp.custom import AllModelsResponse, RenderModelEntry
 from sqlmesh.lsp.tests_ranges import get_test_ranges
+from sqlmesh.lsp.helpers import to_lsp_range
 from sqlmesh.lsp.uri import URI
 from lsprotocol import types
+from sqlmesh.utils import yaml
+from sqlmesh.utils.lineage import get_yaml_model_name_ranges
 
 
 @dataclass
@@ -298,6 +303,42 @@ def get_code_actions(
 
         return code_actions if code_actions else None
 
+    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]:
+        """Get code lenses for external models YAML files."""
+        ranges = get_yaml_model_name_ranges(uri.to_path())
+        if not ranges:
+            return []
+
+        code_lenses: 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_lenses.append(
+                types.CodeLens(
+                    range=to_lsp_range(range),
+                    command=command,
+                )
+            )
+
+        return code_lenses if code_lenses else []
+
     def list_of_models_for_rendering(self) -> t.List[ModelForRendering]:
         """Get a list of models for rendering.
 
@@ -399,3 +440,72 @@ def diagnostic_to_lsp_diagnostic(
             code=diagnostic.rule.name,
             code_description=types.CodeDescription(href=rule_uri),
         )
+
+    def update_external_model_columns(self, ls: LanguageServer, uri: URI, model_name: str) -> bool:
+        """
+        Update the columns for an external model in the YAML file. Returns True if changed, False if didn't because
+        of the columns already being up to date.
+
+        Errors still throw exceptions to be handled by the caller.
+        """
+        models = yaml.load(uri.to_path())
+        if not isinstance(models, list):
+            raise ValueError(
+                f"Expected a list of models in {uri.to_path()}, but got {type(models).__name__}"
+            )
+
+        existing_model = next((model for model in models if model.get("name") == model_name), None)
+        if existing_model is None:
+            raise ValueError(f"Could not find model {model_name} in {uri.to_path()}")
+
+        existing_model_columns = existing_model.get("columns")
+
+        # Get the adapter and fetch columns
+        adapter = self.context.engine_adapter
+        # Get columns for the model
+        new_columns = get_columns(
+            adapter=adapter,
+            dialect=self.context.config.model_defaults.dialect,
+            table=model_name,
+            strict=True,
+        )
+        # Compare existing columns and matching types and if they are the same, do not update
+        if existing_model_columns is not None:
+            if existing_model_columns == new_columns:
+                return False
+
+        # Model index to update
+        model_index = next(
+            (i for i, model in enumerate(models) if model.get("name") == model_name), None
+        )
+        if model_index is None:
+            raise ValueError(f"Could not find model {model_name} in {uri.to_path()}")
+
+        # Get end of the file to set the edit range
+        with open(uri.to_path(), "r", encoding="utf-8") as file:
+            read_file = file.read()
+
+        end_line = read_file.count("\n")
+        end_character = len(read_file.splitlines()[-1]) if end_line > 0 else 0
+
+        models[model_index]["columns"] = new_columns
+        edit = types.TextDocumentEdit(
+            text_document=types.OptionalVersionedTextDocumentIdentifier(
+                uri=uri.value,
+                version=None,
+            ),
+            edits=[
+                types.TextEdit(
+                    range=types.Range(
+                        start=types.Position(line=0, character=0),
+                        end=types.Position(
+                            line=end_line,
+                            character=end_character,
+                        ),
+                    ),
+                    new_text=yaml.dump(models),
+                )
+            ],
+        )
+        ls.apply_edit(types.WorkspaceEdit(document_changes=[edit]))
+        return True

sqlmesh/lsp/context.py

@@ -24,6 +24,7 @@
     ApiResponseGetModels,
 )
 
+from sqlmesh.lsp.commands import EXTERNAL_MODEL_UPDATE_COLUMNS
 from sqlmesh.lsp.completions import get_sql_completions
 from sqlmesh.lsp.context import (
     LSPContext,
@@ -368,6 +369,47 @@ def function_call(ls: LanguageServer, params: t.Any) -> t.Dict[str, t.Any]:
 
             self.server.feature(name)(create_function_call(method))
 
+        # Check if command is already registered to prevent duplicate registration errors
+        if EXTERNAL_MODEL_UPDATE_COLUMNS not in getattr(self.server, "_commands", {}):
+
+            @self.server.command(EXTERNAL_MODEL_UPDATE_COLUMNS)
+            def command_external_models_update_columns(ls: LanguageServer, raw: t.Any) -> None:
+                try:
+                    if not isinstance(raw, list):
+                        raise ValueError("Invalid command parameters")
+                    if len(raw) != 1:
+                        raise ValueError("Command expects exactly one parameter")
+                    model_name = raw[0]
+                    if not isinstance(model_name, str):
+                        raise ValueError("Command parameter must be a string")
+
+                    context = self._context_get_or_load()
+                    if not isinstance(context, LSPContext):
+                        raise ValueError("Context is not loaded or invalid")
+                    model = context.context.get_model(model_name)
+                    if model is None:
+                        raise ValueError(f"External model '{model_name}' not found")
+                    if model._path is None:
+                        raise ValueError(f"External model '{model_name}' does not have a file path")
+                    uri = URI.from_path(model._path)
+                    updated = context.update_external_model_columns(
+                        ls=ls,
+                        uri=uri,
+                        model_name=model_name,
+                    )
+                    if updated:
+                        ls.show_message(
+                            f"Updated columns for '{model_name}'",
+                            types.MessageType.Info,
+                        )
+                    else:
+                        ls.show_message(
+                            f"Columns for '{model_name}' are already up to date",
+                        )
+                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."""
@@ -750,6 +792,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 if code_lenses else []
+            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

sqlmesh/lsp/main.py

@@ -387,18 +387,38 @@ def _get_yaml_model_range(path: Path, model_name: str) -> t.Optional[Range]:
     Returns:
         The Range of the model block in the YAML file, or None if not found
     """
+    model_name_ranges = get_yaml_model_name_ranges(path)
+    if model_name_ranges is None:
+        return None
+    return model_name_ranges.get(model_name, None)
+
+
+def get_yaml_model_name_ranges(path: Path) -> t.Optional[t.Dict[str, Range]]:
+    """
+    Get the ranges of all model names in a YAML file.
+
+    Args:
+        path: Path to the YAML file
+
+    Returns:
+        A dictionary mapping model names to their ranges in the YAML file.
+    """
     yaml = YAML()
     with path.open("r", encoding="utf-8") as f:
         data = yaml.load(f)
 
     if not isinstance(data, list):
         return None
 
+    model_name_ranges = {}
     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
+        if isinstance(item, dict):
             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)
-    return None
+            name = item.get("name")
+            if not name :
+                continue
+            model_name_ranges[name] = Range(start=start, end=end)
+
+    return model_name_ranges