fix(lsp): overwriting the formatted file (#4547)

Author: benfdking

sqlmesh/core/context.py

@@ -1129,33 +1129,15 @@ def format(
 
             with open(target._path, "r+", encoding="utf-8") as file:
                 before = file.read()
-                expressions = parse(before, default_dialect=self.config_for_node(target).dialect)
-                if transpile and is_meta_expression(expressions[0]):
-                    for prop in expressions[0].expressions:
-                        if prop.name.lower() == "dialect":
-                            prop.replace(
-                                exp.Property(
-                                    this="dialect",
-                                    value=exp.Literal.string(transpile or target.dialect),
-                                )
-                            )
-
-                format_config = self.config_for_node(target).format
-                after = format_model_expressions(
-                    expressions,
-                    transpile or target.dialect,
-                    rewrite_casts=(
-                        rewrite_casts
-                        if rewrite_casts is not None
-                        else not format_config.no_rewrite_casts
-                    ),
-                    **{**format_config.generator_options, **kwargs},
-                )
 
-                if append_newline is None:
-                    append_newline = format_config.append_newline
-                if append_newline:
-                    after += "\n"
+                after = self._format(
+                    target,
+                    before,
+                    transpile=transpile,
+                    rewrite_casts=rewrite_casts,
+                    append_newline=append_newline,
+                    **kwargs,
+                )
 
                 if not check:
                     file.seek(0)
@@ -1174,6 +1156,44 @@ def format(
 
         return True
 
+    def _format(
+        self,
+        target: Model | Audit,
+        before: str,
+        *,
+        transpile: t.Optional[str] = None,
+        rewrite_casts: t.Optional[bool] = None,
+        append_newline: t.Optional[bool] = None,
+        **kwargs: t.Any,
+    ) -> str:
+        expressions = parse(before, default_dialect=self.config_for_node(target).dialect)
+        if transpile and is_meta_expression(expressions[0]):
+            for prop in expressions[0].expressions:
+                if prop.name.lower() == "dialect":
+                    prop.replace(
+                        exp.Property(
+                            this="dialect",
+                            value=exp.Literal.string(transpile or target.dialect),
+                        )
+                    )
+
+        format_config = self.config_for_node(target).format
+        after = format_model_expressions(
+            expressions,
+            transpile or target.dialect,
+            rewrite_casts=(
+                rewrite_casts if rewrite_casts is not None else not format_config.no_rewrite_casts
+            ),
+            **{**format_config.generator_options, **kwargs},
+        )
+
+        if append_newline is None:
+            append_newline = format_config.append_newline
+        if append_newline:
+            after += "\n"
+
+        return after
+
     @python_api_analytics
     def plan(
         self,

sqlmesh/lsp/main.py

@@ -1,6 +1,7 @@
 #!/usr/bin/env python
 """A Language Server Protocol (LSP) server for SQL with SQLMesh integration, refactored without globals."""
 
+from itertools import chain
 import logging
 import typing as t
 from pathlib import Path
@@ -20,7 +21,11 @@
     ApiResponseGetModels,
 )
 from sqlmesh.lsp.completions import get_sql_completions
-from sqlmesh.lsp.context import LSPContext, ModelTarget, render_model as render_model_context
+from sqlmesh.lsp.context import (
+    LSPContext,
+    ModelTarget,
+    render_model as render_model_context,
+)
 from sqlmesh.lsp.custom import (
     ALL_MODELS_FEATURE,
     RENDER_MODEL_FEATURE,
@@ -213,15 +218,29 @@ def formatting(
                 uri = URI(params.text_document.uri)
                 self._ensure_context_for_document(uri)
                 document = ls.workspace.get_text_document(params.text_document.uri)
+                before = document.source
                 if self.lsp_context is None:
                     raise RuntimeError(f"No context found for document: {document.path}")
 
-                # Perform formatting using the loaded context
-                self.lsp_context.context.format(paths=(str(uri.to_path()),))
-                with open(uri.to_path(), "r+", encoding="utf-8") as file:
-                    new_text = file.read()
-
-                # Return a single edit that replaces the entire file.
+                target = next(
+                    (
+                        target
+                        for target in chain(
+                            self.lsp_context.context._models.values(),
+                            self.lsp_context.context._audits.values(),
+                        )
+                        if target._path is not None
+                        and target._path.suffix == ".sql"
+                        and (target._path.samefile(uri.to_path()))
+                    ),
+                    None,
+                )
+                if target is None:
+                    return []
+                after = self.lsp_context.context._format(
+                    target=target,
+                    before=before,
+                )
                 return [
                     types.TextEdit(
                         range=types.Range(
@@ -231,7 +250,7 @@ def formatting(
                                 character=len(document.lines[-1]) if document.lines else 0,
                             ),
                         ),
-                        new_text=new_text,
+                        new_text=after,
                     )
                 ]
             except Exception as e: