store plan flags

Author: toriwei

sqlmesh/core/context.py

@@ -1161,11 +1161,11 @@ def plan(
         end: t.Optional[TimeLike] = None,
         execution_time: t.Optional[TimeLike] = None,
         create_from: t.Optional[str] = None,
-        skip_tests: bool = False,
+        skip_tests: t.Optional[bool] = None,
         restate_models: t.Optional[t.Iterable[str]] = None,
-        no_gaps: bool = False,
-        skip_backfill: bool = False,
-        empty_backfill: bool = False,
+        no_gaps: t.Optional[bool] = None,
+        skip_backfill: t.Optional[bool] = None,
+        empty_backfill: t.Optional[bool] = None,
         forward_only: t.Optional[bool] = None,
         allow_destructive_models: t.Optional[t.Collection[str]] = None,
         no_prompts: t.Optional[bool] = None,
@@ -1178,9 +1178,9 @@ def plan(
         categorizer_config: t.Optional[CategorizerConfig] = None,
         enable_preview: t.Optional[bool] = None,
         no_diff: t.Optional[bool] = None,
-        run: bool = False,
-        diff_rendered: bool = False,
-        skip_linter: bool = False,
+        run: t.Optional[bool] = None,
+        diff_rendered: t.Optional[bool] = None,
+        skip_linter: t.Optional[bool] = None,
     ) -> Plan:
         """Interactively creates a plan.
 
@@ -1230,6 +1230,18 @@ def plan(
         Returns:
             The populated Plan object.
         """
+
+        flags = {
+            k: v for k, v in locals().items() if k not in {"self", "environment"} if v is not None
+        }
+        skip_tests = skip_tests or False
+        no_gaps = no_gaps or False
+        skip_backfill = skip_backfill or False
+        empty_backfill = empty_backfill or False
+        run = run or False
+        diff_rendered = diff_rendered or False
+        skip_linter = skip_linter or False
+
         plan_builder = self.plan_builder(
             environment,
             start=start,
@@ -1253,6 +1265,7 @@ def plan(
             run=run,
             diff_rendered=diff_rendered,
             skip_linter=skip_linter,
+            flags=flags,
         )
 
         if no_auto_categorization:
@@ -1295,6 +1308,7 @@ def plan_builder(
         run: bool = False,
         diff_rendered: bool = False,
         skip_linter: bool = False,
+        flags: t.Optional[t.Dict[str, t.Any]] = None,
     ) -> PlanBuilder:
         """Creates a plan builder.
 
@@ -1469,6 +1483,7 @@ def plan_builder(
             engine_schema_differ=self.engine_adapter.SCHEMA_DIFFER,
             interval_end_per_model=max_interval_end_per_model,
             console=self.console,
+            flags=flags,
         )
 
     def apply(

sqlmesh/core/plan/builder.py

@@ -107,6 +107,7 @@ def __init__(
         ensure_finalized_snapshots: bool = False,
         interval_end_per_model: t.Optional[t.Dict[str, int]] = None,
         console: t.Optional[PlanBuilderConsole] = None,
+        flags: t.Optional[t.Dict[str, t.Any]] = None,
     ):
         self._context_diff = context_diff
         self._no_gaps = no_gaps
@@ -134,6 +135,7 @@ def __init__(
         self._engine_schema_differ = engine_schema_differ
         self._console = console or get_console()
         self._choices: t.Dict[SnapshotId, SnapshotChangeCategory] = {}
+        self._flags = flags
 
         self._start = start
         if not self._start and (
@@ -280,6 +282,7 @@ def build(self) -> Plan:
             execution_time=self._execution_time,
             end_bounded=self._end_bounded,
             ensure_finalized_snapshots=self._ensure_finalized_snapshots,
+            flags=self._flags,
         )
         self._latest_plan = plan
         return plan

sqlmesh/core/plan/definition.py

@@ -63,6 +63,8 @@ class Plan(PydanticModel, frozen=True):
     effective_from: t.Optional[TimeLike] = None
     execution_time: t.Optional[TimeLike] = None
 
+    flags: t.Optional[t.Dict[str, t.Any]] = None
+
     @cached_property
     def start(self) -> TimeLike:
         if self.provided_start is not None:

tests/core/test_plan.py

@@ -3105,3 +3105,10 @@ def test_set_choice_for_forward_only_model(make_snapshot):
         plan.snapshots[updated_snapshot.snapshot_id].change_category
         == SnapshotChangeCategory.FORWARD_ONLY
     )
+
+
+def test_flags(sushi_context: Context):
+    plan = sushi_context.plan(no_prompts=True, run=True, execution_time="2025-01-01")
+    sushi_context.apply(plan)
+
+    assert plan.flags == {"no_prompts": True, "run": True, "execution_time": "2025-01-01"}