feat(vscode): go to definition for standalone audits (#4344)

Author: benfdking

.gitignore

@@ -155,3 +155,6 @@ tests/_version.py
 # spark
 metastore_db/
 spark-warehouse/
+
+# claude
+.claude/

sqlmesh/lsp/completions.py

@@ -2,7 +2,7 @@
 from sqlglot import Dialect, Tokenizer
 from sqlmesh.lsp.custom import AllModelsResponse
 import typing as t
-from sqlmesh.lsp.context import LSPContext
+from sqlmesh.lsp.context import AuditTarget, LSPContext, ModelTarget
 
 
 def get_sql_completions(context: t.Optional[LSPContext], file_uri: str) -> AllModelsResponse:
@@ -24,11 +24,20 @@ def get_models(context: t.Optional[LSPContext], file_uri: t.Optional[str]) -> t.
     """
     if context is None:
         return set()
-    all_models = set(model for models in context.map.values() for model in models)
-    if file_uri is not None:
-        models_file_refers_to = context.map[file_uri]
-        for model in models_file_refers_to:
-            all_models.discard(model)
+
+    all_models = set()
+    # Extract model names from ModelInfo objects
+    for file_info in context.map.values():
+        if isinstance(file_info, ModelTarget):
+            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]
+        if isinstance(file_info, ModelTarget):
+            for model in file_info.names:
+                all_models.discard(model)
+
     return all_models
 
 
@@ -43,16 +52,25 @@ def get_keywords(context: t.Optional[LSPContext], file_uri: t.Optional[str]) ->
     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:
-        models = context.map[file_uri]
-        if models:
-            model = models[0]
-            model_from_context = context.context.get_model(model)
-            if model_from_context is not None:
-                if model_from_context.dialect:
-                    return get_keywords_from_tokenizer(model_from_context.dialect)
+    if file_uri is not None and context is not None and file_uri in context.map:
+        file_info = context.map[file_uri]
+
+        # Handle ModelInfo objects
+        if isinstance(file_info, ModelTarget) and file_info.names:
+            model_name = file_info.names[0]
+            model_from_context = context.context.get_model(model_name)
+            if model_from_context is not None and model_from_context.dialect:
+                return get_keywords_from_tokenizer(model_from_context.dialect)
+
+        # Handle AuditInfo objects
+        elif isinstance(file_info, AuditTarget) and file_info.name:
+            audit = context.context.standalone_audits.get(file_info.name)
+            if audit is not None and audit.dialect:
+                return get_keywords_from_tokenizer(audit.dialect)
+
     if context is not None:
         return get_keywords_from_tokenizer(context.context.default_dialect)
+
     return get_keywords_from_tokenizer(None)
 
 

sqlmesh/lsp/context.py

@@ -1,20 +1,53 @@
-from collections import defaultdict
+from dataclasses import dataclass
 from pathlib import Path
 from sqlmesh.core.context import Context
 import typing as t
 
 
+@dataclass
+class ModelTarget:
+    """Information about models in a file."""
+
+    names: t.List[str]
+
+
+@dataclass
+class AuditTarget:
+    """Information about standalone audits in a file."""
+
+    name: str
+
+
 class LSPContext:
     """
-    A context that is used for linting. It contains the context and a reverse map of file uri to model names .
+    A context that is used for linting. It contains the context and a reverse map of file uri to
+    model names and standalone audit names.
     """
 
     def __init__(self, context: Context) -> None:
         self.context = context
-        map: t.Dict[str, t.List[str]] = defaultdict(list)
+
+        # Add models to the map
+        model_map: t.Dict[str, ModelTarget] = {}
         for model in context.models.values():
             if model._path is not None:
                 path = Path(model._path).resolve()
-                map[f"file://{path.as_posix()}"].append(model.name)
+                uri = f"file://{path.as_posix()}"
+                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[str, AuditTarget] = {}
+        for audit in context.standalone_audits.values():
+            if audit._path is not None:
+                path = Path(audit._path).resolve()
+                uri = f"file://{path.as_posix()}"
+                if uri not in audit_map:
+                    audit_map[uri] = AuditTarget(name=audit.name)
 
-        self.map = map
+        self.map: t.Dict[str, t.Union[ModelTarget, AuditTarget]] = {
+            **model_map,
+            **audit_map,
+        }

sqlmesh/lsp/main.py

@@ -12,7 +12,7 @@
 from sqlmesh.core.context import Context
 from sqlmesh.core.linter.definition import AnnotatedRuleViolation
 from sqlmesh.lsp.completions import get_sql_completions
-from sqlmesh.lsp.context import LSPContext
+from sqlmesh.lsp.context import LSPContext, ModelTarget
 from sqlmesh.lsp.custom import ALL_MODELS_FEATURE, AllModelsRequest, AllModelsResponse
 from sqlmesh.lsp.reference import get_model_definitions_for_a_path
 
@@ -91,8 +91,10 @@ def did_open(ls: LanguageServer, params: types.DidOpenTextDocumentParams) -> Non
             models = context.map[params.text_document.uri]
             if models is None:
                 return
+            if not isinstance(models, ModelTarget):
+                return
             self.lint_cache[params.text_document.uri] = context.context.lint_models(
-                models,
+                models.names,
                 raise_on_error=False,
             )
             ls.publish_diagnostics(
@@ -108,8 +110,10 @@ def did_change(ls: LanguageServer, params: types.DidChangeTextDocumentParams) ->
             models = context.map[params.text_document.uri]
             if models is None:
                 return
+            if not isinstance(models, ModelTarget):
+                return
             self.lint_cache[params.text_document.uri] = context.context.lint_models(
-                models,
+                models.names,
                 raise_on_error=False,
             )
             ls.publish_diagnostics(
@@ -125,8 +129,10 @@ def did_save(ls: LanguageServer, params: types.DidSaveTextDocumentParams) -> Non
             models = context.map[params.text_document.uri]
             if models is None:
                 return
+            if not isinstance(models, ModelTarget):
+                return
             self.lint_cache[params.text_document.uri] = context.context.lint_models(
-                models,
+                models.names,
                 raise_on_error=False,
             )
             ls.publish_diagnostics(

sqlmesh/lsp/reference.py

@@ -3,7 +3,7 @@
 
 from sqlmesh.core.dialect import normalize_model_name
 from sqlmesh.core.model.definition import SqlModel
-from sqlmesh.lsp.context import LSPContext
+from sqlmesh.lsp.context import LSPContext, ModelTarget, AuditTarget
 from sqlglot import exp
 
 from sqlmesh.utils.pydantic import PydanticModel
@@ -20,7 +20,7 @@ def get_model_definitions_for_a_path(
     """
     Get the model references for a given path.
 
-    Works for models and audits.
+    Works for models and standalone audits.
     Works for targeting sql and python models.
 
     Steps:
@@ -31,39 +31,63 @@ 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 model
+    # Ensure the path is a sql file
     if not document_uri.endswith(".sql"):
         return []
 
-    # Get the model
-    models = lint_context.map[document_uri]
-    if not models:
+    # Get the file info from the context map
+    if document_uri not in lint_context.map:
         return []
-    model = lint_context.context.get_model(model_or_snapshot=models[0], raise_if_missing=False)
-    if model is None or not isinstance(model, SqlModel):
+
+    file_info = lint_context.map[document_uri]
+
+    # Process based on whether it's a model or standalone audit
+    if isinstance(file_info, ModelTarget):
+        # It's a model
+        model = lint_context.context.get_model(
+            model_or_snapshot=file_info.names[0], raise_if_missing=False
+        )
+        if model is None or not isinstance(model, SqlModel):
+            return []
+
+        query = model.query
+        dialect = model.dialect
+        depends_on = model.depends_on
+        file_path = model._path
+    elif isinstance(file_info, AuditTarget):
+        # It's a standalone audit
+        audit = lint_context.context.standalone_audits.get(file_info.name)
+        if audit is None:
+            return []
+
+        query = audit.query
+        dialect = audit.dialect
+        depends_on = audit.depends_on
+        file_path = audit._path
+    else:
         return []
 
     # Find all possible references
     references = []
-    tables = list(model.query.find_all(exp.Table))
+
+    # Get SQL query and find all table references
+    tables = list(query.find_all(exp.Table))
     if len(tables) == 0:
         return []
 
-    read_file = open(model._path, "r").readlines()
+    read_file = open(file_path, "r").readlines()
 
     for table in tables:
-        depends_on = model.depends_on
-
         # Normalize the table reference
         unaliased = table.copy()
         if unaliased.args.get("alias") is not None:
             unaliased.set("alias", None)
-        reference_name = unaliased.sql(dialect=model.dialect)
+        reference_name = unaliased.sql(dialect=dialect)
         try:
             normalized_reference_name = normalize_model_name(
                 reference_name,
                 default_catalog=lint_context.context.default_catalog,
-                dialect=model.dialect,
+                dialect=dialect,
             )
             if normalized_reference_name not in depends_on:
                 continue

tests/lsp/test_context.py

@@ -1,6 +1,6 @@
 import pytest
 from sqlmesh.core.context import Context
-from sqlmesh.lsp.context import LSPContext
+from sqlmesh.lsp.context import LSPContext, ModelTarget
 
 
 @pytest.mark.fast
@@ -16,4 +16,7 @@ def test_lsp_context():
     active_customers_key = next(
         key for key in lsp_context.map.keys() if key.endswith("models/active_customers.sql")
     )
-    assert lsp_context.map[active_customers_key] == ["sushi.active_customers"]
+
+    # Check that the value is a ModelInfo with the expected model name
+    assert isinstance(lsp_context.map[active_customers_key], ModelTarget)
+    assert "sushi.active_customers" in lsp_context.map[active_customers_key].names

tests/lsp/test_reference.py

@@ -1,6 +1,6 @@
 import pytest
 from sqlmesh.core.context import Context
-from sqlmesh.lsp.context import LSPContext
+from sqlmesh.lsp.context import LSPContext, ModelTarget, AuditTarget
 from sqlmesh.lsp.reference import get_model_definitions_for_a_path
 
 
@@ -9,11 +9,16 @@ def test_reference() -> None:
     context = Context(paths=["examples/sushi"])
     lsp_context = LSPContext(context)
 
+    # Find model URIs
     active_customers_uri = next(
-        uri for uri, models in lsp_context.map.items() if "sushi.active_customers" in models
+        uri
+        for uri, info in lsp_context.map.items()
+        if isinstance(info, ModelTarget) and "sushi.active_customers" in info.names
     )
     sushi_customers_uri = next(
-        uri for uri, models in lsp_context.map.items() if "sushi.customers" in models
+        uri
+        for uri, info in lsp_context.map.items()
+        if isinstance(info, ModelTarget) and "sushi.customers" in info.names
     )
 
     references = get_model_definitions_for_a_path(lsp_context, active_customers_uri)
@@ -35,7 +40,9 @@ def test_reference_with_alias() -> None:
     lsp_context = LSPContext(context)
 
     waiter_revenue_by_day_uri = next(
-        uri for uri, models in lsp_context.map.items() if "sushi.waiter_revenue_by_day" in models
+        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)
@@ -52,6 +59,37 @@ def test_reference_with_alias() -> None:
     assert get_string_from_range(read_file, references[2].range) == "sushi.items"
 
 
+@pytest.mark.fast
+def test_standalone_audit_reference() -> None:
+    context = Context(paths=["examples/sushi"])
+    lsp_context = LSPContext(context)
+
+    # Find the standalone audit URI
+    audit_uri = 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(
+        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)
+
+    assert len(references) == 1
+    assert references[0].uri == items_uri
+
+    # Check that the reference in the correct range is sushi.items
+    path = audit_uri.removeprefix("file://")
+    read_file = open(path, "r").readlines()
+    referenced_text = get_string_from_range(read_file, references[0].range)
+    assert referenced_text == "sushi.items"
+
+
 def get_string_from_range(file_lines, range_obj) -> str:
     start_line = range_obj.start.line
     end_line = range_obj.end.line