feat: add fixes to linting errors (#4800)

Author: benfdking

sqlmesh/core/linter/definition.py

@@ -8,7 +8,7 @@
 from collections.abc import Iterator, Iterable, Set, Mapping, Callable
 from functools import reduce
 from sqlmesh.core.model import Model
-from sqlmesh.core.linter.rule import Rule, RuleViolation, Range
+from sqlmesh.core.linter.rule import Rule, RuleViolation, Range, Fix
 from sqlmesh.core.console import LinterConsole, get_console
 
 if t.TYPE_CHECKING:
@@ -75,6 +75,7 @@ def lint_model(
                 model=model,
                 violation_type="error",
                 violation_range=violation.violation_range,
+                fixes=violation.fixes,
             )
             for violation in error_violations
         ] + [
@@ -84,6 +85,7 @@ def lint_model(
                 model=model,
                 violation_type="warning",
                 violation_range=violation.violation_range,
+                fixes=violation.fixes,
             )
             for violation in warn_violations
         ]
@@ -152,7 +154,8 @@ def __init__(
         model: Model,
         violation_type: t.Literal["error", "warning"],
         violation_range: t.Optional[Range] = None,
+        fixes: t.Optional[t.List[Fix]] = None,
     ) -> None:
-        super().__init__(rule, violation_msg, violation_range)
+        super().__init__(rule, violation_msg, violation_range, fixes)
         self.model = model
         self.violation_type = violation_type

sqlmesh/core/linter/rule.py

@@ -39,6 +39,22 @@ class Range:
     end: Position
 
 
+@dataclass(frozen=True)
+class TextEdit:
+    """A text edit to apply to a file."""
+
+    range: Range
+    new_text: str
+
+
+@dataclass(frozen=True)
+class Fix:
+    """A fix that can be applied to resolve a rule violation."""
+
+    title: str
+    edits: t.List[TextEdit]
+
+
 class _Rule(abc.ABCMeta):
     def __new__(cls: Type[_Rule], clsname: str, bases: t.Tuple, attrs: t.Dict) -> _Rule:
         attrs["name"] = clsname.lower()
@@ -66,10 +82,14 @@ def violation(
         self,
         violation_msg: t.Optional[str] = None,
         violation_range: t.Optional[Range] = None,
+        fixes: t.Optional[t.List[Fix]] = None,
     ) -> RuleViolation:
         """Create a RuleViolation instance for this rule"""
         return RuleViolation(
-            rule=self, violation_msg=violation_msg or self.summary, violation_range=violation_range
+            rule=self,
+            violation_msg=violation_msg or self.summary,
+            violation_range=violation_range,
+            fixes=fixes,
         )
 
     def get_definition_location(self) -> RuleLocation:
@@ -103,11 +123,16 @@ def __repr__(self) -> str:
 
 class RuleViolation:
     def __init__(
-        self, rule: Rule, violation_msg: str, violation_range: t.Optional[Range] = None
+        self,
+        rule: Rule,
+        violation_msg: str,
+        violation_range: t.Optional[Range] = None,
+        fixes: t.Optional[t.List[Fix]] = None,
     ) -> None:
         self.rule = rule
         self.violation_msg = violation_msg
         self.violation_range = violation_range
+        self.fixes = fixes or []
 
     def __repr__(self) -> str:
         return f"{self.rule.name}: {self.violation_msg}"

sqlmesh/core/linter/rules/builtin.py

@@ -8,7 +8,7 @@
 from sqlglot.helper import subclasses
 
 from sqlmesh.core.linter.helpers import TokenPositionDetails
-from sqlmesh.core.linter.rule import Rule, RuleViolation, Range
+from sqlmesh.core.linter.rule import Rule, RuleViolation, Range, Fix, TextEdit
 from sqlmesh.core.linter.definition import RuleSet
 from sqlmesh.core.model import Model, SqlModel
 
@@ -22,7 +22,8 @@ def check_model(self, model: Model) -> t.Optional[RuleViolation]:
             return None
         if model.query.is_star:
             violation_range = self._get_range(model)
-            return self.violation(violation_range=violation_range)
+            fixes = self._create_fixes(model, violation_range)
+            return self.violation(violation_range=violation_range, fixes=fixes)
         return None
 
     def _get_range(self, model: SqlModel) -> t.Optional[Range]:
