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,20 @@
 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
 
 
 @dataclass
@@ -298,6 +302,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]:
+        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_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 +439,73 @@ 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:
+                ls.show_message("Columns already up to date")
+                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

@@ -23,6 +23,7 @@
     ApiResponseGetLineage,
     ApiResponseGetModels,
 )
+from sqlmesh.lsp.commands import EXTERNAL_MODEL_UPDATE_COLUMNS
 from sqlmesh.lsp.completions import get_sql_completions
 from sqlmesh.lsp.context import (
     LSPContext,
@@ -367,6 +368,44 @@ 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 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."""
@@ -749,6 +788,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

@@ -886,6 +886,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.
@@ -904,11 +932,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