pr feedback

Author: themisvaltinos

sqlmesh/core/context.py

@@ -117,7 +117,6 @@
     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,11 +1582,6 @@ 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,
@@ -1682,6 +1676,7 @@ def plan_builder(
             end_override_per_model=max_interval_end_per_model,
             console=self.console,
             user_provided_flags=user_provided_flags,
+            selected_models=model_selector.expand_model_selections(select_models or "*"),
             explain=explain or False,
             ignore_cron=ignore_cron or False,
         )
@@ -2488,9 +2483,6 @@ 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/core/environment.py

@@ -312,6 +312,7 @@ def execute_environment_statements(
     start: t.Optional[TimeLike] = None,
     end: t.Optional[TimeLike] = None,
     execution_time: t.Optional[TimeLike] = None,
+    selected_models: t.Optional[t.Set[str]] = None,
 ) -> None:
     try:
         rendered_expressions = [
@@ -327,6 +328,7 @@ def execute_environment_statements(
                 execution_time=execution_time,
                 environment_naming_info=environment_naming_info,
                 engine_adapter=adapter,
+                selected_models=selected_models,
             )
         ]
     except Exception as e:

sqlmesh/core/plan/builder.py

@@ -129,6 +129,7 @@ def __init__(
         end_override_per_model: t.Optional[t.Dict[str, datetime]] = None,
         console: t.Optional[PlanBuilderConsole] = None,
         user_provided_flags: t.Optional[t.Dict[str, UserProvidedFlags]] = None,
+        selected_models: t.Optional[t.Set[str]] = None,
     ):
         self._context_diff = context_diff
         self._no_gaps = no_gaps
@@ -169,6 +170,7 @@ def __init__(
         self._console = console or get_console()
         self._choices: t.Dict[SnapshotId, SnapshotChangeCategory] = {}
         self._user_provided_flags = user_provided_flags
+        self._selected_models = selected_models
         self._explain = explain
 
         self._start = start
@@ -347,6 +349,7 @@ def build(self) -> Plan:
             ensure_finalized_snapshots=self._ensure_finalized_snapshots,
             ignore_cron=self._ignore_cron,
             user_provided_flags=self._user_provided_flags,
+            selected_models=self._selected_models,
         )
         self._latest_plan = plan
         return plan

sqlmesh/core/plan/definition.py

@@ -70,6 +70,8 @@ class Plan(PydanticModel, frozen=True):
     execution_time_: t.Optional[TimeLike] = Field(default=None, alias="execution_time")
 
     user_provided_flags: t.Optional[t.Dict[str, UserProvidedFlags]] = None
+    selected_models: t.Optional[t.Set[str]] = None
+    """Models that have been selected for this plan (used for dbt selected_resouces)"""
 
     @cached_property
     def start(self) -> TimeLike:
@@ -282,6 +284,7 @@ def to_evaluatable(self) -> EvaluatablePlan:
             },
             environment_statements=self.context_diff.environment_statements,
             user_provided_flags=self.user_provided_flags,
+            selected_models=self.selected_models,
         )
 
     @cached_property
@@ -319,6 +322,7 @@ class EvaluatablePlan(PydanticModel):
     disabled_restatement_models: t.Set[str]
     environment_statements: t.Optional[t.List[EnvironmentStatements]] = None
     user_provided_flags: t.Optional[t.Dict[str, UserProvidedFlags]] = None
+    selected_models: t.Optional[t.Set[str]] = None
 
     def is_selected_for_backfill(self, model_fqn: str) -> bool:
         return self.models_to_backfill is None or model_fqn in self.models_to_backfill

sqlmesh/core/plan/evaluator.py

@@ -137,6 +137,7 @@ def visit_before_all_stage(self, stage: stages.BeforeAllStage, plan: Evaluatable
             start=plan.start,
             end=plan.end,
             execution_time=plan.execution_time,
+            selected_models=plan.selected_models,
         )
 
     def visit_after_all_stage(self, stage: stages.AfterAllStage, plan: EvaluatablePlan) -> None:
