refactor(lsp): refactor lsp calls into context (#4616)

benfdking · web-flow · commit 1da137a0bb91 · 2025-06-02T16:58:10.000+01:00
diff --git a/sqlmesh/lsp/context.py b/sqlmesh/lsp/context.py
@@ -4,6 +4,7 @@
 import typing as t
 
 from sqlmesh.core.model.definition import SqlModel
+from sqlmesh.core.linter.definition import AnnotatedRuleViolation
 from sqlmesh.lsp.custom import RenderModelEntry
 from sqlmesh.lsp.uri import URI
 
@@ -28,8 +29,14 @@ class LSPContext:
     model names and standalone audit names.
     """
 
+    map: t.Dict[Path, t.Union[ModelTarget, AuditTarget]]
+    _render_cache: t.Dict[Path, t.List[RenderModelEntry]]
+    _lint_cache: t.Dict[Path, t.List[AnnotatedRuleViolation]]
+
     def __init__(self, context: Context) -> None:
         self.context = context
+        self._render_cache = {}
+        self._lint_cache = {}
 
         # Add models to the map
         model_map: t.Dict[Path, ModelTarget] = {}
@@ -54,36 +61,90 @@ def __init__(self, context: Context) -> None:
             **audit_map,
         }
 
+    def render_model(self, uri: URI) -> t.List[RenderModelEntry]:
+        """Get rendered models for a file, using cache when available.
 
-def render_model(context: LSPContext, uri: URI) -> t.Iterator[RenderModelEntry]:
-    target = context.map[uri.to_path()]
-    if isinstance(target, AuditTarget):
-        audit = context.context.standalone_audits[target.name]
-        definition = audit.render_definition(
-            include_python=False,
-            render_query=True,
-        )
-        rendered_query = [render.sql(dialect=audit.dialect, pretty=True) for render in definition]
-        yield RenderModelEntry(
-            name=audit.name,
-            fqn=audit.fqn,
-            description=audit.description,
-            rendered_query="\n\n".join(rendered_query),
-        )
-    if isinstance(target, ModelTarget):
-        for name in target.names:
-            model = context.context.get_model(name)
-            if isinstance(model, SqlModel):
-                rendered_query = [
-                    render.sql(dialect=model.dialect, pretty=True)
-                    for render in model.render_definition(
-                        include_python=False,
-                        render_query=True,
+        Args:
+            uri: The URI of the file to render.
+
+        Returns:
+            List of rendered model entries.
+        """
+        path = uri.to_path()
+
+        # Check cache first
+        if path in self._render_cache:
+            return self._render_cache[path]
+
+        # If not cached, render and cache
+        entries: t.List[RenderModelEntry] = []
+        target = self.map.get(path)
+
+        if isinstance(target, AuditTarget):
+            audit = self.context.standalone_audits[target.name]
+            definition = audit.render_definition(
+                include_python=False,
+                render_query=True,
+            )
+            rendered_query = [
+                render.sql(dialect=audit.dialect, pretty=True) for render in definition
+            ]
+            entry = RenderModelEntry(
+                name=audit.name,
+                fqn=audit.fqn,
+                description=audit.description,
+                rendered_query="\n\n".join(rendered_query),
+            )
+            entries.append(entry)
+
+        elif isinstance(target, ModelTarget):
+            for name in target.names:
+                model = self.context.get_model(name)
+                if isinstance(model, SqlModel):
+                    rendered_query = [
+                        render.sql(dialect=model.dialect, pretty=True)
+                        for render in model.render_definition(
+                            include_python=False,
+                            render_query=True,
+                        )
+                    ]
+                    entry = RenderModelEntry(
+                        name=model.name,
+                        fqn=model.fqn,
+                        description=model.description,
+                        rendered_query="\n\n".join(rendered_query),
                     )
-                ]
-                yield RenderModelEntry(
-                    name=model.name,
-                    fqn=model.fqn,
-                    description=model.description,
-                    rendered_query="\n\n".join(rendered_query),
-                )
+                    entries.append(entry)
+
+        # Store in cache
+        self._render_cache[path] = entries
+        return entries
+
+    def lint_model(self, uri: URI) -> t.List[AnnotatedRuleViolation]:
+        """Get lint diagnostics for a model, using cache when available.
+
+        Args:
+            uri: The URI of the file to lint.
+
+        Returns:
+            List of annotated rule violations.
+        """
+        path = uri.to_path()
+
+        # Check cache first
+        if path in self._lint_cache:
+            return self._lint_cache[path]
+
+        # If not cached, lint and cache
+        target = self.map.get(path)
+        if target is None or not isinstance(target, ModelTarget):
+            return []
+
+        diagnostics = self.context.lint_models(
+            target.names,
+            raise_on_error=False,
+        )
+
+        # Store in cache
+        self._lint_cache[path] = diagnostics
+        return diagnostics
diff --git a/sqlmesh/lsp/main.py b/sqlmesh/lsp/main.py
@@ -24,7 +24,6 @@
 from sqlmesh.lsp.context import (
     LSPContext,
     ModelTarget,
-    render_model as render_model_context,
 )
 from sqlmesh.lsp.custom import (
     ALL_MODELS_FEATURE,
@@ -58,10 +57,6 @@ def __init__(
         self.context_class = context_class
         self.lsp_context: t.Optional[LSPContext] = None
 
-        # Cache stores tuples of (diagnostics, diagnostic_version)
-        self.lint_cache: t.Dict[URI, t.Tuple[t.List[AnnotatedRuleViolation], int]] = {}
-        self._diagnostic_version_counter: int = 0
-
         self.client_supports_pull_diagnostics = False
         # Register LSP features (e.g., formatting, hover, etc.)
         self._register_features()
@@ -120,7 +115,7 @@ def all_models(ls: LanguageServer, params: AllModelsRequest) -> AllModelsRespons
         def render_model(ls: LanguageServer, params: RenderModelRequest) -> RenderModelResponse:
             uri = URI(params.textDocumentUri)
             context = self._context_get_or_load(uri)
-            return RenderModelResponse(models=list(render_model_context(context, uri)))
+            return RenderModelResponse(models=context.render_model(uri))
 
         @self.server.feature(API_FEATURE)
         def api(ls: LanguageServer, request: ApiRequest) -> t.Dict[str, t.Any]:
@@ -173,17 +168,11 @@ def did_open(ls: LanguageServer, params: types.DidOpenTextDocumentParams) -> Non
             if models is None or not isinstance(models, ModelTarget):
                 return
 
-            if self.lint_cache.get(uri) is None:
-                diagnostics = context.context.lint_models(
-                    models.names,
-                    raise_on_error=False,
-                )
-                self._diagnostic_version_counter += 1
-                self.lint_cache[uri] = (diagnostics, self._diagnostic_version_counter)
+            # Get diagnostics from context (which handles caching)
+            diagnostics = context.lint_model(uri)
 
             # Only publish diagnostics if client doesn't support pull diagnostics
             if not self.client_supports_pull_diagnostics:
-                diagnostics, _ = self.lint_cache[uri]
                 ls.publish_diagnostics(
                     params.text_document.uri,
                     SQLMeshLanguageServer._diagnostics_to_lsp_diagnostics(diagnostics),
@@ -197,13 +186,8 @@ def did_change(ls: LanguageServer, params: types.DidChangeTextDocumentParams) ->
             if models is None or not isinstance(models, ModelTarget):
                 return
 
-            # Always update the cache
-            diagnostics = context.context.lint_models(
-                models.names,
-                raise_on_error=False,
-            )
-            self._diagnostic_version_counter += 1
-            self.lint_cache[uri] = (diagnostics, self._diagnostic_version_counter)
+            # Get diagnostics from context (which handles caching)
+            diagnostics = context.lint_model(uri)
 
             # Only publish diagnostics if client doesn't support pull diagnostics
             if not self.client_supports_pull_diagnostics:
@@ -220,13 +204,8 @@ def did_save(ls: LanguageServer, params: types.DidSaveTextDocumentParams) -> Non
             if models is None or not isinstance(models, ModelTarget):
                 return
 
-            # Always update the cache
-            diagnostics = context.context.lint_models(
-                models.names,
-                raise_on_error=False,
-            )
-            self._diagnostic_version_counter += 1
-            self.lint_cache[uri] = (diagnostics, self._diagnostic_version_counter)
+            # Get diagnostics from context (which handles caching)
+            diagnostics = context.lint_model(uri)
 
             # Only publish diagnostics if client doesn't support pull diagnostics
             if not self.client_supports_pull_diagnostics:
@@ -445,31 +424,16 @@ def workspace_diagnostic(
                 return types.WorkspaceDiagnosticReport(items=[])
 
     def _get_diagnostics_for_uri(self, uri: URI) -> t.Tuple[t.List[types.Diagnostic], int]:
-        """Get diagnostics for a specific URI, returning (diagnostics, result_id)."""
-        # Check if we have cached diagnostics
-        if uri in self.lint_cache:
-            diagnostics, result_id = self.lint_cache[uri]
-            return SQLMeshLanguageServer._diagnostics_to_lsp_diagnostics(diagnostics), result_id
+        """Get diagnostics for a specific URI, returning (diagnostics, result_id).
 
-        # Try to get diagnostics by loading context and linting
+        Since we no longer track version numbers, we always return 0 as the result_id.
+        This means pull diagnostics will always fetch fresh results.
+        """
         try:
             context = self._context_get_or_load(uri)
-            models = context.map[uri.to_path()]
-            if models is None or not isinstance(models, ModelTarget):
-                return [], 0
-
-            # Lint the models and cache the results
-            diagnostics = context.context.lint_models(
-                models.names,
-                raise_on_error=False,
-            )
-            self._diagnostic_version_counter += 1
-            self.lint_cache[uri] = (diagnostics, self._diagnostic_version_counter)
-            return SQLMeshLanguageServer._diagnostics_to_lsp_diagnostics(
-                diagnostics
-            ), self._diagnostic_version_counter
+            diagnostics = context.lint_model(uri)
+            return SQLMeshLanguageServer._diagnostics_to_lsp_diagnostics(diagnostics), 0
         except Exception:
-            # If we can't get diagnostics, return empty list with no result ID
             return [], 0
 
     def _context_get_or_load(self, document_uri: URI) -> LSPContext:
@@ -523,7 +487,7 @@ def _ensure_context_for_document(
                         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
+                        # Re-check context for the document now that it's loaded
                         return self._ensure_context_for_document(document_uri)
                     except Exception as e:
                         self.server.show_message(
