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

Author: erindru

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:

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

sqlmesh_dbt/operations.py

@@ -185,7 +185,7 @@ def _plan_builder_options(
             options.update(
                 dict(
                     # Add every selected model as a restatement to force them to get repopulated from scratch
-                    restate_models=list(self.context.models)
+                    restate_models=[m.dbt_fqn for m in self.context.models.values() if m.dbt_fqn]
                     if not select_models
                     else select_models,
                     # by default in SQLMesh, restatements only operate on what has been committed to state.
@@ -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]

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

tests/dbt/cli/test_operations.py

@@ -138,7 +138,7 @@ def test_run_option_mapping(jaffle_shop_duckdb: Path):
     assert plan.selected_models_to_backfill is None
     assert {s.name for s in plan.snapshots} == {k for k in operations.context.snapshots}
 
-    plan = operations.run(select=["main.stg_orders+"])
+    plan = operations.run(select=["stg_orders+"])
     assert plan.environment.name == "prod"
     assert console.no_prompts is True
     assert console.no_diff is True
@@ -155,7 +155,7 @@ def test_run_option_mapping(jaffle_shop_duckdb: Path):
         plan.selected_models_to_backfill | {standalone_audit_name}
     )
 
-    plan = operations.run(select=["main.stg_orders+"], exclude=["main.customers"])
+    plan = operations.run(select=["stg_orders+"], exclude=["customers"])
     assert plan.environment.name == "prod"
     assert console.no_prompts is True
     assert console.no_diff is True
@@ -171,7 +171,7 @@ def test_run_option_mapping(jaffle_shop_duckdb: Path):
         plan.selected_models_to_backfill | {standalone_audit_name}
     )
 
-    plan = operations.run(exclude=["main.customers"])
+    plan = operations.run(exclude=["customers"])
     assert plan.environment.name == "prod"
     assert console.no_prompts is True
     assert console.no_diff is True
@@ -238,7 +238,7 @@ def test_run_option_mapping_dev(jaffle_shop_duckdb: Path):
     assert plan.skip_backfill is True
     assert plan.selected_models_to_backfill == {'"jaffle_shop"."main"."new_model"'}
 
-    plan = operations.run(environment="dev", select=["main.stg_orders+"])
+    plan = operations.run(environment="dev", select=["stg_orders+"])
     assert plan.environment.name == "dev"
     assert console.no_prompts is True
     assert console.no_diff is True
@@ -325,7 +325,7 @@ def test_run_option_full_refresh_with_selector(jaffle_shop_duckdb: Path):
     console = PlanCapturingConsole()
     operations.context.console = console
 
-    plan = operations.run(select=["main.stg_customers"], full_refresh=True)
+    plan = operations.run(select=["stg_customers"], full_refresh=True)
     assert len(plan.restatements) == 1
     assert list(plan.restatements)[0].name == '"jaffle_shop"."main"."stg_customers"'
 

tests/dbt/cli/test_run.py

@@ -27,7 +27,7 @@ def test_run_with_selectors(jaffle_shop_duckdb: Path, invoke_cli: t.Callable[...
         assert result.exit_code == 0
         assert "main.orders" in result.output
 
-    result = invoke_cli(["run", "--select", "main.raw_customers+", "--exclude", "main.orders"])
+    result = invoke_cli(["run", "--select", "raw_customers+", "--exclude", "orders"])
 
     assert result.exit_code == 0
     assert not result.exception