@@ -150,6 +151,7 @@ def visit_after_all_stage(self, stage: stages.AfterAllStage, plan: EvaluatablePl
             start=plan.start,
             end=plan.end,
             execution_time=plan.execution_time,
+            selected_models=plan.selected_models,
         )
 
     def visit_create_snapshot_records_stage(
@@ -257,6 +259,7 @@ def visit_backfill_stage(self, stage: stages.BackfillStage, plan: EvaluatablePla
             allow_destructive_snapshots=plan.allow_destructive_models,
             allow_additive_snapshots=plan.allow_additive_models,
             selected_snapshot_ids=stage.selected_snapshot_ids,
+            selected_models=plan.selected_models,
         )
         if errors:
             raise PlanError("Plan application failed.")

sqlmesh/core/scheduler.py

@@ -416,6 +416,7 @@ def run_merged_intervals(
         start: t.Optional[TimeLike] = None,
         end: t.Optional[TimeLike] = None,
         allow_destructive_snapshots: t.Optional[t.Set[str]] = None,
+        selected_models: t.Optional[t.Set[str]] = None,
         allow_additive_snapshots: t.Optional[t.Set[str]] = None,
         selected_snapshot_ids: t.Optional[t.Set[SnapshotId]] = None,
         run_environment_statements: bool = False,
@@ -472,6 +473,7 @@ def run_merged_intervals(
                 start=start,
                 end=end,
                 execution_time=execution_time,
+                selected_models=selected_models,
             )
 
         snapshots_to_create = {
@@ -526,6 +528,7 @@ def run_node(node: SchedulingUnit) -> None:
                             allow_destructive_snapshots=allow_destructive_snapshots,
                             allow_additive_snapshots=allow_additive_snapshots,
                             target_table_exists=snapshot.snapshot_id not in snapshots_to_create,
+                            selected_models=selected_models,
                         )
 
                     evaluation_duration_ms = now_timestamp() - execution_start_ts
@@ -595,6 +598,7 @@ def run_node(node: SchedulingUnit) -> None:
                     start=start,
                     end=end,
                     execution_time=execution_time,
+                    selected_models=selected_models,
                 )
 
             self.state_sync.recycle()
@@ -798,6 +802,7 @@ def _run_or_audit(
             run_environment_statements=run_environment_statements,
             audit_only=audit_only,
             auto_restatement_triggers=auto_restatement_triggers,
+            selected_models=selected_snapshots or {s.name for s in merged_intervals},
         )
 
         return CompletionStatus.FAILURE if errors else CompletionStatus.SUCCESS

sqlmesh/dbt/adapter.py

@@ -180,6 +180,10 @@ def graph(self) -> t.Any:
             }
         )
 
+    @property
+    def selected_resources(self) -> t.List[str]:
+        return []
+
 
 class ParsetimeAdapter(BaseAdapter):
     def get_relation(self, database: str, schema: str, identifier: str) -> t.Optional[BaseRelation]:
@@ -280,6 +284,13 @@ def __init__(
     def graph(self) -> t.Any:
         return self.jinja_globals.get("flat_graph", super().graph)
 
+    @property
+    def selected_resources(self) -> t.List[str]:
+        selected_models = self.jinja_globals.get("selected_models")
+        if selected_models:
+            return [self._dbt_model_id(model) for model in sorted(selected_models)]
+        return []
+
     def get_relation(
         self, database: t.Optional[str], schema: str, identifier: str
     ) -> t.Optional[BaseRelation]:
@@ -504,3 +515,7 @@ def _normalize(self, input_table: exp.Table) -> exp.Table:
             normalized_table.set("db", normalized_table.this)
             normalized_table.set("this", None)
         return normalized_table
+
+    def _dbt_model_id(self, sqlmesh_model_name: str) -> str:
+        parts = [part.strip('"') for part in sqlmesh_model_name.split(".")]
+        return f"model.{parts[0]}.{parts[-1]}"

sqlmesh/dbt/builtin.py

@@ -459,7 +459,7 @@ def create_builtin_globals(
             "run_query": sql_execution.run_query,
             "statement": sql_execution.statement,
             "graph": adapter.graph,
-            "selected_resources": get_selected_resources(),
+            "selected_resources": adapter.selected_resources,
         }
     )
 
@@ -486,33 +486,3 @@ 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 = []