Feat: Add an option to not run upstream dependencies when selecting models in the command (#3413)

Author: izeigerman

sqlmesh/cli/main.py

@@ -457,6 +457,11 @@ def plan(
     type=int,
     help="If set, the command will exit with the specified code if the run is interrupted by an update to the target environment.",
 )
+@click.option(
+    "--no-auto-upstream",
+    is_flag=True,
+    help="Do not automatically include upstream models. Only applicable when --select-model is used. Note: this may result in missing / invalid data for the selected models.",
+)
 @click.pass_context
 @error_handler
 @cli_analytics

sqlmesh/core/context.py

@@ -597,6 +597,7 @@ def run(
         ignore_cron: bool = False,
         select_models: t.Optional[t.Collection[str]] = None,
         exit_on_env_update: t.Optional[int] = None,
+        no_auto_upstream: bool = False,
     ) -> bool:
         """Run the entire dag through the scheduler.
 
@@ -611,6 +612,7 @@ def run(
                 upstream dependencies of selected models will also be evaluated.
             exit_on_env_update: If set, exits with the provided code if the run is interrupted by an update
                 to the target environment.
+            no_auto_upstream: Whether to not force upstream models to run. Only applicable when using `select_models`.
 
         Returns:
             True if the run was successful, False otherwise.
@@ -679,6 +681,7 @@ def _has_environment_changed() -> bool:
                     ignore_cron=ignore_cron,
                     select_models=select_models,
                     circuit_breaker=_has_environment_changed,
+                    no_auto_upstream=no_auto_upstream,
                 )
                 done = True
             except CircuitBreakerError:
@@ -1886,6 +1889,7 @@ def _run(
         ignore_cron: bool,
         select_models: t.Optional[t.Collection[str]],
         circuit_breaker: t.Optional[t.Callable[[], bool]],
+        no_auto_upstream: bool,
     ) -> bool:
         scheduler = self.scheduler(environment=environment)
         snapshots = scheduler.snapshots
@@ -1898,7 +1902,9 @@ def _run(
             for fqn, model in models.items():
                 dag.add(fqn, model.depends_on)
             model_selector = self._new_selector(models=models, dag=dag)
-            select_models = set(dag.subdag(*model_selector.expand_model_selections(select_models)))
+            select_models = set(model_selector.expand_model_selections(select_models))
+            if not no_auto_upstream:
+                select_models = set(dag.subdag(*select_models))
 
         return scheduler.run(
             environment,

sqlmesh/magics.py

@@ -462,6 +462,11 @@ def plan(self, context: Context, line: str) -> None:
         type=int,
         help="If set, the command will exit with the specified code if the run is interrupted by an update to the target environment.",
     )
+    @argument(
+        "--no-auto-upstream",
+        action="store_true",
+        help="Do not automatically include upstream models. Only applicable when --select-model is used. Note: this may result in missing / invalid data for the selected models.",
+    )
     @line_magic
     @pass_sqlmesh_context
     def run_dag(self, context: Context, line: str) -> None:
@@ -476,6 +481,7 @@ def run_dag(self, context: Context, line: str) -> None:
             ignore_cron=args.ignore_cron,
             select_models=args.select_model,
             exit_on_env_update=args.exit_on_env_update,
+            no_auto_upstream=args.no_auto_upstream,
         )
         if not success:
             raise SQLMeshError("Error Running DAG. Check logs for details.")

tests/core/test_integration.py

@@ -1159,6 +1159,42 @@ def test_run_with_select_models(
         }
 
 
+@freeze_time("2023-01-08 15:00:00")
+def test_run_with_select_models_no_auto_upstream(
+    init_and_plan_context: t.Callable,
+):
+    context, _ = init_and_plan_context("examples/sushi")
+
+    model = context.get_model("sushi.waiter_revenue_by_day")
+    model = SqlModel.parse_obj({**model.dict(), "audits": []})
+    context.upsert_model(model)
+
+    context.plan("prod", no_prompts=True, skip_tests=True, auto_apply=True)
+
+    with freeze_time("2023-01-09 00:00:00"):
+        assert context.run(select_models=["*waiter_revenue_by_day"], no_auto_upstream=True)
+
+        snapshots = context.state_sync.state_sync.get_snapshots(context.snapshots.values())
+        # Only waiter_revenue_by_day should be backfilled up to 2023-01-09.
+        assert {s.name: s.intervals[0][1] for s in snapshots.values() if s.intervals} == {
+            '"memory"."sushi"."waiter_revenue_by_day"': to_timestamp("2023-01-09"),
+            '"memory"."sushi"."order_items"': to_timestamp("2023-01-08"),
+            '"memory"."sushi"."orders"': to_timestamp("2023-01-08"),
+            '"memory"."sushi"."items"': to_timestamp("2023-01-08"),
+            '"memory"."sushi"."customer_revenue_lifetime"': to_timestamp("2023-01-08"),
+            '"memory"."sushi"."customer_revenue_by_day"': to_timestamp("2023-01-08"),
+            '"memory"."sushi"."waiter_names"': to_timestamp("2023-01-08"),
+            '"memory"."sushi"."raw_marketing"': to_timestamp("2023-01-08"),
+            '"memory"."sushi"."marketing"': to_timestamp("2023-01-08"),
+            '"memory"."sushi"."waiter_as_customer_by_day"': to_timestamp("2023-01-08"),
+            '"memory"."sushi"."top_waiters"': to_timestamp("2023-01-08"),
+            '"memory"."raw"."demographics"': to_timestamp("2023-01-08"),
+            "assert_item_price_above_zero": to_timestamp("2023-01-08"),
+            '"memory"."sushi"."active_customers"': to_timestamp("2023-01-08"),
+            '"memory"."sushi"."customers"': to_timestamp("2023-01-08"),
+        }
+
+
 @freeze_time("2023-01-08 15:00:00")
 def test_select_models(init_and_plan_context: t.Callable):
     context, plan = init_and_plan_context("examples/sushi")