Update sqlmesh/lsp/main.py

Co-authored-by: Jo <46752250+georgesittas@users.noreply.github.com>

Author: benfdking

sqlmesh/lsp/main.py

@@ -15,61 +15,65 @@
 
 logger = logging.getLogger(__name__)
 
-CONTEXTS: t.Dict[str, Context] = {}
-"""A mapping of workspace paths to SQLMesh contexts."""
+CONTEXT: t.Optional[
+    t.Tuple[
+        Context,
+        t.Dict[str, t.Tuple[Context, Model]],
+    ]
+] = 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 by searching for a config.py or config.yml file in the parent directories."""
-    if document.uri in _CACHE:
-        return document
-    _CACHE.add(document.uri)
+    # If the context is already loaded, return the document, if it is part of the same context
+    if CONTEXT is not None:
+        CONTEXT[0].reload()
+        if document.uri in CONTEXT[1]:
+            return document
+        else:
+            for model in CONTEXT[1]._models:
+                path = model._path.resolve()
+                if path == document.path:
+                    CONTEXT = (CONTEXT[0], CONTEXT[1] | {document.uri: (CONTEXT[0], model)})
+                    return document
+            for audit in CONTEXT[1]._audits:
+                path = audit._path.resolve()
+                if path == document.path:
+                    CONTEXT = (CONTEXT[0], CONTEXT[1] | {document.uri: (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, model)
-                            for model in handle.models.values()
-                        }
-                    )
+                    handle = Context(paths=[path])
+                    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:
         document = ensure_context_for_document(ls.workspace.get_document(params.text_document.uri))
-        context, _ = PATHS_TO_MODELS.get(document.path, (None, None))
+        context = CONTEXT
         if context is None:
             raise Exception(f"No context found for document: {document.path}")
         context.format(paths=(Path(document.path),))