@@ -37,6 +38,28 @@ def _get_range(self, model: SqlModel) -> t.Optional[Range]:
 
         return None
 
+    def _create_fixes(
+        self, model: SqlModel, violation_range: t.Optional[Range]
+    ) -> t.Optional[t.List[Fix]]:
+        """Create fixes for the SELECT * violation."""
+        if not violation_range:
+            return None
+        columns = model.columns_to_types
+        if not columns:
+            return None
+        new_text = ", ".join(columns.keys())
+        return [
+            Fix(
+                title="Replace SELECT * with explicit column list",
+                edits=[
+                    TextEdit(
+                        range=violation_range,
+                        new_text=new_text,
+                    )
+                ],
+            )
+        ]
+
 
 class InvalidSelectStarExpansion(Rule):
     def check_model(self, model: Model) -> t.Optional[RuleViolation]:

sqlmesh/lsp/context.py

@@ -1,5 +1,6 @@
 from dataclasses import dataclass
 from pathlib import Path
+import uuid
 from sqlmesh.core.context import Context
 import typing as t
 
@@ -8,6 +9,7 @@
 from sqlmesh.lsp.custom import ModelForRendering
 from sqlmesh.lsp.custom import AllModelsResponse, RenderModelEntry
 from sqlmesh.lsp.uri import URI
+from lsprotocol import types
 
 
 @dataclass
@@ -33,8 +35,14 @@ class LSPContext:
     map: t.Dict[Path, t.Union[ModelTarget, AuditTarget]]
     _render_cache: t.Dict[Path, t.List[RenderModelEntry]]
     _lint_cache: t.Dict[Path, t.List[AnnotatedRuleViolation]]
+    _version_id: str
+    """
+    This is a version ID for the context. It is used to track changes to the context. It can be used to 
+    return a version number to the LSP client.
+    """
 
     def __init__(self, context: Context) -> None:
+        self._version_id = str(uuid.uuid4())
         self.context = context
         self._render_cache = {}
         self._lint_cache = {}
@@ -62,6 +70,11 @@ def __init__(self, context: Context) -> None:
             **audit_map,
         }
 
+    @property
+    def version_id(self) -> str:
+        """Get the version ID for the context."""
+        return self._version_id
+
     def render_model(self, uri: URI) -> t.List[RenderModelEntry]:
         """Get rendered models for a file, using cache when available.
 
@@ -150,6 +163,86 @@ def lint_model(self, uri: URI) -> t.List[AnnotatedRuleViolation]:
         self._lint_cache[path] = diagnostics
         return diagnostics
 
+    def get_code_actions(
+        self, uri: URI, params: types.CodeActionParams
+    ) -> t.Optional[t.List[t.Union[types.Command, types.CodeAction]]]:
+        """Get code actions for a file."""
+
+        # Get the violations (which contain the fixes)
+        violations = self.lint_model(uri)
+
+        # Convert violations to a map for quick lookup
+        # Use a hashable representation of Range as the key
+        violation_map: t.Dict[
+            t.Tuple[str, t.Tuple[int, int, int, int]], AnnotatedRuleViolation
+        ] = {}
+        for violation in violations:
+            if violation.violation_range:
+                lsp_diagnostic = self.diagnostic_to_lsp_diagnostic(violation)
+                if lsp_diagnostic:
+                    # Create a hashable key from the diagnostic message and range
+                    key = (
+                        lsp_diagnostic.message,
+                        (
+                            lsp_diagnostic.range.start.line,
+                            lsp_diagnostic.range.start.character,
+                            lsp_diagnostic.range.end.line,
+                            lsp_diagnostic.range.end.character,
+                        ),
+                    )
+                    violation_map[key] = violation
+
+        # Get diagnostics in the requested range
+        diagnostics = params.context.diagnostics if params.context else []
+
+        code_actions: t.List[t.Union[types.Command, types.CodeAction]] = []
+
+        for diagnostic in diagnostics:
+            # Find the corresponding violation
+            key = (
+                diagnostic.message,
+                (
+                    diagnostic.range.start.line,
+                    diagnostic.range.start.character,
+                    diagnostic.range.end.line,
+                    diagnostic.range.end.character,
+                ),
+            )
+            found_violation = violation_map.get(key)
+
+            if found_violation is not None and found_violation.fixes:
+                # Create code actions for each fix
+                for fix in found_violation.fixes:
+                    # Convert our Fix to LSP TextEdits
+                    text_edits = []
+                    for edit in fix.edits:
+                        text_edits.append(
+                            types.TextEdit(
+                                range=types.Range(
+                                    start=types.Position(
+                                        line=edit.range.start.line,
+                                        character=edit.range.start.character,
+                                    ),
+                                    end=types.Position(
+                                        line=edit.range.end.line,
+                                        character=edit.range.end.character,
+                                    ),
+                                ),
+                                new_text=edit.new_text,
+                            )
+                        )
+
+                    # Create the code action
+                    code_action = types.CodeAction(
+                        title=fix.title,
+                        kind=types.CodeActionKind.QuickFix,
+                        diagnostics=[diagnostic],
+                        edit=types.WorkspaceEdit(changes={params.text_document.uri: text_edits}),
+                    )
+                    code_actions.append(code_action)
+
+        return code_actions if code_actions else None
+
     def list_of_models_for_rendering(self) -> t.List[ModelForRendering]:
         """Get a list of models for rendering.
 
