Feat(sqlmesh_dbt): Select based on dbt name, not sqlmesh name

erindru · erindru · commit 78f553062ea9 · 2025-09-21T23:02:58.000Z
diff --git a/sqlmesh/core/context.py b/sqlmesh/core/context.py
@@ -348,6 +348,8 @@ class GenericContext(BaseContext, t.Generic[C]):
         load: Whether or not to automatically load all models and macros (default True).
         console: The rich instance used for printing out CLI command results.
         users: A list of users to make known to SQLMesh.
+        dbt_mode: A flag to indicate we are running in 'dbt mode' which means that things like
+            model selections should use the dbt names and not the native SQLMesh names
     """
 
     CONFIG_TYPE: t.Type[C]
@@ -368,6 +370,7 @@ def __init__(
         load: bool = True,
         users: t.Optional[t.List[User]] = None,
         config_loader_kwargs: t.Optional[t.Dict[str, t.Any]] = None,
+        dbt_mode: bool = False,
     ):
         self.configs = (
             config
@@ -390,6 +393,7 @@ def __init__(
         self._engine_adapter: t.Optional[EngineAdapter] = None
         self._linters: t.Dict[str, Linter] = {}
         self._loaded: bool = False
+        self._dbt_mode = dbt_mode
 
         self.path, self.config = t.cast(t.Tuple[Path, C], next(iter(self.configs.items())))
 
@@ -2901,6 +2905,7 @@ def _new_selector(
             default_catalog=self.default_catalog,
             dialect=self.default_dialect,
             cache_dir=self.cache_dir,
+            dbt_mode=self._dbt_mode,
         )
 
     def _register_notification_targets(self) -> None:
diff --git a/sqlmesh/core/selector.py b/sqlmesh/core/selector.py
@@ -3,6 +3,7 @@
 import fnmatch
 import typing as t
 from pathlib import Path
+from itertools import zip_longest
 
 from sqlglot import exp
 from sqlglot.errors import ParseError
@@ -36,6 +37,7 @@ def __init__(
         default_catalog: t.Optional[str] = None,
         dialect: t.Optional[str] = None,
         cache_dir: t.Optional[Path] = None,
+        dbt_mode: bool = False,
     ):
         self._state_reader = state_reader
         self._models = models
@@ -44,6 +46,7 @@ def __init__(
         self._default_catalog = default_catalog
         self._dialect = dialect
         self._git_client = GitClient(context_path)
+        self._dbt_mode = dbt_mode
 
         if dag is None:
             self._dag: DAG[str] = DAG()
@@ -167,13 +170,13 @@ def get_model(fqn: str) -> t.Optional[Model]:
     def expand_model_selections(
         self, model_selections: t.Iterable[str], models: t.Optional[t.Dict[str, Model]] = None
     ) -> t.Set[str]:
-        """Expands a set of model selections into a set of model names.
+        """Expands a set of model selections into a set of model fqns that can be looked up in the Context.
 
         Args:
             model_selections: A set of model selections.
 
         Returns:
-            A set of model names.
+            A set of model fqns.
         """
 
         node = parse(" | ".join(f"({s})" for s in model_selections))
@@ -194,10 +197,9 @@ def evaluate(node: exp.Expression) -> t.Set[str]:
                     return {
                         fqn
                         for fqn, model in all_models.items()
-                        if fnmatch.fnmatchcase(model.name, node.this)
+                        if fnmatch.fnmatchcase(self._model_name(model), node.this)
                     }
-                fqn = normalize_model_name(pattern, self._default_catalog, self._dialect)
-                return {fqn} if fqn in all_models else set()
+                return self._pattern_to_model_fqns(pattern, all_models)
             if isinstance(node, exp.And):
                 return evaluate(node.left) & evaluate(node.right)
             if isinstance(node, exp.Or):
@@ -241,6 +243,59 @@ def evaluate(node: exp.Expression) -> t.Set[str]:
 
         return evaluate(node)
 
