feat: publish linter errors through lsp (#4152)

benfdking · web-flow · commit bb73daf40d24 · 2025-04-17T13:45:02.000+01:00
diff --git a/examples/sushi/config.py b/examples/sushi/config.py
@@ -13,6 +13,7 @@
     PlanConfig,
     SparkConnectionConfig,
 )
+from sqlmesh.core.config.linter import LinterConfig
 from sqlmesh.core.notification_target import (
     BasicSMTPNotificationTarget,
     SlackApiNotificationTarget,
@@ -41,6 +42,15 @@
     },
     default_gateway="duckdb",
     model_defaults=model_defaults,
+    linter=LinterConfig(
+        enabled=False,
+        rules=[
+            "ambiguousorinvalidcolumn",
+            "invalidselectstarexpansion",
+            "noselectstar",
+            "nomissingaudits",
+        ],
+    ),
 )
 
 bigquery_config = Config(
diff --git a/sqlmesh/core/context.py b/sqlmesh/core/context.py
@@ -76,7 +76,7 @@
 from sqlmesh.core.engine_adapter import EngineAdapter
 from sqlmesh.core.environment import Environment, EnvironmentNamingInfo, EnvironmentStatements
 from sqlmesh.core.loader import Loader
-from sqlmesh.core.linter.definition import Linter
+from sqlmesh.core.linter.definition import AnnotatedRuleViolation, Linter
 from sqlmesh.core.linter.rules import BUILTIN_RULES
 from sqlmesh.core.macros import ExecutableOrMacro, macro
 from sqlmesh.core.metric import Metric, rewrite
@@ -2466,22 +2466,31 @@ 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[AnnotatedRuleViolation]:
         found_error = False
 
         model_list = (
             list(self.get_model(model) for model in models) if models else self.models.values()
         )
+        all_violations = []
         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
+                lint_violation, violations = (
+                    linter.lint_model(model, console=self.console) or found_error
+                )
+                if lint_violation:
+                    found_error = True
+                all_violations.extend(violations)
 
-        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
@@ -9,6 +9,7 @@
 from functools import reduce
 from sqlmesh.core.model import Model
 from sqlmesh.core.linter.rule import Rule, RuleViolation
+from sqlmesh.core.console import LinterConsole, get_console
 
 
 def select_rules(all_rules: RuleSet, rule_names: t.Set[str]) -> RuleSet:
@@ -50,9 +51,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[AnnotatedRuleViolation]]:
         if not self.enabled:
-            return False
+            return False, []
 
         ignored_rules = select_rules(self.all_rules, model.ignored_rules)
 
