Repository: TobikoData/sqlmesh
Pull request #4573 (2025-05-29 06:31:24)

Fix: expand the 'tag:' selector based on local models and not remote ones
This addresses an issue where if you add a tag to a local model that has previously been run / committed to state and run `sqlmesh plan --select-model 'tag:<new tag>'`, the model doesnt get included.


Patch:
diff --git a/examples/sushi/macros/utils.py b/examples/sushi/macros/utils.py
index fbc4c04e..a76bc3bf 100644
--- a/examples/sushi/macros/utils.py
+++ b/examples/sushi/macros/utils.py
@@ -26,6 +26,7 @@ def sql_literal(
     column_str: str,
     column_quoted: str,
 ):
+    """A macro that accepts various types of SQL literals and returns the column."""
     assert isinstance(column, str)
     assert isinstance(str_lit, str)
     assert str_lit == "'x'"
diff --git a/sqlmesh/core/context.py b/sqlmesh/core/context.py
index 1f472226..238e1995 100644
--- a/sqlmesh/core/context.py
+++ b/sqlmesh/core/context.py
@@ -932,19 +932,22 @@ class GenericContext(BaseContext, t.Generic[C]):
 
         return snapshot
 
-    def config_for_path(self, path: Path) -> Config:
+    def config_for_path(self, path: Path) -> t.Tuple[Config, Path]:
+        """Returns the config and path of the said project for a given file path."""
         for config_path, config in self.configs.items():
             try:
                 path.relative_to(config_path)
-                return config
+                return config, config_path
             except ValueError:
                 pass
-        return self.config
+        return self.config, self.path
 
     def config_for_node(self, node: str | Model | Audit) -> Config:
         if isinstance(node, str):
-            return self.config_for_path(self.get_snapshot(node, raise_if_missing=True).node._path)  # type: ignore
-        return self.config_for_path(node._path)  # type: ignore
+            return self.config_for_path(self.get_snapshot(node, raise_if_missing=True).node._path)[
+                0
+            ]  # type: ignore
+        return self.config_for_path(node._path)[0]  # type: ignore
 
     @property
     def models(self) -> MappingProxyType[str, Model]:
diff --git a/sqlmesh/core/selector.py b/sqlmesh/core/selector.py
index 4519afe4..ffc5af3e 100644
--- a/sqlmesh/core/selector.py
+++ b/sqlmesh/core/selector.py
@@ -93,7 +93,7 @@ class Selector:
             }
 
         all_selected_models = self.expand_model_selections(
-            model_selections, models={**self._models, **env_models}
+            model_selections, models={**env_models, **self._models}
         )
 
         dag: DAG[str] = DAG()
diff --git a/sqlmesh/lsp/reference.py b/sqlmesh/lsp/reference.py
index 58409cca..b8e6e9c4 100644
--- a/sqlmesh/lsp/reference.py
+++ b/sqlmesh/lsp/reference.py
@@ -1,6 +1,8 @@
 from lsprotocol.types import Range, Position
 import typing as t
+from pathlib import Path
 
+from sqlmesh.core.audit import StandaloneAudit
 from sqlmesh.core.dialect import normalize_model_name
 from sqlmesh.core.model.definition import SqlModel
 from sqlmesh.lsp.context import LSPContext, ModelTarget, AuditTarget
@@ -10,6 +12,10 @@ from sqlglot.optimizer.scope import build_scope
 from sqlmesh.lsp.uri import URI
 from sqlmesh.utils.pydantic import PydanticModel
 from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
+import ast
+from sqlmesh.core.model import Model
+from sqlmesh import macro
+import inspect
 
 
 class Reference(PydanticModel):