+    def _model_fqn(self, model: Model) -> str:
+        if self._dbt_mode:
+            dbt_fqn = model.dbt_fqn
+            if dbt_fqn is None:
+                raise SQLMeshError("Expecting dbt node information to be populated; it wasnt")
+            return dbt_fqn
+        return model.fqn
+
+    def _model_name(self, model: Model) -> str:
+        if self._dbt_mode:
+            # dbt always matches on the fqn, not the name
+            return self._model_fqn(model)
+        return model.name
+
+    def _pattern_to_model_fqns(self, pattern: str, all_models: t.Dict[str, Model]) -> t.Set[str]:
+        # note: all_models should be keyed by sqlmesh fqn, not dbt fqn
+        if not self._dbt_mode:
+            fqn = normalize_model_name(pattern, self._default_catalog, self._dialect)
+            return {fqn} if fqn in all_models else set()
+
+        # a pattern like "staging.customers" should match a model called "jaffle_shop.staging.customers"
+        # but not a model called "jaffle_shop.customers.staging"
+        # also a pattern like "aging" should not match "staging" so we need to consider components; not substrings
+        pattern_components = pattern.split(".")
+        first_pattern_component = pattern_components[0]
+        matches = set()
+        for fqn, model in all_models.items():
+            if not model.dbt_fqn:
+                continue
+
+            dbt_fqn_components = model.dbt_fqn.split(".")
+            try:
+                starting_idx = dbt_fqn_components.index(first_pattern_component)
+            except ValueError:
+                continue
+            for pattern_component, fqn_component in zip_longest(
+                pattern_components, dbt_fqn_components[starting_idx:]
+            ):
+                if pattern_component and not fqn_component:
+                    # the pattern still goes but we have run out of fqn components to match; no match
+                    break
+                if fqn_component and not pattern_component:
+                    # all elements of the pattern have matched elements of the fqn; match
+                    matches.add(fqn)
+                    break
+                if pattern_component != fqn_component:
+                    # the pattern explicitly doesnt match a component; no match
+                    break
+            else:
+                # called if no explicit break, indicating all components of the pattern matched all components of the fqn
+                matches.add(fqn)
+        return matches
+
 
 class SelectorDialect(Dialect):
     IDENTIFIERS_CAN_START_WITH_DIGIT = True
diff --git a/sqlmesh_dbt/operations.py b/sqlmesh_dbt/operations.py
@@ -250,6 +250,8 @@ def create(
             paths=[project_dir],
             config_loader_kwargs=dict(profile=profile, target=target, variables=vars),
             load=True,
+            # dbt mode enables selectors to use dbt model fqn's rather than SQLMesh model names
+            dbt_mode=True,
         )
 
         dbt_loader = sqlmesh_context._loaders[0]
diff --git a/tests/dbt/cli/test_list.py b/tests/dbt/cli/test_list.py
@@ -19,7 +19,7 @@ def test_list(jaffle_shop_duckdb: Path, invoke_cli: t.Callable[..., Result]):
 
 
 def test_list_select(jaffle_shop_duckdb: Path, invoke_cli: t.Callable[..., Result]):
-    result = invoke_cli(["list", "--select", "main.raw_customers+"])
+    result = invoke_cli(["list", "--select", "raw_customers+"])
 
     assert result.exit_code == 0
     assert not result.exception