@@ -186,3 +279,68 @@ def get_completions(
         from sqlmesh.lsp.completions import get_sql_completions
 
         return get_sql_completions(self, uri, file_content)
+
+    @staticmethod
+    def diagnostics_to_lsp_diagnostics(
+        diagnostics: t.List[AnnotatedRuleViolation],
+    ) -> t.List[types.Diagnostic]:
+        """
+        Converts a list of AnnotatedRuleViolations to a list of LSP diagnostics. It will remove duplicates based on the message and range.
+        """
+        lsp_diagnostics = {}
+        for diagnostic in diagnostics:
+            lsp_diagnostic = LSPContext.diagnostic_to_lsp_diagnostic(diagnostic)
+            if lsp_diagnostic is not None:
+                # Create a unique key combining message and range
+                diagnostic_key = (
+                    lsp_diagnostic.message,
+                    lsp_diagnostic.range.start.line,
+                    lsp_diagnostic.range.start.character,
+                    lsp_diagnostic.range.end.line,
+                    lsp_diagnostic.range.end.character,
+                )
+                if diagnostic_key not in lsp_diagnostics:
+                    lsp_diagnostics[diagnostic_key] = lsp_diagnostic
+        return list(lsp_diagnostics.values())
+
+    @staticmethod
+    def diagnostic_to_lsp_diagnostic(
+        diagnostic: AnnotatedRuleViolation,
+    ) -> t.Optional[types.Diagnostic]:
+        if diagnostic.model._path is None:
+            return None
+        if not diagnostic.violation_range:
+            with open(diagnostic.model._path, "r", encoding="utf-8") as file:
+                lines = file.readlines()
+            diagnostic_range = types.Range(
+                start=types.Position(line=0, character=0),
+                end=types.Position(line=len(lines) - 1, character=len(lines[-1])),
+            )
+        else:
+            diagnostic_range = types.Range(
+                start=types.Position(
+                    line=diagnostic.violation_range.start.line,
+                    character=diagnostic.violation_range.start.character,
+                ),
+                end=types.Position(
+                    line=diagnostic.violation_range.end.line,
+                    character=diagnostic.violation_range.end.character,
+                ),
+            )
+
+        # Get rule definition location for diagnostics link
+        rule_location = diagnostic.rule.get_definition_location()
+        rule_uri_wihout_extension = URI.from_path(rule_location.file_path)
+        rule_uri = f"{rule_uri_wihout_extension.value}#L{rule_location.start_line}"
+
+        # Use URI format to create a link for "related information"
+        return types.Diagnostic(
+            range=diagnostic_range,
+            message=diagnostic.violation_msg,
+            severity=types.DiagnosticSeverity.Error
+            if diagnostic.violation_type == "error"
+            else types.DiagnosticSeverity.Warning,
+            source="sqlmesh",
+            code=diagnostic.rule.name,
+            code_description=types.CodeDescription(href=rule_uri),
+        )