feat: add linting to the lsp

benfdking · benfdking · commit d1c54934ce61 · 2025-04-16T16:28:06.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=True,
+        rules=[
+            "ambiguousorinvalidcolumn",
+            "invalidselectstarexpansion",
+            "noselectstar",
+            "nomissingaudits",
+        ],
+    ),
 )
 
 bigquery_config = Config(
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
 
 
@@ -31,14 +32,85 @@ def __init__(
         """
         self.server = LanguageServer(server_name, version)
         self.context_class = context_class
-        self.context: t.Optional[Context] = None
+        # A tuple of (context, reverse_map) where the reverse_map is uri to model name
+        self.context_and_reverse_map: t.Optional[t.Tuple[Context, t.Dict[str, str]]] = 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_and_reverse_map is None:
+                self.ensure_context_for_document(
+                    ls.workspace.get_document(params.text_document.uri)
+                )
+            if self.context_and_reverse_map is None:
+                raise RuntimeError("No context found")
+            model = self.context_and_reverse_map[1][params.text_document.uri]
+            if model is None:
+                return
+            self.lint_cache[params.text_document.uri] = self.context_and_reverse_map[0].lint_models(
+                [model],
+                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:
+            if self.context_and_reverse_map is None:
+                self.ensure_context_for_document(
+                    ls.workspace.get_document(params.text_document.uri)
+                )
+            if self.context_and_reverse_map is None:
+                raise RuntimeError("No context found")
+            model = self.context_and_reverse_map[1][params.text_document.uri]
+            if model is None:
+                return
+            self.lint_cache[params.text_document.uri] = self.context_and_reverse_map[0].lint_models(
+                [model],
+                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_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_and_reverse_map is None:
+                self.ensure_context_for_document(
+                    ls.workspace.get_document(params.text_document.uri)
+                )
+            if self.context_and_reverse_map is None:
+                raise RuntimeError("No context found")
+            model = self.context_and_reverse_map[1][params.text_document.uri]
+            if model is None:
+                return
+            self.lint_cache[params.text_document.uri] = self.context_and_reverse_map[0].lint_models(
+                [model]
+            )
+            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
@@ -49,11 +121,11 @@ def formatting(
                     ls.workspace.get_document(params.text_document.uri)
                 )
 
-                if self.context is None:
+                if self.context_and_reverse_map 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.context_and_reverse_map[0].format(paths=(Path(document.path),))
                 with open(document.path, "r+", encoding="utf-8") as file:
                     new_text = file.read()
 
@@ -80,8 +152,16 @@ def ensure_context_for_document(self, document: TextDocument) -> TextDocument:
         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
+        if self.context_and_reverse_map is not None:
+            context, _ = self.context_and_reverse_map
+            context.load()  # Reload or refresh context
+            self.context_and_reverse_map = (
+                context,
+                {
+                    f"file://{Path(model._path).resolve().as_posix()}": model.name
+                    for model in context._models.values()
+                },
+            )
             return document
 
         # No context yet: try to find config and load it
@@ -97,7 +177,15 @@ 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])
+                        context = self.context_class(paths=[path])
+                        self.context_and_reverse_map = (
+                            context,
+                            {
+                                f"file://{Path(model._path).resolve().as_posix()}": model.name
+                                for model in context._models.values()
+                                if model._path is not None
+                            },
+                        )
                         self.server.show_message(f"Context loaded for: {path}")
                         loaded = True
                         # Re-check context for document now that it's loaded
@@ -110,6 +198,45 @@ 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,
+        )
+
+    @staticmethod
+    def _diagnostics_to_lsp_diagnostics(
+        diagnostics: t.List[RuleViolationWithModelAndType],
+    ) -> t.List[types.Diagnostic]:
+        lsp_diagnostics: t.List[types.Diagnostic] = []
+        for diagnostic in diagnostics:
+            if diagnostic is None:
+                continue
+            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."""
         logging.basicConfig(level=logging.DEBUG)