@@ -34,7 +34,7 @@ def test_list_select(jaffle_shop_duckdb: Path, invoke_cli: t.Callable[..., Resul
 
 def test_list_select_exclude(jaffle_shop_duckdb: Path, invoke_cli: t.Callable[..., Result]):
     # single exclude
-    result = invoke_cli(["list", "--select", "main.raw_customers+", "--exclude", "main.orders"])
+    result = invoke_cli(["list", "--select", "raw_customers+", "--exclude", "orders"])
 
     assert result.exit_code == 0
     assert not result.exception
@@ -49,8 +49,8 @@ def test_list_select_exclude(jaffle_shop_duckdb: Path, invoke_cli: t.Callable[..
 
     # multiple exclude
     for args in (
-        ["--select", "main.stg_orders+", "--exclude", "main.customers", "--exclude", "main.orders"],
-        ["--select", "main.stg_orders+", "--exclude", "main.customers main.orders"],
+        ["--select", "stg_orders+", "--exclude", "customers", "--exclude", "orders"],
+        ["--select", "stg_orders+", "--exclude", "customers orders"],
     ):
         result = invoke_cli(["list", *args])
         assert result.exit_code == 0
diff --git a/tests/dbt/cli/test_selectors.py b/tests/dbt/cli/test_selectors.py
@@ -1,6 +1,8 @@
 import typing as t
 import pytest
 from sqlmesh_dbt import selectors
+from sqlmesh.core.context import Context
+from pathlib import Path
 
 
 @pytest.mark.parametrize(
@@ -77,3 +79,181 @@ def test_split_unions_and_intersections(
     expression: str, expected: t.Tuple[t.List[str], t.List[str]]
 ):
     assert selectors._split_unions_and_intersections(expression) == expected
+
+
+@pytest.mark.parametrize(
+    "dbt_select,expected",
+    [
+        (["aging"], set()),
+        (
+            ["staging"],
+            {
+                '"jaffle_shop"."main"."stg_customers"',
+                '"jaffle_shop"."main"."stg_orders"',
+                '"jaffle_shop"."main"."stg_payments"',
+            },
+        ),
+        (["staging.stg_customers"], {'"jaffle_shop"."main"."stg_customers"'}),
+        (["stg_customers.staging"], set()),
+        (
+            ["+customers"],
+            {
+                '"jaffle_shop"."main"."customers"',
+                '"jaffle_shop"."main"."stg_customers"',
+                '"jaffle_shop"."main"."stg_orders"',
+                '"jaffle_shop"."main"."stg_payments"',
+                '"jaffle_shop"."main"."raw_customers"',
+                '"jaffle_shop"."main"."raw_orders"',
+                '"jaffle_shop"."main"."raw_payments"',
+            },
+        ),
+        (["customers+"], {'"jaffle_shop"."main"."customers"'}),
+        (
+            ["customers+", "stg_orders"],
+            {'"jaffle_shop"."main"."customers"', '"jaffle_shop"."main"."stg_orders"'},
+        ),
+        (["tag:agg"], {'"jaffle_shop"."main"."agg_orders"'}),
+        (
+            ["staging.stg_customers", "tag:agg"],
+            {
+                '"jaffle_shop"."main"."stg_customers"',
+                '"jaffle_shop"."main"."agg_orders"',
+            },
+        ),
+        (
+            ["+tag:agg"],
+            {
+                '"jaffle_shop"."main"."agg_orders"',
+                '"jaffle_shop"."main"."orders"',
+                '"jaffle_shop"."main"."stg_orders"',
+                '"jaffle_shop"."main"."stg_payments"',
+                '"jaffle_shop"."main"."raw_orders"',
+                '"jaffle_shop"."main"."raw_payments"',
+            },
+        ),
+        (
+            ["tag:agg+"],
+            {
+                '"jaffle_shop"."main"."agg_orders"',
+            },
+        ),
+    ],
+)
+def test_select_by_dbt_names(
+    jaffle_shop_duckdb: Path,
+    jaffle_shop_duckdb_context: Context,
+    dbt_select: t.List[str],
+    expected: t.Set[str],
+):
+    (jaffle_shop_duckdb / "models" / "agg_orders.sql").write_text("""
+     {{ config(tags=["agg"]) }}
+      select order_date, count(*) as num_orders from {{ ref('orders') }}                                                             
+    """)
+
+    ctx = jaffle_shop_duckdb_context
+    ctx.load()
+    assert '"jaffle_shop"."main"."agg_orders"' in ctx.models
+
+    selector = ctx._new_selector()
+    assert selector._dbt_mode
+
+    sqlmesh_selector = selectors.to_sqlmesh(dbt_select=dbt_select, dbt_exclude=[])
+    assert sqlmesh_selector
+
+    assert selector.expand_model_selections([sqlmesh_selector]) == expected
+
+
+@pytest.mark.parametrize(
+    "dbt_exclude,expected",
+    [
+        (["jaffle_shop"], set()),
+        (
+            ["staging"],
+            {
+                '"jaffle_shop"."main"."agg_orders"',
+                '"jaffle_shop"."main"."customers"',
+                '"jaffle_shop"."main"."orders"',
+                '"jaffle_shop"."main"."raw_customers"',
+                '"jaffle_shop"."main"."raw_orders"',
+                '"jaffle_shop"."main"."raw_payments"',
+            },
+        ),
+        (["+customers"], {'"jaffle_shop"."main"."orders"', '"jaffle_shop"."main"."agg_orders"'}),
+        (
+            ["+tag:agg"],
+            {
+                '"jaffle_shop"."main"."customers"',
+                '"jaffle_shop"."main"."stg_customers"',
+                '"jaffle_shop"."main"."raw_customers"',
+            },
+        ),
+    ],
+)
+def test_exclude_by_dbt_names(
+    jaffle_shop_duckdb: Path,
+    jaffle_shop_duckdb_context: Context,
+    dbt_exclude: t.List[str],
+    expected: t.Set[str],
+):
+    (jaffle_shop_duckdb / "models" / "agg_orders.sql").write_text("""
+     {{ config(tags=["agg"]) }}
+      select order_date, count(*) as num_orders from {{ ref('orders') }}                                                             
+    """)
+
+    ctx = jaffle_shop_duckdb_context
+    ctx.load()
+    assert '"jaffle_shop"."main"."agg_orders"' in ctx.models
+
+    selector = ctx._new_selector()
+    assert selector._dbt_mode
+
+    sqlmesh_selector = selectors.to_sqlmesh(dbt_select=[], dbt_exclude=dbt_exclude)
+    assert sqlmesh_selector
+
+    assert selector.expand_model_selections([sqlmesh_selector]) == expected
+
+
+@pytest.mark.parametrize(
+    "dbt_select,dbt_exclude,expected",
+    [
+        (["jaffle_shop"], ["jaffle_shop"], set()),
+        (
+            ["staging"],
+            ["stg_customers"],
+            {
+                '"jaffle_shop"."main"."stg_orders"',
+                '"jaffle_shop"."main"."stg_payments"',
+            },
+        ),
+        (
+            ["staging.stg_customers", "tag:agg"],
+            ["tag:agg"],
+            {
+                '"jaffle_shop"."main"."stg_customers"',
+            },
+        ),
+    ],
+)
+def test_selection_and_exclusion_by_dbt_names(
+    jaffle_shop_duckdb: Path,
+    jaffle_shop_duckdb_context: Context,
+    dbt_select: t.List[str],
+    dbt_exclude: t.List[str],
+    expected: t.Set[str],
+):
+    (jaffle_shop_duckdb / "models" / "agg_orders.sql").write_text("""
+     {{ config(tags=["agg"]) }}
+      select order_date, count(*) as num_orders from {{ ref('orders') }}                                                             
+    """)
+
+    ctx = jaffle_shop_duckdb_context
+    ctx.load()
+    assert '"jaffle_shop"."main"."agg_orders"' in ctx.models
+
+    selector = ctx._new_selector()
+    assert selector._dbt_mode
+
+    sqlmesh_selector = selectors.to_sqlmesh(dbt_select=dbt_select, dbt_exclude=dbt_exclude)
+    assert sqlmesh_selector
+
+    assert selector.expand_model_selections([sqlmesh_selector]) == expected
diff --git a/tests/dbt/conftest.py b/tests/dbt/conftest.py
@@ -99,7 +99,7 @@ def jaffle_shop_duckdb(copy_to_temp_path: t.Callable[..., t.List[Path]]) -> t.It
 @pytest.fixture
 def jaffle_shop_duckdb_context(jaffle_shop_duckdb: Path) -> Context:
     init_project_if_required(jaffle_shop_duckdb)
-    return Context(paths=[jaffle_shop_duckdb])
+    return Context(paths=[jaffle_shop_duckdb], dbt_mode=True)
 
 
 @pytest.fixture()

Original file line number	Diff line number	Diff line change
`@@ -250,6 +250,8 @@ def create(`
`250`	`250`	`paths=[project_dir],`
`251`	`251`	`config_loader_kwargs=dict(profile=profile, target=target, variables=vars),`
`252`	`252`	`load=True,`
	`253`	`+ # dbt mode enables selectors to use dbt model fqn's rather than SQLMesh model names`
	`254`	`+ dbt_mode=True,`
`253`	`255`	`)`
`254`	`256`
`255`	`257`	`dbt_loader = sqlmesh_context._loaders[0]`