@@ -62,14 +65,31 @@ 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[AnnotatedRuleViolation] = [
+            AnnotatedRuleViolation(
+                rule=violation.rule,
+                violation_msg=violation.violation_msg,
+                model=model,
+                violation_type="error",
+            )
+            for violation in error_violations
+        ] + [
+            AnnotatedRuleViolation(
+                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
 
 
 class RuleSet(Mapping[str, type[Rule]]):
@@ -117,3 +137,16 @@ def intersection(self, *others: RuleSet) -> RuleSet:
 
     def difference(self, *others: RuleSet) -> RuleSet:
         return reduce(lambda lhs, rhs: lhs.__op(op.sub, rhs), (self, *others))
+
+
+class AnnotatedRuleViolation(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
diff --git a/sqlmesh/core/linter/rule.py b/sqlmesh/core/linter/rule.py
@@ -2,7 +2,6 @@
 
 import abc
 
-
 from sqlmesh.core.model import Model
 
 from typing import Type
diff --git a/sqlmesh/lsp/main.py b/sqlmesh/lsp/main.py
@@ -1,16 +1,33 @@
 #!/usr/bin/env python
 """A Language Server Protocol (LSP) server for SQL with SQLMesh integration, refactored without globals."""
 
+from collections import defaultdict
 import logging
 import typing as t
 from pathlib import Path
 
 from lsprotocol import types
 from pygls.server import LanguageServer
-from pygls.workspace import TextDocument
 
 from sqlmesh._version import __version__
 from sqlmesh.core.context import Context
+from sqlmesh.core.linter.definition import AnnotatedRuleViolation
+
+
+class LSPContext:
+    """
+    A context that is used for linting. It contains the context and a reverse map of file uri to model names .
+    """
+
+    def __init__(self, context: Context) -> None:
+        self.context = context
+        map: t.Dict[str, t.List[str]] = defaultdict(list[str])
+        for model in context.models.values():
+            if model._path is None:
+                path = Path(model._path).resolve()
+                map[f"file://{path.as_posix()}"].append(model.name)
+
+        self.map = map
 
 
 class SQLMeshLanguageServer:
@@ -27,29 +44,87 @@ def __init__(
         """
         self.server = LanguageServer(server_name, version)
         self.context_class = context_class
-        self.context: t.Optional[Context] = None
+        self.lsp_context: t.Optional[LSPContext] = None
+        self.lint_cache: t.Dict[str, t.List[AnnotatedRuleViolation]] = {}
 
         # 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:
+            context = self._context_get_or_load(params.text_document.uri)
+            if self.lint_cache.get(params.text_document.uri) is not None:
+                ls.publish_diagnostics(
+                    params.text_document.uri,
+                    SQLMeshLanguageServer._diagnostics_to_lsp_diagnostics(
+                        self.lint_cache[params.text_document.uri]
+                    ),
+                )
+                return
+            models = context.map[params.text_document.uri]
+            if models is None:
+                return
+            self.lint_cache[params.text_document.uri] = context.context.lint_models(
+                models,
+                raise_on_error=False,
+            )
+            ls.publish_diagnostics(
+                params.text_document.uri,
+                SQLMeshLanguageServer._diagnostics_to_lsp_diagnostics(
+                    self.lint_cache[params.text_document.uri]
+                ),
+            )
+
+        @self.server.feature(types.TEXT_DOCUMENT_DID_CHANGE)
+        def did_change(ls: LanguageServer, params: types.DidChangeTextDocumentParams) -> None:
+            context = self._context_get_or_load(params.text_document.uri)
+            models = context.map[params.text_document.uri]
+            if models is None:
+                return
+            self.lint_cache[params.text_document.uri] = context.context.lint_models(
+                models,
+                raise_on_error=False,
+            )
+            ls.publish_diagnostics(
+                params.text_document.uri,
+                SQLMeshLanguageServer._diagnostics_to_lsp_diagnostics(
+                    self.lint_cache[params.text_document.uri]
+                ),
+            )
+
+        @self.server.feature(types.TEXT_DOCUMENT_DID_SAVE)
+        def did_save(ls: LanguageServer, params: types.DidSaveTextDocumentParams) -> None:
+            context = self._context_get_or_load(params.text_document.uri)
+            models = context.map[params.text_document.uri]
+            if models is None:
+                return
+            self.lint_cache[params.text_document.uri] = context.context.lint_models(
+                models,
+                raise_on_error=False,
+            )
+            ls.publish_diagnostics(
+                params.text_document.uri,
+                SQLMeshLanguageServer._diagnostics_to_lsp_diagnostics(
+                    self.lint_cache[params.text_document.uri]
+                ),
+            )
+
         @self.server.feature(types.TEXT_DOCUMENT_FORMATTING)
         def formatting(
             ls: LanguageServer, params: types.DocumentFormattingParams
         ) -> t.List[types.TextEdit]:
             """Format the document using SQLMesh `format_model_expressions`."""
             try:
-                document = self.ensure_context_for_document(
-                    ls.workspace.get_document(params.text_document.uri)
-                )
-
-                if self.context is None:
+                self._ensure_context_for_document(params.text_document.uri)
+                document = ls.workspace.get_document(params.text_document.uri)
+                if self.lsp_context is None:
                     raise RuntimeError(f"No context found for document: {document.path}")
 
                 # Perform formatting using the loaded context
-                self.context.format(paths=(Path(document.path),))
+                self.lsp_context.context.format(paths=(Path(document.path),))
                 with open(document.path, "r+", encoding="utf-8") as file:
                     new_text = file.read()
 
@@ -70,20 +145,31 @@ def formatting(
                 ls.show_message(f"Error formatting SQL: {e}", types.MessageType.Error)
                 return []
 
-    def ensure_context_for_document(self, document: TextDocument) -> TextDocument:
+    def _context_get_or_load(self, document_uri: str) -> LSPContext:
+        if self.lsp_context is None:
+            self._ensure_context_for_document(document_uri)
+        if self.lsp_context is None:
+            raise RuntimeError("No context found")
+        return self.lsp_context
+
+    def _ensure_context_for_document(
+        self,
+        document_uri: str,
+    ) -> None:
         """
         Ensure that a context exists for the given document if applicable by searching
         for a config.py or config.yml file in the parent directories.
         """
-        # If the context is already loaded, check if this document belongs to it.
-        if self.context is not None:
-            self.context.load()  # Reload or refresh context
-            return document
+        if self.lsp_context is not None:
+            context = self.lsp_context
+            context.context.load()  # Reload or refresh context
+            self.lsp_context = LSPContext(context.context)
+            return
 
         # No context yet: try to find config and load it
-        path = Path(document.path).resolve()
+        path = Path(self._uri_to_path(document_uri)).resolve()
         if path.suffix not in (".sql", ".py"):
-            return document
+            return
 
         loaded = False
         # Ascend directories to look for config
@@ -93,18 +179,55 @@ def ensure_context_for_document(self, document: TextDocument) -> TextDocument:
                 if config_path.exists():
                     try:
                         # Use user-provided instantiator to build the context
-                        self.context = self.context_class(paths=[path])
-                        self.server.show_message(f"Context loaded for: {path}")
+                        created_context = self.context_class(paths=[path])
+                        self.lsp_context = LSPContext(created_context)
                         loaded = True
                         # Re-check context for document now that it's loaded
-                        return self.ensure_context_for_document(document)
+                        return self._ensure_context_for_document(document_uri)
                     except Exception as e:
                         self.server.show_message(
                             f"Error loading context: {e}", types.MessageType.Error
                         )
             path = path.parent
 
-        return document
+        return
+
+    @staticmethod
+    def _diagnostic_to_lsp_diagnostic(
+        diagnostic: AnnotatedRuleViolation,
+    ) -> 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,
+        )
+
+    @staticmethod
+    def _diagnostics_to_lsp_diagnostics(
+        diagnostics: t.List[AnnotatedRuleViolation],
+    ) -> t.List[types.Diagnostic]:
+        lsp_diagnostics: t.List[types.Diagnostic] = []
+        for diagnostic in diagnostics:
+            lsp_diagnostic = SQLMeshLanguageServer._diagnostic_to_lsp_diagnostic(diagnostic)
+            if lsp_diagnostic is not None:
+                lsp_diagnostics.append(lsp_diagnostic)
+        return lsp_diagnostics
+
+    @staticmethod
+    def _uri_to_path(uri: str) -> str:
+        """Convert a URI to a path."""
+        if uri.startswith("file://"):
+            return Path(uri[7:]).resolve().as_posix()
+        return Path(uri).resolve().as_posix()
 
     def start(self) -> None:
         """Start the server with I/O transport."""
