fix(vscode): use paths instead of uris (#4448)

Author: benfdking

sqlmesh/lsp/completions.py

@@ -33,8 +33,9 @@ def get_models(context: t.Optional[LSPContext], file_uri: t.Optional[URI]) -> t.
             all_models.update(file_info.names)
 
     # Remove models from the current file
-    if file_uri is not None and file_uri in context.map:
-        file_info = context.map[file_uri]
+    path = file_uri.to_path() if file_uri is not None else None
+    if path is not None and path in context.map:
+        file_info = context.map[path]
         if isinstance(file_info, ModelTarget):
             for model in file_info.names:
                 all_models.discard(model)
@@ -53,8 +54,8 @@ def get_keywords(context: t.Optional[LSPContext], file_uri: t.Optional[URI]) ->
     If both a context and a file_uri are provided, returns the keywords
     for the dialect of the model that the file belongs to.
     """
-    if file_uri is not None and context is not None and file_uri in context.map:
-        file_info = context.map[file_uri]
+    if file_uri is not None and context is not None and file_uri.to_path() in context.map:
+        file_info = context.map[file_uri.to_path()]
 
         # Handle ModelInfo objects
         if isinstance(file_info, ModelTarget) and file_info.names:

sqlmesh/lsp/context.py

@@ -1,9 +1,8 @@
 from dataclasses import dataclass
+from pathlib import Path
 from sqlmesh.core.context import Context
 import typing as t
 
-from sqlmesh.lsp.uri import URI
-
 
 @dataclass
 class ModelTarget:
@@ -29,24 +28,24 @@ def __init__(self, context: Context) -> None:
         self.context = context
 
         # Add models to the map
-        model_map: t.Dict[URI, ModelTarget] = {}
+        model_map: t.Dict[Path, ModelTarget] = {}
         for model in context.models.values():
             if model._path is not None:
-                uri = URI.from_path(model._path)
+                uri = model._path
                 if uri in model_map:
                     model_map[uri].names.append(model.name)
                 else:
                     model_map[uri] = ModelTarget(names=[model.name])
 
         # Add standalone audits to the map
-        audit_map: t.Dict[URI, AuditTarget] = {}
+        audit_map: t.Dict[Path, AuditTarget] = {}
         for audit in context.standalone_audits.values():
             if audit._path is not None:
-                uri = URI.from_path(audit._path)
+                uri = audit._path
                 if uri not in audit_map:
                     audit_map[uri] = AuditTarget(name=audit.name)
 
-        self.map: t.Dict[URI, t.Union[ModelTarget, AuditTarget]] = {
+        self.map: t.Dict[Path, t.Union[ModelTarget, AuditTarget]] = {
             **model_map,
             **audit_map,
         }

sqlmesh/lsp/main.py

@@ -116,7 +116,7 @@ def did_open(ls: LanguageServer, params: types.DidOpenTextDocumentParams) -> Non
                     SQLMeshLanguageServer._diagnostics_to_lsp_diagnostics(self.lint_cache[uri]),
                 )
                 return
-            models = context.map[uri]
+            models = context.map[uri.to_path()]
             if models is None:
                 return
             if not isinstance(models, ModelTarget):
@@ -134,7 +134,7 @@ def did_open(ls: LanguageServer, params: types.DidOpenTextDocumentParams) -> Non
         def did_change(ls: LanguageServer, params: types.DidChangeTextDocumentParams) -> None:
             uri = URI(params.text_document.uri)
             context = self._context_get_or_load(uri)
-            models = context.map[uri]
+            models = context.map[uri.to_path()]
             if models is None:
                 return
             if not isinstance(models, ModelTarget):
@@ -152,7 +152,7 @@ def did_change(ls: LanguageServer, params: types.DidChangeTextDocumentParams) ->
         def did_save(ls: LanguageServer, params: types.DidSaveTextDocumentParams) -> None:
             uri = URI(params.text_document.uri)
             context = self._context_get_or_load(uri)
-            models = context.map[uri]
+            models = context.map[uri.to_path()]
             if models is None:
                 return
             if not isinstance(models, ModelTarget):

sqlmesh/lsp/reference.py

@@ -88,15 +88,14 @@ def get_model_definitions_for_a_path(
     - Try get_model before normalization
     - Match to models that the model refers to
     """
-    # Ensure the path is a sql file
-    if document_uri.to_path().suffix != ".sql":
+    path = document_uri.to_path()
+    if path.suffix != ".sql":
         return []
     # Get the file info from the context map
-    if document_uri not in lint_context.map:
+    if path not in lint_context.map:
         return []
 
-    file_info = lint_context.map[document_uri]
-
+    file_info = lint_context.map[path]
     # Process based on whether it's a model or standalone audit
     if isinstance(file_info, ModelTarget):
         # It's a model

tests/lsp/test_completions.py

@@ -3,6 +3,7 @@
 from sqlmesh.core.context import Context
 from sqlmesh.lsp.completions import get_keywords_from_tokenizer, get_sql_completions
 from sqlmesh.lsp.context import LSPContext
+from sqlmesh.lsp.uri import URI
 
 
 TOKENIZER_KEYWORDS = set(Tokenizer.KEYWORDS.keys())
@@ -36,9 +37,7 @@ def test_get_sql_completions_with_context_and_file_uri():
     context = Context(paths=["examples/sushi"])
     lsp_context = LSPContext(context)
 
-    file_uri = next(
-        key for key in lsp_context.map.keys() if str(key.to_path()).endswith("active_customers.sql")
-    )
-    completions = get_sql_completions(lsp_context, file_uri)
+    file_uri = next(key for key in lsp_context.map.keys() if key.name == "active_customers.sql")
+    completions = get_sql_completions(lsp_context, URI.from_path(file_uri))
     assert len(completions.keywords) > len(TOKENIZER_KEYWORDS)
     assert "sushi.active_customers" not in completions.models

tests/lsp/test_context.py

@@ -14,7 +14,7 @@ def test_lsp_context():
 
     # find one model in the map
     active_customers_key = next(
-        key for key in lsp_context.map.keys() if str(key.to_path()).endswith("active_customers.sql")
+        key for key in lsp_context.map.keys() if key.name == "active_customers.sql"
     )
 
     # Check that the value is a ModelInfo with the expected model name

tests/lsp/test_reference.py

@@ -3,6 +3,7 @@
 from sqlmesh.core.context import Context
 from sqlmesh.lsp.context import LSPContext, ModelTarget, AuditTarget
 from sqlmesh.lsp.reference import get_model_definitions_for_a_path, by_position
+from sqlmesh.lsp.uri import URI
 
 
 @pytest.mark.fast
@@ -11,21 +12,22 @@ def test_reference() -> None:
     lsp_context = LSPContext(context)
 
     # Find model URIs
-    active_customers_uri = next(
-        uri
-        for uri, info in lsp_context.map.items()
+    active_customers_path = next(
+        path
+        for path, info in lsp_context.map.items()
         if isinstance(info, ModelTarget) and "sushi.active_customers" in info.names
     )
-    sushi_customers_uri = next(
-        uri
-        for uri, info in lsp_context.map.items()
+    sushi_customers_path = next(
+        path
+        for path, info in lsp_context.map.items()
         if isinstance(info, ModelTarget) and "sushi.customers" in info.names
     )
 
+    active_customers_uri = URI.from_path(active_customers_path)
     references = get_model_definitions_for_a_path(lsp_context, active_customers_uri)
 
     assert len(references) == 1
-    assert references[0].uri == sushi_customers_uri.value
+    assert URI(references[0].uri) == URI.from_path(sushi_customers_path)
 
     # Check that the reference in the correct range is sushi.customers
     path = active_customers_uri.to_path()
@@ -42,17 +44,18 @@ def test_reference_with_alias() -> None:
     context = Context(paths=["examples/sushi"])
     lsp_context = LSPContext(context)
 
-    waiter_revenue_by_day_uri = next(
+    waiter_revenue_by_day_path = next(
         uri
         for uri, info in lsp_context.map.items()
         if isinstance(info, ModelTarget) and "sushi.waiter_revenue_by_day" in info.names
     )
 
-    references = get_model_definitions_for_a_path(lsp_context, waiter_revenue_by_day_uri)
+    references = get_model_definitions_for_a_path(
+        lsp_context, URI.from_path(waiter_revenue_by_day_path)
+    )
     assert len(references) == 3
 
-    path = waiter_revenue_by_day_uri.to_path()
-    with open(path, "r") as file:
+    with open(waiter_revenue_by_day_path, "r") as file:
         read_file = file.readlines()
 
     assert references[0].uri.endswith("orders.py")
@@ -70,27 +73,25 @@ def test_standalone_audit_reference() -> None:
     lsp_context = LSPContext(context)
 
     # Find the standalone audit URI
-    audit_uri = next(
+    audit_path = next(
         uri
         for uri, info in lsp_context.map.items()
         if isinstance(info, AuditTarget) and info.name == "assert_item_price_above_zero"
     )
-
     # Find the items model URI
-    items_uri = next(
+    items_path = next(
         uri
         for uri, info in lsp_context.map.items()
         if isinstance(info, ModelTarget) and "sushi.items" in info.names
     )
 
-    references = get_model_definitions_for_a_path(lsp_context, audit_uri)
+    references = get_model_definitions_for_a_path(lsp_context, URI.from_path(audit_path))
 
     assert len(references) == 1
-    assert references[0].uri == items_uri.value
+    assert references[0].uri == URI.from_path(items_path).value
 
     # Check that the reference in the correct range is sushi.items
-    path = audit_uri.to_path()
-    with open(path, "r") as file:
+    with open(audit_path, "r") as file:
         read_file = file.readlines()
     referenced_text = get_string_from_range(read_file, references[0].range)
     assert referenced_text == "sushi.items"
@@ -123,19 +124,20 @@ def test_filter_references_by_position() -> None:
     lsp_context = LSPContext(context)
 
     # Use a file with multiple references (waiter_revenue_by_day)
-    waiter_revenue_by_day_uri = next(
-        uri
-        for uri, info in lsp_context.map.items()
+    waiter_revenue_by_day_path = next(
+        path
+        for path, info in lsp_context.map.items()
         if isinstance(info, ModelTarget) and "sushi.waiter_revenue_by_day" in info.names
     )
 
     # Get all references in the file
-    all_references = get_model_definitions_for_a_path(lsp_context, waiter_revenue_by_day_uri)
+    all_references = get_model_definitions_for_a_path(
+        lsp_context, URI.from_path(waiter_revenue_by_day_path)
+    )
     assert len(all_references) == 3
 
     # Get file contents to locate positions for testing
-    path = waiter_revenue_by_day_uri.to_path()
-    with open(path, "r") as file:
+    with open(waiter_revenue_by_day_path, "r") as file:
         read_file = file.readlines()
 
     # Test positions for each reference