@@ -72,6 +78,11 @@ def get_references(
         A list of references at the given position
     """
     references = get_model_definitions_for_a_path(lint_context, document_uri)
+
+    # Get macro references before filtering by position
+    macro_references = get_macro_definitions_for_a_path(lint_context, document_uri)
+    references.extend(macro_references)
+
     filtered_references = list(filter(by_position(position), references))
     return filtered_references
 
@@ -134,9 +145,12 @@ def get_model_definitions_for_a_path(
     with open(file_path, "r", encoding="utf-8") as file:
         read_file = file.readlines()
 
-    # Build scope tree to properly handle nested CTEs
-    query = normalize_identifiers(query.copy(), dialect=dialect)
-    root_scope = build_scope(query)
+    # Build a scope tree to properly handle nested CTEs
+    try:
+        query = normalize_identifiers(query.copy(), dialect=dialect)
+        root_scope = build_scope(query)
+    except Exception:
+        root_scope = None
 
     if root_scope:
         # Traverse all scopes to find CTE definitions and table references
@@ -287,3 +301,180 @@ def _range_from_token_position_details(
         start=Position(line=start_line_0, character=start_col_0),
         end=Position(line=end_line_0, character=end_col_0),
     )
+
+
+def get_macro_definitions_for_a_path(
+    lsp_context: LSPContext, document_uri: URI
+) -> t.List[Reference]:
+    """
+    Get macro references for a given path.
+
+    This function finds all macro invocations (e.g., @ADD_ONE, @MULTIPLY) in a SQL file
+    and creates references to their definitions in the Python macro files.
+
+    Args:
+        lsp_context: The LSP context containing macro definitions
+        document_uri: The URI of the document to search for macro invocations
+
+    Returns:
+        A list of Reference objects for each macro invocation found
+    """
+    path = document_uri.to_path()
+    if path.suffix != ".sql":
+        return []
+
+    # Get the file info from the context map
+    if path not in lsp_context.map:
+        return []
+
+    file_info = lsp_context.map[path]
+    # Process based on whether it's a model or standalone audit
+    if isinstance(file_info, ModelTarget):
+        # It's a model
+        target: t.Optional[t.Union[Model, StandaloneAudit]] = lsp_context.context.get_model(
+            model_or_snapshot=file_info.names[0], raise_if_missing=False
+        )
+        if target is None or not isinstance(target, SqlModel):
+            return []
+        query = target.query
+        file_path = target._path
+    elif isinstance(file_info, AuditTarget):
+        # It's a standalone audit
+        target = lsp_context.context.standalone_audits.get(file_info.name)
+        if target is None:
+            return []
+        query = target.query
+        file_path = target._path
+    else:
+        return []
+
+    references = []
+    config_for_model, config_path = lsp_context.context.config_for_path(
+        file_path,
+    )
+
+    with open(file_path, "r", encoding="utf-8") as file:
+        read_file = file.readlines()
+
+    for node in query.find_all(exp.Anonymous):
+        macro_name = node.name.lower()
+        reference = get_macro_reference(
+            node=node,
+            target=target,
+            read_file=read_file,
+            config_path=config_path,
+            macro_name=macro_name,
+        )
+        if reference is not None:
+            references.append(reference)
+
+    return references
+
+
+def get_macro_reference(
+    target: t.Union[Model, StandaloneAudit],
+    read_file: t.List[str],
+    config_path: t.Optional[Path],
+    node: exp.Expression,
+    macro_name: str,
+) -> t.Optional[Reference]:
+    # Get the file path where the macro is defined
+    try:
+        # Get the position of the macro invocation in the source file first
+        if hasattr(node, "meta") and node.meta:
+            token_details = TokenPositionDetails.from_meta(node.meta)
+            macro_range = _range_from_token_position_details(token_details, read_file)
+
+            # Check if it's a built-in method
+            if builtin := get_built_in_macro_reference(macro_name, macro_range):
+                return builtin
+        else:
+            # Skip if we can't get the position
+            return None
+
+        # Find the macro definition information
+        macro_def = target.python_env.get(macro_name)
+        if macro_def is None:
+            return None
+
+        function_name = macro_def.name
+        if not function_name:
+            return None
+        if not macro_def.path:
+            return None
+        if not config_path:
+            return None
+        path = Path(config_path).joinpath(macro_def.path)
+
+        # Parse the Python file to find the function definition
+        with open(path, "r") as f:
+            tree = ast.parse(f.read())
+        with open(path, "r") as f:
+            output_read_line = f.readlines()
+
+        # Find the function definition by name
+        start_line = None
+        end_line = None
+        get_length_of_end_line = None
+        docstring = None
+        for ast_node in ast.walk(tree):
+            if isinstance(ast_node, ast.FunctionDef) and ast_node.name == function_name:
+                start_line = ast_node.lineno
+                end_line = ast_node.end_lineno
+                get_length_of_end_line = (
+                    len(output_read_line[end_line - 1])
+                    if end_line is not None and end_line - 1 < len(read_file)
+                    else 0
+                )
+                # Extract docstring if present
+                docstring = ast.get_docstring(ast_node)
+                break
+
+        if start_line is None or end_line is None or get_length_of_end_line is None:
+            return None
+
+        # Create a reference to the macro definition
+        macro_uri = URI.from_path(path)
+
+        return Reference(
+            uri=macro_uri.value,
+            range=macro_range,
+            target_range=Range(
+                start=Position(line=start_line - 1, character=0),
+                end=Position(line=end_line - 1, character=get_length_of_end_line),
+            ),
+            markdown_description=docstring,
+        )
+    except Exception:
+        return None
+
+
+def get_built_in_macro_reference(macro_name: str, macro_range: Range) -> t.Optional[Reference]:
+    """
+    Get a reference to a built-in macro by its name.
+
+    Args:
+        macro_name: The name of the built-in macro (e.g., 'each', 'sql_literal')
+        macro_range: The range of the macro invocation in the source file
+    """
+    built_in_macros = macro.get_registry()
+    built_in_macro = built_in_macros.get(macro_name)
+    if built_in_macro is None:
+        return None
+
+    func = built_in_macro.func
+    filename = inspect.getfile(func)
+    source_lines, line_number = inspect.getsourcelines(func)
+
+    # Calculate the end line number by counting the number of source lines
+    end_line_number = line_number + len(source_lines) - 1
+
+    return Reference(
+        uri=URI.from_path(Path(filename)).value,
+        range=macro_range,
+        target_range=Range(
+            start=Position(line=line_number - 1, character=0),
+            end=Position(line=end_line_number - 1, character=0),
+        ),
+        markdown_description=func.__doc__ if func.__doc__ else None,
+    )
diff --git a/web/server/api/endpoints/files.py b/web/server/api/endpoints/files.py
index 7e4d3652..d843e65c 100644
--- a/web/server/api/endpoints/files.py
+++ b/web/server/api/endpoints/files.py
@@ -60,7 +60,7 @@ async def write_file(
         replace_file(settings.project_path / path, settings.project_path / path_or_new_path)
     else:
         full_path = settings.project_path / path
-        config = context.config_for_path(Path(path_or_new_path)) if context else None
+        config, _ = context.config_for_path(Path(path_or_new_path)) if context else (None, None)
         if (
             config
             and config.ui.format_on_save
diff --git a/web/server/utils.py b/web/server/utils.py
index cda6a7d9..e2238d1b 100644
--- a/web/server/utils.py
+++ b/web/server/utils.py
@@ -64,11 +64,10 @@ def validate_path(path: str, settings: Settings = Depends(get_settings)) -> str:
     if any(
         full_path.match(pattern)
         for pattern in (
-            context.config_for_path(Path(path)).ignore_patterns if context else c.IGNORE_PATTERNS
+            context.config_for_path(Path(path))[0].ignore_patterns if context else c.IGNORE_PATTERNS
         )
     ):
         raise HTTPException(status_code=HTTP_404_NOT_FOUND)
-
     return path
 
 


Test patch:
diff --git a/tests/core/test_selector.py b/tests/core/test_selector.py
index 7c4e7713..9f3bc9f6 100644
--- a/tests/core/test_selector.py
+++ b/tests/core/test_selector.py
@@ -657,6 +657,55 @@ def test_select_models_with_external_parent(mocker: MockerFixture):
     assert expanded_selections == {added_model.fqn}
 
 
+def test_select_models_local_tags_take_precedence_over_remote(
+    mocker: MockerFixture, make_snapshot: t.Callable
+) -> None:
+    existing_model = SqlModel(
+        name="db.existing",
+        query=d.parse_one("SELECT 1 AS a"),
+    )
+
+    existing_snapshot = make_snapshot(existing_model)
+    existing_snapshot.categorize_as(SnapshotChangeCategory.BREAKING)
+
+    env_name = "test_env"
+
+    state_reader_mock = mocker.Mock()
+    state_reader_mock.get_environment.return_value = Environment(
+        name=env_name,
+        snapshots=[existing_snapshot.table_info],
+        start_at="2023-01-01",
+        end_at="2023-02-01",
+        plan_id="test_plan_id",
+    )
+    state_reader_mock.get_snapshots.return_value = {
+        existing_snapshot.snapshot_id: existing_snapshot
+    }
+
+    local_models: UniqueKeyDict[str, Model] = UniqueKeyDict("models")
+    local_new = SqlModel(
+        name="db.new",
+        tags=["a"],
+        query=d.parse_one("SELECT 1 as a"),
+    )
+    local_existing = existing_model.copy(update={"tags": ["a"]})  # type: ignore
+    local_models[local_existing.fqn] = local_existing
+    local_models[local_new.fqn] = local_new
+
+    selector = Selector(state_reader_mock, local_models)
+
+    selected = selector.select_models(["tag:a"], env_name)
+
+    # both should get selected because they both now have the 'a' tag locally, even though one exists in remote state without the 'a' tag
+    _assert_models_equal(
+        selected,
+        {
+            local_existing.fqn: local_existing,
+            local_new.fqn: local_new,
+        },
+    )
+
+
 def _assert_models_equal(actual: t.Dict[str, Model], expected: t.Dict[str, Model]) -> None:
     assert set(actual) == set(expected)
     for name, model in actual.items():
diff --git a/tests/lsp/test_reference_macro.py b/tests/lsp/test_reference_macro.py
new file mode 100644
index 00000000..35cf317b
--- /dev/null
+++ b/tests/lsp/test_reference_macro.py
@@ -0,0 +1,30 @@
+import pytest
+from sqlmesh.core.context import Context
+from sqlmesh.lsp.context import LSPContext, ModelTarget
+from sqlmesh.lsp.reference import get_macro_definitions_for_a_path
+from sqlmesh.lsp.uri import URI
+
+
+@pytest.mark.fast
+def test_macro_references() -> None:
+    """Test that macro references (e.g., @ADD_ONE, @MULTIPLY) have proper go-to-definition support."""
+    context = Context(paths=["examples/sushi"])
+    lsp_context = LSPContext(context)
+
+    # Find the top_waiters model that uses macros
+    top_waiters_path = next(
+        path
+        for path, info in lsp_context.map.items()
+        if isinstance(info, ModelTarget) and "sushi.top_waiters" in info.names
+    )
+
+    top_waiters_uri = URI.from_path(top_waiters_path)
+    macro_references = get_macro_definitions_for_a_path(lsp_context, top_waiters_uri)
+
+    # We expect 3 macro references: @ADD_ONE, @MULTIPLY, @SQL_LITERAL
+    assert len(macro_references) == 3
+
+    # Check that all references point to the utils.py file
+    for ref in macro_references:
+        assert ref.uri.endswith("sushi/macros/utils.py")
+        assert ref.target_range is not None
diff --git a/tests/lsp/test_reference_macro_multi.py b/tests/lsp/test_reference_macro_multi.py
new file mode 100644
index 00000000..3dfa588e
--- /dev/null
+++ b/tests/lsp/test_reference_macro_multi.py
@@ -0,0 +1,25 @@
+import pytest
+from sqlmesh.core.context import Context
+from sqlmesh.lsp.context import LSPContext, ModelTarget
+from sqlmesh.lsp.reference import get_macro_definitions_for_a_path
+from sqlmesh.lsp.uri import URI
+
+
+@pytest.mark.fast
+def test_macro_references_multirepo() -> None:
+    context = Context(paths=["examples/multi/repo_1", "examples/multi/repo_2"])
+    lsp_context = LSPContext(context)
+
+    d_path = next(
+        path
+        for path, info in lsp_context.map.items()
+        if isinstance(info, ModelTarget) and "silver.d" in info.names
+    )
+
+    d = URI.from_path(d_path)
+    macro_references = get_macro_definitions_for_a_path(lsp_context, d)
+
+    assert len(macro_references) == 2
+    for ref in macro_references:
+        assert ref.uri.endswith("multi/repo_2/macros/__init__.py")
+        assert ref.target_range is not None
diff --git a/vscode/extension/tests/go_to_definition.spec.ts b/vscode/extension/tests/go_to_definition.spec.ts
new file mode 100644
index 00000000..7047f06f
--- /dev/null
+++ b/vscode/extension/tests/go_to_definition.spec.ts
@@ -0,0 +1,38 @@
+import { test, expect } from '@playwright/test';
+import path from 'path';
+import fs from 'fs-extra';
+import os from 'os';
+import { startVSCode, SUSHI_SOURCE_PATH } from './utils';
+
+test('Go to definition for macro', async () => {
+    const tempDir = await fs.mkdtemp(path.join(os.tmpdir(), 'vscode-test-sushi-'));
+    await fs.copy(SUSHI_SOURCE_PATH, tempDir);
+
+    try {
+      const { window, close } = await startVSCode(tempDir);
+
+      //   Wait for the models folder to be visible
+      await window.waitForSelector('text=models');
+
+      // Click on the models folder, excluding external_models
+      await window.getByRole('treeitem', { name: 'models', exact: true }).locator('a').click();
+
+      // Open the customer_revenue_lifetime model
+      await window.getByRole('treeitem', { name: 'top_waiters.sql', exact: true }).locator('a').click();
+
+      await window.waitForSelector('text=grain');
+      await window.waitForSelector('text=Loaded SQLMesh Context')
+
+      // Render the model
+      window.locator("text=@MULTIPLY").click({
+        modifiers: ["Meta"]
+      })
+
+      // Check if the model is rendered by check if "`oi`.`order_id` AS `order_id`," is in the window
+      await expect(window.locator('text=def multiply(')).toBeVisible();
+
+      await close();
+    } finally {
+      await fs.removeSync(tempDir);
+    }
+  });
\ No newline at end of file
diff --git a/vscode/extension/tests/utils.ts b/vscode/extension/tests/utils.ts
index bceceb0f..e836d61e 100644
--- a/vscode/extension/tests/utils.ts
+++ b/vscode/extension/tests/utils.ts
@@ -40,11 +40,8 @@ export const startVSCode = async (workspaceDir: string): Promise<{
       args,
     });
     const window = await electronApp.firstWindow();
-    await window.waitForLoadState('domcontentloaded');
-    await window.waitForLoadState('networkidle');
-    await window.waitForTimeout(2_000);
     return { window, close: async () => {
       await electronApp.close();
-      await fs.remove(userDataDir);
+      await fs.removeSync(userDataDir);
     } };
   }