temp [ci skip]

Author: benfdking

examples/sushi/models/customers.sql

@@ -32,7 +32,7 @@ WITH current_marketing AS (
 )
 SELECT DISTINCT
   o.customer_id::INT AS customer_id, /* this comment should not be registered */
-  m.status,
+  m.status ,
   d.zip
 FROM sushi.orders AS o
 LEFT JOIN current_marketing AS m

sqlmesh/core/context.py

@@ -1078,11 +1078,19 @@ def format(
     ) -> bool:
         """Format all SQL models and audits."""
         filtered_targets = [
-            target for target in self._models.values() | self._audits.values()
-            if target._path is not None and target._path.suffix == ".sql"
+            target
+            for target in chain(self._models.values(), self._audits.values())
+            if target._path is not None
+            and target._path.suffix == ".sql"
             and (not paths or any(target._path.samefile(p) for p in paths))
         ]
+        unformatted_file_paths = []
+
         for target in filtered_targets:
+            if (
+                target._path is None
+            ):  # introduced to satisfy type checker as still want to pull filter out as many targets as possible before loop
+                continue
             with open(target._path, "r+", encoding="utf-8") as file:
                 before = file.read()
                 expressions = parse(before, default_dialect=self.config_for_node(target).dialect)
@@ -1129,11 +1137,6 @@ def format(
             return False
 
         return True
-    
-        
-
-    
-
 
     @python_api_analytics
     def plan(

sqlmesh/lsp/main.py

@@ -3,7 +3,6 @@
 
 import logging
 import typing as t
-import weakref
 from contextlib import suppress
 from pathlib import Path
 
@@ -12,70 +11,84 @@
 from pygls.server import LanguageServer
 from pygls.workspace import TextDocument
 from sqlmesh._version import __version__
+from sqlmesh.core.dialect import Audit
 from sqlmesh.core.model import Model
 
 logger = logging.getLogger(__name__)
 
-CONTEXTS: t.Dict[str, Context] = {}
-"""A mapping of workspace paths to SQLMesh contexts."""
+GLOBAL_CONTEXT: t.Optional[
+    t.Tuple[
+        Context,
+        t.Dict[str, t.Tuple[Context, t.Union[Model, Audit]]],
+    ]
+] = None
 
-PATHS_TO_MODELS: t.Dict[str, t.Tuple[Context, Model]] = {}
-"""A mapping of file paths to SQLMesh (context, model) tuples."""
 
 server = LanguageServer("sqlmesh_lsp", __version__)
 
-_CACHE: t.Set[str] = set()
-"""A cache of URIs for which we have already ensured a context exists."""
-
 
 def ensure_context_for_document(document: TextDocument) -> TextDocument:
-    """Ensure that a context exists for the given document if applicable."""
-    if document.uri in _CACHE:
-        return document
-    _CACHE.add(document.uri)
+    """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, return the document, if it is part of the same context
+    global GLOBAL_CONTEXT
+    if GLOBAL_CONTEXT is not None:
+        GLOBAL_CONTEXT[0].load()
+        if document.uri in GLOBAL_CONTEXT[1]:
+            return document
+        else:
+            for model in GLOBAL_CONTEXT[0]._models.values():
+                if model._path is None:
+                    continue
+                path = model._path.resolve()
+                if path == document.path:
+                    GLOBAL_CONTEXT = (
+                        GLOBAL_CONTEXT[0],
+                        GLOBAL_CONTEXT[1] | {document.uri: (GLOBAL_CONTEXT[0], model)},
+                    )
+                    return document
+            for audit in GLOBAL_CONTEXT[0]._audits.values():
+                if audit._path is None:
+                    continue
+                path = audit._path.resolve()
+                if path == document.path:
+                    GLOBAL_CONTEXT = (
+                        GLOBAL_CONTEXT[0],
+                        GLOBAL_CONTEXT[1] | {document.uri: (GLOBAL_CONTEXT[0], audit)},
+                    )
+                    return document
+            return document
+
+    # If there is no context, load the context and then call this function again
     path = Path(document.path).resolve()
     if path.suffix not in (".sql", ".py"):
         return document
-    initial_path = path
-    while path.parents:
-        if str(path) in CONTEXTS:
-            return document
-        path = path.parent
-    path = initial_path
     loaded = False
     while path.parents and not loaded:
         for ext in ("py", "yml", "yaml"):
             config_path = path / f"config.{ext}"
             if config_path.exists():
                 with suppress(Exception):
-                    handle = Context(paths=[f"{path}"])
-                    CONTEXTS[str(path)] = handle
-                    PATHS_TO_MODELS.update(
-                        {
-                            str(model._path.resolve()): (handle, weakref.proxy(model))
-                            for model in handle.models.values()
-                        }
-                    )
+                    handle = Context(paths=[path])
+                    GLOBAL_CONTEXT = (handle, {})
                     server.show_message(f"Context loaded for: {path}")
                     loaded = True
-                    break
+                    return ensure_context_for_document(document)
         path = path.parent
+
     return document
 
 
 @server.feature(types.TEXT_DOCUMENT_FORMATTING)
 def formatting(
     ls: LanguageServer, params: types.DocumentFormattingParams
 ) -> t.List[types.TextEdit]:
-    """Format the document based using SQLMesh format_model_expressions."""
+    """Format the document using SQLMesh format_model_expressions."""
     try:
-        logger.info(f"Formatting document: {params.text_document.uri}")
         document = ensure_context_for_document(ls.workspace.get_document(params.text_document.uri))
-        context, _ = PATHS_TO_MODELS.get(document.path, (None, None))
+        context = GLOBAL_CONTEXT
         if context is None:
-            logger.error(f"No context found for document: {document.path}")
-            return []
-        context.format(paths=(Path(document.path),))
+            raise Exception(f"No context found for document: {document.path}")
+        context[0].format(paths=(Path(document.path),))
         with open(document.path, "r+", encoding="utf-8") as file:
             return [
                 types.TextEdit(