feat: nomissingexternalmodels returns fix

- nomissingexternalmodels lint rule now returns a fix that can be used
  by the IDE

Author: benfdking

sqlmesh/core/linter/rules/builtin.py

@@ -7,13 +7,14 @@
 from sqlglot.expressions import Star
 from sqlglot.helper import subclasses
 
+from sqlmesh.core.constants import EXTERNAL_MODELS_YAML
 from sqlmesh.core.dialect import normalize_model_name
 from sqlmesh.core.linter.helpers import (
     TokenPositionDetails,
     get_range_of_model_block,
     read_range_from_string,
 )
-from sqlmesh.core.linter.rule import Rule, RuleViolation, Range, Fix, TextEdit
+from sqlmesh.core.linter.rule import Rule, RuleViolation, Range, Fix, TextEdit, Position
 from sqlmesh.core.linter.definition import RuleSet
 from sqlmesh.core.model import Model, SqlModel, ExternalModel
 from sqlmesh.utils.lineage import extract_references_from_query, ExternalModelReference
@@ -185,12 +186,14 @@ def check_model(
         violations = []
         for ref_name, ref in external_references.items():
             if ref_name in not_registered_external_models:
+                fix = self.create_fix(ref_name)
                 violations.append(
                     RuleViolation(
                         rule=self,
                         violation_msg=f"Model '{model.fqn}' depends on unregistered external model '{ref_name}'. "
                         "Please register it in the external models file. This can be done by running 'sqlmesh create_external_models'.",
                         violation_range=ref.range,
+                        fixes=[fix] if fix else [],
                     )
                 )
 
@@ -212,5 +215,44 @@ def _standard_error_message(
             "Please register them in the external models file. This can be done by running 'sqlmesh create_external_models'.",
         )
 
+    def create_fix(self, model_name: str) -> t.Optional[Fix]:
+        """
+        Add an external model to the external models file.
+        - If no external models file exists, it will create one with the model.
+        - If the model already exists, it will not add it again.
+        """
+        root = self.context.path
+        if not root:
+            return None
+
+        external_models_path = root / EXTERNAL_MODELS_YAML
+        if not external_models_path.exists():
+            return None
+
+        # Figure out the position to insert the new external model at the end of the file, whether
+        # needs new line or not.
+        with open(external_models_path, "r", encoding="utf-8") as file:
+            lines = file.read()
+
+        # If a file ends in newline, we can add the new model directly.
+        if not lines.endswith("\n"):
+            new_text = f"\n- name: '{model_name}'\n"
+        else:
+            new_text = f"- name: '{model_name}'\n"
+
+        split_lines = lines.splitlines()
+        position = Position(line=len(split_lines) - 1, character=len(split_lines[-1]))
+
+        return Fix(
+            title="Add external model",
+            edits=[
+                TextEdit(
+                    path=external_models_path,
+                    range=Range(start=position, end=position),
+                    new_text=new_text,
+                )
+            ],
+        )
+
 
 BUILTIN_RULES = RuleSet(subclasses(__name__, Rule, (Rule,)))

sqlmesh/lsp/context.py

@@ -1,21 +1,16 @@
 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, ExternalModel
+from sqlmesh.core.model.definition import SqlModel
 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
@@ -307,36 +302,6 @@ 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 ranges is None:
-            return []
-        return [
-            types.CodeLens(
-                range=to_lsp_range(range),
-                command=types.Command(
-                    title="Update Columns",
-                    command=EXTERNAL_MODEL_UPDATE_COLUMNS,
-                    arguments=[
-                        name,
-                    ],
-                ),
-            )
-            for name, range in ranges.items()
-        ]
-
     def list_of_models_for_rendering(self) -> t.List[ModelForRendering]:
         """Get a list of models for rendering.
 
@@ -438,74 +403,3 @@ 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.
-
-        In this case, the model name is the name of the external model as is defined in the YAML file, not any other version of it.
-
-        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/main.py

@@ -24,7 +24,6 @@
     ApiResponseGetModels,
 )
 
-from sqlmesh.lsp.commands import EXTERNAL_MODEL_UPDATE_COLUMNS
 from sqlmesh.lsp.completions import get_sql_completions
 from sqlmesh.lsp.context import (
     LSPContext,
@@ -369,44 +368,6 @@ 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."""
@@ -789,17 +750,6 @@ 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