move to plan_builder, edit possible flags, update EvaluatablePlan

Author: toriwei

sqlmesh/core/context.py

@@ -1230,18 +1230,6 @@ 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,
@@ -1265,7 +1253,6 @@ def plan(
             run=run,
             diff_rendered=diff_rendered,
             skip_linter=skip_linter,
-            flags=flags,
         )
 
         if no_auto_categorization:
@@ -1291,11 +1278,11 @@ def plan_builder(
         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_auto_categorization: t.Optional[bool] = None,
@@ -1305,10 +1292,9 @@ def plan_builder(
         backfill_models: t.Optional[t.Collection[str]] = None,
         categorizer_config: t.Optional[CategorizerConfig] = None,
         enable_preview: t.Optional[bool] = None,
-        run: bool = False,
-        diff_rendered: bool = False,
-        skip_linter: bool = False,
-        flags: t.Optional[t.Dict[str, t.Any]] = None,
+        run: t.Optional[bool] = None,
+        diff_rendered: t.Optional[bool] = None,
+        skip_linter: t.Optional[bool] = None,
     ) -> PlanBuilder:
         """Creates a plan builder.
 
@@ -1349,6 +1335,56 @@ def plan_builder(
         Returns:
             The plan builder.
         """
+        kwargs: t.Dict[
+            str,
+            t.Union[
+                t.Optional[TimeLike],
+                t.Optional[str],
+                t.Optional[bool],
+                t.Optional[t.Iterable[str]],
+                t.Optional[t.Collection[str]],
+            ],
+        ] = {
+            "start": start,
+            "end": end,
+            "execution_time": execution_time,
+            "create_from": create_from,
+            "skip_tests": skip_tests,
+            "restate_models": restate_models,
+            "no_gaps": no_gaps,
+            "skip_backfill": skip_backfill,
+            "empty_backfill": empty_backfill,
+            "forward_only": forward_only,
+            "allow_destructive_models": allow_destructive_models,
+            "no_auto_categorization": no_auto_categorization,
+            "effective_from": effective_from,
+            "include_unmodified": include_unmodified,
+            "select_models": select_models,
+            "backfill_models": backfill_models,
+            "enable_preview": enable_preview,
+            "run": run,
+            "diff_rendered": diff_rendered,
+            "skip_linter": skip_linter,
+        }
+        user_defined_flags: t.Dict[
+            str,
+            t.Union[
+                t.Optional[TimeLike],
+                t.Optional[str],
+                t.Optional[bool],
+                t.Optional[t.Iterable[str]],
+                t.Optional[t.Collection[str]],
+            ],
+        ] = {k: v for k, v in kwargs.items() 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
+
         environment = environment or self.config.default_target_environment
         environment = Environment.sanitize_name(environment)
         is_dev = environment != c.PROD
@@ -1483,7 +1519,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,
+            user_defined_flags=user_defined_flags,
         )
 
     def apply(

sqlmesh/core/plan/builder.py

@@ -107,7 +107,18 @@ 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,
+        user_defined_flags: t.Optional[
+            t.Dict[
+                str,
+                t.Union[
+                    t.Optional[TimeLike],
+                    t.Optional[str],
+                    t.Optional[bool],
+                    t.Optional[t.Iterable[str]],
+                    t.Optional[t.Collection[str]],
+                ],
+            ]
+        ] = None,
     ):
         self._context_diff = context_diff
         self._no_gaps = no_gaps
@@ -135,7 +146,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._user_defined_flags = user_defined_flags
 
         self._start = start
         if not self._start and (
@@ -282,7 +293,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,
+            user_defined_flags=self._user_defined_flags,
         )
         self._latest_plan = plan
         return plan

sqlmesh/core/plan/definition.py

@@ -63,7 +63,18 @@ 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
+    user_defined_flags: t.Optional[
+        t.Dict[
+            str,
+            t.Union[
+                t.Optional[TimeLike],
+                t.Optional[str],
+                t.Optional[bool],
+                t.Optional[t.Iterable[str]],
+                t.Optional[t.Collection[str]],
+            ],
+        ]
+    ] = None
 
     @cached_property
     def start(self) -> TimeLike:
@@ -264,6 +275,7 @@ def to_evaluatable(self) -> EvaluatablePlan:
                 if s.is_model and s.model.disable_restatement
             },
             environment_statements=self.context_diff.environment_statements,
+            user_defined_flags=self.user_defined_flags,
         )
 
     @cached_property
@@ -296,6 +308,18 @@ class EvaluatablePlan(PydanticModel):
     execution_time: t.Optional[TimeLike] = None
     disabled_restatement_models: t.Set[str]
     environment_statements: t.Optional[t.List[EnvironmentStatements]] = None
+    user_defined_flags: t.Optional[
+        t.Dict[
+            str,
+            t.Union[
+                t.Optional[TimeLike],
+                t.Optional[str],
+                t.Optional[bool],
+                t.Optional[t.Iterable[str]],
+                t.Optional[t.Collection[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

tests/core/test_plan.py

@@ -3107,8 +3107,48 @@ def test_set_choice_for_forward_only_model(make_snapshot):
     )
 
 
-def test_flags(sushi_context: Context):
-    plan = sushi_context.plan(no_prompts=True, run=True, execution_time="2025-01-01")
-    sushi_context.apply(plan)
+def test_user_defined_flags(sushi_context: Context):
+    expected_flags = {
+        "run": True,
+        "execution_time": "2025-01-01",
+    }
+    plan_a = sushi_context.plan(no_prompts=True, run=True, execution_time="2025-01-01")
+    assert plan_a.user_defined_flags == expected_flags
+    evaluatable_plan = plan_a.to_evaluatable()
+    assert evaluatable_plan.user_defined_flags == expected_flags
 
-    assert plan.flags == {"no_prompts": True, "run": True, "execution_time": "2025-01-01"}
+    plan_b = sushi_context.plan()
+    assert plan_b.user_defined_flags == {}
+    evaluatable_plan_b = plan_b.to_evaluatable()
+    assert evaluatable_plan_b.user_defined_flags == {}
+
+    context_diff = ContextDiff(
+        environment="test_environment",
+        is_new_environment=True,
+        is_unfinalized_environment=False,
+        normalize_environment_name=True,
+        create_from="prod",
+        create_from_env_exists=True,
+        added=set(),
+        removed_snapshots={},
+        modified_snapshots={},
+        snapshots={},
+        new_snapshots={},
+        previous_plan_id=None,
+        previously_promoted_snapshot_ids=set(),
+        previous_finalized_snapshots=None,
+        previous_gateway_managed_virtual_layer=False,
+        gateway_managed_virtual_layer=False,
+    )
+    plan_builder = PlanBuilder(
+        context_diff,
+        DuckDBEngineAdapter.SCHEMA_DIFFER,
+        forward_only=True,
+        user_defined_flags={"forward_only": True},
+    ).build()
+    assert plan_builder.user_defined_flags == {"forward_only": True}
+    plan_builder = PlanBuilder(
+        context_diff,
+        DuckDBEngineAdapter.SCHEMA_DIFFER,
+    ).build()
+    assert plan_builder.user_defined_flags == None