Feat(dbt): Add support for selected resources context variable

Author: themisvaltinos

sqlmesh/core/context.py

@@ -117,6 +117,7 @@
     run_tests,
 )
 from sqlmesh.core.user import User
+from sqlmesh.dbt.builtin import set_selected_resources
 from sqlmesh.utils import UniqueKeyDict, Verbosity
 from sqlmesh.utils.concurrency import concurrent_apply_to_values
 from sqlmesh.utils.dag import DAG
@@ -1583,6 +1584,11 @@ def plan_builder(
                 "Selector did not return any models. Please check your model selection and try again."
             )
 
+        if self._project_type != c.NATIVE:
+            set_selected_resources(
+                models=model_selector.expand_model_selections(select_models or "*")
+            )
+
         snapshots = self._snapshots(models_override)
         context_diff = self._context_diff(
             environment or c.PROD,
@@ -2482,6 +2488,9 @@ def _run(
                 select_models, no_auto_upstream, snapshots.values()
             )
 
+        if self._project_type != c.NATIVE:
+            set_selected_resources(models=select_models or set([s.name for s in snapshots.keys()]))
+
         completion_status = scheduler.run(
             environment,
             start=start,

sqlmesh/dbt/builtin.py

@@ -540,6 +540,7 @@ def create_builtin_globals(
             "run_query": sql_execution.run_query,
             "statement": sql_execution.statement,
             "graph": adapter.graph,
+            "selected_resources": get_selected_resources(),
         }
     )
 
@@ -567,3 +568,33 @@ def _relation_info_to_relation(
         }
     )
     return relation_type.create(**relation_info, quote_policy=quote_policy)
+
+
+_selected_resources: t.List[str] = []
+
+
+def set_selected_resources(
+    models: t.Optional[t.Set[str]] = None,
+) -> None:
+    global _selected_resources
+    resources = []
+
+    if models:
+        for model in models:
+            resources.append(dbt_model_id(model))
+
+    _selected_resources = sorted(resources)
+
+
+def dbt_model_id(sqlmesh_model_name: str) -> str:
+    parts = [part.strip('"') for part in sqlmesh_model_name.split(".")]
+    return f"model.{parts[0]}.{parts[-1]}"
+
+
+def get_selected_resources() -> t.List[str]:
+    return _selected_resources
+
+
+def clear_selected_resources() -> None:
+    global _selected_resources
+    _selected_resources = []

tests/dbt/test_transformation.py

@@ -45,6 +45,12 @@
 from sqlmesh.core.state_sync.db.snapshot import _snapshot_to_json
 from sqlmesh.dbt.builtin import _relation_info_to_relation, Config
 from sqlmesh.dbt.common import Dependencies
+from sqlmesh.dbt.builtin import (
+    _relation_info_to_relation,
+    dbt_model_id,
+    clear_selected_resources,
+    get_selected_resources,
+)
 from sqlmesh.dbt.column import (
     ColumnConfig,
     column_descriptions_to_sqlmesh,
@@ -2352,3 +2358,68 @@ def test_dynamic_var_names_in_macro(sushi_test_project: Project):
     )
     converted_model = model_config.to_sqlmesh(context)
     assert "dynamic_test_var" in converted_model.jinja_macros.global_objs["vars"]  # type: ignore
+
+
+def test_selected_resources_with_selectors():
+    sushi_context = Context(paths=["tests/fixtures/dbt/sushi_test"])
+
+    # A plan with a specific model selection
+    clear_selected_resources()
+    sushi_context.plan_builder(select_models=["sushi.customers"])
+
+    selected = get_selected_resources()
+    assert "model.memory.customers" in selected
+    assert len(selected) == 1
+
+    # Plan without model selections
+    clear_selected_resources()
+    sushi_context.plan_builder()
+    selected = get_selected_resources()
+    assert sorted(
+        [
+            "model.memory.customer_revenue_by_day",
+            "model.memory.customers",
+            "model.memory.items",
+            "model.memory.items_check_snapshot",
+            "model.memory.items_no_hard_delete_snapshot",
+            "model.memory.items_snapshot",
+            "model.memory.order_items",
+            "model.memory.orders",
+            "model.memory.simple_model_a",
+            "model.memory.simple_model_b",
+            "model.memory.top_waiters",
+            "model.memory.waiter_as_customer_by_day",
+            "model.memory.waiter_names",
+            "model.memory.waiter_revenue_by_day_v1",
+            "model.memory.waiter_revenue_by_day_v2",
+            "model.memory.waiters",
+        ]
+    ) == sorted(selected)
+
+    # Test with downstream models as well
+    clear_selected_resources()
+    sushi_context.plan_builder(select_models=["sushi.customers+"])
+    selected = get_selected_resources()
+    assert sorted(["model.memory.customers", "model.memory.waiter_as_customer_by_day"]) == sorted(
+        selected
+    )
+
+    # Test wildcard selection
+    clear_selected_resources()
+    sushi_context.plan_builder(select_models=["sushi.waiter_*"])
+    selected = get_selected_resources()
+    assert sorted(
+        [
+            "model.memory.waiter_as_customer_by_day",
+            "model.memory.waiter_names",
+            "model.memory.waiter_revenue_by_day_v1",
+            "model.memory.waiter_revenue_by_day_v2",
+        ]
+    ) == sorted(selected)
+    clear_selected_resources()
+
+
+def test_dbt_model_id_conversion():
+    assert dbt_model_id("jaffle_shop.main.customers") == "model.jaffle_shop.customers"
+    assert dbt_model_id("jaffle_shop.main.orders") == "model.jaffle_shop.orders"
+    assert dbt_model_id('"jaffle_shop"."customers"') == "model.jaffle_shop.customers"