temp linting

Author: benfdking

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,88 @@ 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: 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)
-                )
-
+                self._ensure_context_for_document(params.text_document.uri)
+                document = ls.workspace.get_document(params.text_document.uri)
                 if self.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.context.context.format(paths=(Path(document.path),))
                 with open(document.path, "r+", encoding="utf-8") as file:
                     new_text = file.read()
 
@@ -70,20 +146,32 @@ 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.context is None:
+            self._ensure_context_for_document(document_uri)
+        if self.context is None:
+            raise RuntimeError("No context found")
+        return self.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
+            context = self.context
+            context.context.load()  # Reload or refresh context
+            self.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 +181,57 @@ 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.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:
+            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."""