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

@@ -2,18 +2,20 @@
 
 from __future__ import annotations
 
+import os
 import typing as t
 
 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 +187,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 +216,46 @@ 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.
+        split_lines = lines.splitlines()
+        if lines.endswith(os.linesep):
+            new_text = f"- name: '{model_name}'\n"
+            position = Position(line=len(split_lines), character=0)
+        else:
+            new_text = f"\n- name: '{model_name}'\n"
+            position = Position(
+                line=len(split_lines) - 1, character=len(split_lines[-1]) if split_lines else 0
+            )
+
+        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,)))

tests/core/linter/test_builtin.py

@@ -1,6 +1,7 @@
 import os
 
 from sqlmesh import Context
+from sqlmesh.core.linter.rule import Position, Range
 
 
 def test_no_missing_external_models(tmp_path, copy_to_temp_path) -> None:
@@ -44,8 +45,124 @@ def test_no_missing_external_models(tmp_path, copy_to_temp_path) -> None:
     # Lint the models
     lints = context.lint_models(raise_on_error=False)
     assert len(lints) == 1
-    assert lints[0].violation_range is not None
+    lint = lints[0]
+    assert lint.violation_range is not None
     assert (
-        lints[0].violation_msg
+        lint.violation_msg
         == """Model '"memory"."sushi"."customers"' depends on unregistered external model '"memory"."raw"."demographics"'. Please register it in the external models file. This can be done by running 'sqlmesh create_external_models'."""
     )
+    assert len(lint.fixes) == 0
+
+
+def test_no_missing_external_models_with_existing_file_ending_in_newline(
+    tmp_path, copy_to_temp_path
+) -> None:
+    sushi_paths = copy_to_temp_path("examples/sushi")
+    sushi_path = sushi_paths[0]
+
+    # Overwrite the external_models.yaml file to end with a random file and a newline
+    os.remove(sushi_path / "external_models.yaml")
+    with open(sushi_path / "external_models.yaml", "w") as f:
+        f.write("- name: memory.raw.test\n")
+
+    # Override the config.py to turn on lint
+    with open(sushi_path / "config.py", "r") as f:
+        read_file = f.read()
+
+    before = """    linter=LinterConfig(
+        enabled=False,
+        rules=[
+            "ambiguousorinvalidcolumn",
+            "invalidselectstarexpansion",
+            "noselectstar",
+            "nomissingaudits",
+            "nomissingowner",
+            "nomissingexternalmodels",
+        ],
+    ),"""
+    after = """linter=LinterConfig(enabled=True, rules=["nomissingexternalmodels"]),"""
+    read_file = read_file.replace(before, after)
+    assert after in read_file
+    with open(sushi_path / "config.py", "w") as f:
+        f.writelines(read_file)
+
+    # Load the context with the temporary sushi path
+    context = Context(paths=[sushi_path])
+
+    # Lint the models
+    lints = context.lint_models(raise_on_error=False)
+    assert len(lints) == 1
+    lint = lints[0]
+    assert lint.violation_range is not None
+    assert (
+        lint.violation_msg
+        == """Model '"memory"."sushi"."customers"' depends on unregistered external model '"memory"."raw"."demographics"'. Please register it in the external models file. This can be done by running 'sqlmesh create_external_models'."""
+    )
+    assert len(lint.fixes) == 1
+    fix = lint.fixes[0]
+    assert len(fix.edits) == 1
+    edit = fix.edits[0]
+    assert edit.new_text == """- name: '"memory"."raw"."demographics"'\n"""
+    assert edit.range == Range(
+        start=Position(line=1, character=0),
+        end=Position(line=1, character=0),
+    )
+    fix_path = sushi_path / "external_models.yaml"
+    assert edit.path == fix_path
+
+
+def test_no_missing_external_models_with_existing_file_not_ending_in_newline(
+    tmp_path, copy_to_temp_path
+) -> None:
+    sushi_paths = copy_to_temp_path("examples/sushi")
+    sushi_path = sushi_paths[0]
+
+    # Overwrite the external_models.yaml file to end with a random file and a newline
+    os.remove(sushi_path / "external_models.yaml")
+    with open(sushi_path / "external_models.yaml", "w") as f:
+        f.write("- name: memory.raw.test")
+
+    # Override the config.py to turn on lint
+    with open(sushi_path / "config.py", "r") as f:
+        read_file = f.read()
+
+    before = """    linter=LinterConfig(
+        enabled=False,
+        rules=[
+            "ambiguousorinvalidcolumn",
+            "invalidselectstarexpansion",
+            "noselectstar",
+            "nomissingaudits",
+            "nomissingowner",
+            "nomissingexternalmodels",
+        ],
+    ),"""
+    after = """linter=LinterConfig(enabled=True, rules=["nomissingexternalmodels"]),"""
+    read_file = read_file.replace(before, after)
+    assert after in read_file
+    with open(sushi_path / "config.py", "w") as f:
+        f.writelines(read_file)
+
+    # Load the context with the temporary sushi path
+    context = Context(paths=[sushi_path])
+
+    # Lint the models
+    lints = context.lint_models(raise_on_error=False)
+    assert len(lints) == 1
+    lint = lints[0]
+    assert lint.violation_range is not None
+    assert (
+        lint.violation_msg
+        == """Model '"memory"."sushi"."customers"' depends on unregistered external model '"memory"."raw"."demographics"'. Please register it in the external models file. This can be done by running 'sqlmesh create_external_models'."""
+    )
+    assert len(lint.fixes) == 1
+    fix = lint.fixes[0]
+    assert len(fix.edits) == 1
+    edit = fix.edits[0]
+    assert edit.new_text == """\n- name: '"memory"."raw"."demographics"'\n"""
+    assert edit.range == Range(
+        start=Position(line=0, character=23),
+        end=Position(line=0, character=23),
+    )
+    fix_path = sushi_path / "external_models.yaml"
+    assert edit.path == fix_path