feat: add linting to the lsp

benfdking · benfdking · commit f8c517a45159 · 2025-04-16T15:53:19.000+01:00
diff --git a/sqlmesh/core/context.py b/sqlmesh/core/context.py
@@ -75,6 +75,7 @@
 )
 from sqlmesh.core.engine_adapter import EngineAdapter
 from sqlmesh.core.environment import Environment, EnvironmentNamingInfo, EnvironmentStatements
+from sqlmesh.core.linter.rule import RuleViolationWithModelAndType
 from sqlmesh.core.loader import Loader
 from sqlmesh.core.linter.definition import Linter
 from sqlmesh.core.linter.rules import BUILTIN_RULES
@@ -2460,7 +2461,8 @@ def _get_models_for_interval_end(
     def lint_models(
         self,
         models: t.Optional[t.Iterable[t.Union[str, Model]]] = None,
-    ) -> None:
+        raise_on_error: bool = True,
+    ) -> t.List[RuleViolationWithModelAndType]:
         found_error = False
 
         model_list = (
@@ -2469,13 +2471,17 @@ def lint_models(
         for model in model_list:
             # Linter may be `None` if the context is not loaded yet
             if linter := self._linters.get(model.project):
-                found_error = linter.lint_model(model, console=self.console) or found_error
+                found_error, all_violations = (
+                    linter.lint_model(model, console=self.console) or found_error
+                )
 
-        if found_error:
+        if raise_on_error and found_error:
             raise LinterError(
                 "Linter detected errors in the code. Please fix them before proceeding."
             )
 
+        return all_violations
+
     def load_model_tests(
         self, tests: t.Optional[t.List[str]] = None, patterns: list[str] | None = None
     ) -> t.List[ModelTestMetadata]:
diff --git a/sqlmesh/core/linter/definition.py b/sqlmesh/core/linter/definition.py
@@ -8,7 +8,7 @@
 
 from sqlmesh.utils.errors import raise_config_error
 from sqlmesh.core.console import LinterConsole, get_console, Console
-from sqlmesh.core.linter.rule import RuleSet
+from sqlmesh.core.linter.rule import RuleSet, RuleViolationWithModelAndType
 
 
 def select_rules(all_rules: RuleSet, rule_names: t.Set[str]) -> RuleSet:
@@ -50,9 +50,11 @@ def from_rules(cls, all_rules: RuleSet, config: LinterConfig) -> Linter:
 
         return Linter(config.enabled, all_rules, rules, warn_rules)
 
-    def lint_model(self, model: Model, console: LinterConsole = get_console()) -> bool:
+    def lint_model(
+        self, model: Model, console: LinterConsole = get_console()
+    ) -> t.Tuple[bool, t.List[RuleViolationWithModelAndType]]:
         if not self.enabled:
-            return False
+            return False, []
 
         ignored_rules = select_rules(self.all_rules, model.ignored_rules)
 
@@ -62,11 +64,28 @@ def lint_model(self, model: Model, console: LinterConsole = get_console()) -> bo
         error_violations = rules.check_model(model)
         warn_violations = warn_rules.check_model(model)
 
+        all_violations: t.List[RuleViolationWithModelAndType] = [
+            RuleViolationWithModelAndType(
+                rule=violation.rule,
+                violation_msg=violation.violation_msg,
+                model=model,
+                violation_type="error",
+            )
+            for violation in error_violations
+        ] + [
+            RuleViolationWithModelAndType(
+                rule=violation.rule,
+                violation_msg=violation.violation_msg,
+                model=model,
+                violation_type="warning",
+            )
+            for violation in warn_violations
+        ]
+
         if warn_violations:
             console.show_linter_violations(warn_violations, model)
-
         if error_violations:
             console.show_linter_violations(error_violations, model, is_error=True)
-            return True
+            return True, all_violations
 
-        return False
+        return False, all_violations
diff --git a/sqlmesh/core/linter/rule.py b/sqlmesh/core/linter/rule.py
@@ -50,6 +50,19 @@ def __repr__(self) -> str:
         return f"{self.rule.name}: {self.violation_msg}"
 
 
+class RuleViolationWithModelAndType(RuleViolation):
+    def __init__(
+        self,
+        rule: Rule,
+        violation_msg: str,
+        model: Model,
+        violation_type: t.Literal["error", "warning"],
+    ) -> None:
+        super().__init__(rule, violation_msg)
+        self.model = model
+        self.violation_type = violation_type
+
+
 class RuleSet(Mapping[str, type[Rule]]):
     def __init__(self, rules: Iterable[type[Rule]] = ()) -> None:
         self._underlying = {rule.name: rule for rule in rules}
diff --git a/sqlmesh/lsp/main.py b/sqlmesh/lsp/main.py
@@ -14,6 +14,7 @@
 from sqlmesh._version import __version__
 from sqlmesh.core.audit.definition import ModelAudit
 from sqlmesh.core.context import Context
+from sqlmesh.core.linter.rule import RuleViolationWithModelAndType
 from sqlmesh.core.model import Model
 
 
@@ -32,13 +33,78 @@ def __init__(
         self.server = LanguageServer(server_name, version)
         self.context_class = context_class
         self.context: t.Optional[Context] = None
+        self.lint_cache: t.Dict[str, t.List[RuleViolationWithModelAndType]] = {}
 
         # Register LSP features (e.g., formatting, hover, etc.)
         self._register_features()
 
     def _register_features(self) -> None:
         """Register LSP features on the internal LanguageServer instance."""
 
+        @self.server.feature(types.TEXT_DOCUMENT_DID_OPEN)
+        def did_open(ls: LanguageServer, params: types.DidOpenTextDocumentParams) -> None:
+            if self.context is None:
+                self.ensure_context_for_document(
+                    ls.workspace.get_document(params.text_document.uri)
+                )
+            if self.context is None:
+                raise RuntimeError("No context found")
+            self.lint_cache[params.text_document.uri] = self.context.lint_models(
+                [params.text_document.uri]
+            )
+            ls.publish_diagnostics(
+                params.text_document.uri,
+                [
+                    self._diagnostic_to_lsp_diagnostic(diagnostic)
+                    for diagnostic in self.lint_cache[params.text_document.uri]
+                    if diagnostic is not None
+                ],
+            )
+
+        @self.server.feature(types.TEXT_DOCUMENT_DID_CHANGE)
+        def did_change(ls: LanguageServer, params: types.DidChangeTextDocumentParams) -> None:
+            if self.context is None:
+                self.ensure_context_for_document(
+                    ls.workspace.get_document(params.text_document.uri)
+                )
+            if self.context is None:
+                raise RuntimeError("No context found")
+            self.lint_cache[params.text_document.uri] = self.context.lint_models(
+                [params.text_document.uri]
+            )
+            ls.publish_diagnostics(
+                params.text_document.uri,
+                [
+                    self._diagnostic_to_lsp_diagnostic(diagnostic)
+                    for diagnostic in self.lint_cache[params.text_document.uri]
+                    if diagnostic is not None
+                ],
+            )
+
+        @self.server.feature(types.TEXT_DOCUMENT_DID_CLOSE)
+        def did_close(ls: LanguageServer, params: types.DidCloseTextDocumentParams) -> None:
+            self.lint_cache.pop(params.text_document.uri, None)
+
+        @self.server.feature(types.TEXT_DOCUMENT_DID_SAVE)
+        def did_save(ls: LanguageServer, params: types.DidSaveTextDocumentParams) -> None:
+            if self.context is None:
+                self.ensure_context_for_document(
+                    ls.workspace.get_document(params.text_document.uri)
+                )
+            if self.context is None:
+                raise RuntimeError("No context found")
+            self.lint_cache[params.text_document.uri] = self.context.lint_models(
+                [params.text_document.uri]
+            )
+            ls.publish_diagnostics(
+                params.text_document.uri,
+                [
+                    self._diagnostic_to_lsp_diagnostic(diagnostic)
+                    for diagnostic in self.lint_cache[params.text_document.uri]
+                    if diagnostic is not None
+                ],
+            )
+
         @self.server.feature(types.TEXT_DOCUMENT_FORMATTING)
         def formatting(
             ls: LanguageServer, params: types.DocumentFormattingParams
@@ -110,6 +176,25 @@ def ensure_context_for_document(self, document: TextDocument) -> TextDocument:
 
         return document
 
+    @staticmethod
+    def _diagnostic_to_lsp_diagnostic(
+        diagnostic: RuleViolationWithModelAndType,
+    ) -> t.Optional[types.Diagnostic]:
+        if diagnostic.model._path is None:
+            return None
+        with open(diagnostic.model._path, "r", encoding="utf-8") as file:
+            lines = file.readlines()
+        return types.Diagnostic(
+            range=types.Range(
+                start=types.Position(line=0, character=0),
+                end=types.Position(line=len(lines), character=len(lines[-1])),
+            ),
+            message=diagnostic.violation_msg,
+            severity=types.DiagnosticSeverity.Error
+            if diagnostic.violation_type == "error"
+            else types.DiagnosticSeverity.Warning,
+        )
+
     def start(self) -> None:
         """Start the server with I/O transport."""
         logging.basicConfig(level=logging.DEBUG)
