feat: store user provided plan flags (#4319)

Author: toriwei

sqlmesh/core/context.py

@@ -88,6 +88,7 @@
     NotificationTargetManager,
 )
 from sqlmesh.core.plan import Plan, PlanBuilder, SnapshotIntervals
+from sqlmesh.core.plan.definition import UserProvidedFlags
 from sqlmesh.core.reference import ReferenceGraph
 from sqlmesh.core.scheduler import Scheduler, CompletionStatus
 from sqlmesh.core.schema_loader import create_external_models_file
@@ -1161,11 +1162,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 +1179,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.
 
@@ -1278,11 +1279,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,
@@ -1292,9 +1293,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,
+        run: t.Optional[bool] = None,
+        diff_rendered: t.Optional[bool] = None,
+        skip_linter: t.Optional[bool] = None,
     ) -> PlanBuilder:
         """Creates a plan builder.
 
@@ -1335,6 +1336,42 @@ def plan_builder(
         Returns:
             The plan builder.
         """
+        kwargs: t.Dict[str, t.Optional[UserProvidedFlags]] = {
+            "start": start,
+            "end": end,
+            "execution_time": execution_time,
+            "create_from": create_from,
+            "skip_tests": skip_tests,
+            "restate_models": list(restate_models) if restate_models is not None else None,
+            "no_gaps": no_gaps,
+            "skip_backfill": skip_backfill,
+            "empty_backfill": empty_backfill,
+            "forward_only": forward_only,
+            "allow_destructive_models": list(allow_destructive_models)
+            if allow_destructive_models is not None
+            else None,
+            "no_auto_categorization": no_auto_categorization,
+            "effective_from": effective_from,
+            "include_unmodified": include_unmodified,
+            "select_models": list(select_models) if select_models is not None else None,
+            "backfill_models": list(backfill_models) if backfill_models is not None else None,
+            "enable_preview": enable_preview,
+            "run": run,
+            "diff_rendered": diff_rendered,
+            "skip_linter": skip_linter,
+        }
+        user_provided_flags: t.Dict[str, UserProvidedFlags] = {
+            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
@@ -1469,6 +1506,7 @@ def plan_builder(
             engine_schema_differ=self.engine_adapter.SCHEMA_DIFFER,
             interval_end_per_model=max_interval_end_per_model,
             console=self.console,
+            user_provided_flags=user_provided_flags,
         )
 
     def apply(

sqlmesh/core/plan/builder.py

@@ -15,7 +15,12 @@
 )
 from sqlmesh.core.context_diff import ContextDiff
 from sqlmesh.core.environment import EnvironmentNamingInfo
-from sqlmesh.core.plan.definition import Plan, SnapshotMapping, earliest_interval_start
+from sqlmesh.core.plan.definition import (
+    Plan,
+    SnapshotMapping,
+    UserProvidedFlags,
+    earliest_interval_start,
+)
 from sqlmesh.core.schema_diff import SchemaDiffer, has_drop_alteration, get_dropped_column_names
 from sqlmesh.core.snapshot import (
     DeployabilityIndex,
@@ -107,6 +112,7 @@ def __init__(
         ensure_finalized_snapshots: bool = False,
         interval_end_per_model: t.Optional[t.Dict[str, int]] = None,
         console: t.Optional[PlanBuilderConsole] = None,
+        user_provided_flags: t.Optional[t.Dict[str, UserProvidedFlags]] = None,
     ):
         self._context_diff = context_diff
         self._no_gaps = no_gaps
@@ -134,6 +140,7 @@ def __init__(
         self._engine_schema_differ = engine_schema_differ
         self._console = console or get_console()
         self._choices: t.Dict[SnapshotId, SnapshotChangeCategory] = {}
+        self._user_provided_flags = user_provided_flags
 
         self._start = start
         if not self._start and (
@@ -280,6 +287,7 @@ def build(self) -> Plan:
             execution_time=self._execution_time,
             end_bounded=self._end_bounded,
             ensure_finalized_snapshots=self._ensure_finalized_snapshots,
+            user_provided_flags=self._user_provided_flags,
         )
         self._latest_plan = plan
         return plan

sqlmesh/core/plan/definition.py

@@ -28,6 +28,7 @@
 from sqlmesh.utils.pydantic import PydanticModel
 
 SnapshotMapping = t.Dict[SnapshotId, t.Set[SnapshotId]]
+UserProvidedFlags = t.Union[TimeLike, str, bool, t.List[str]]
 
 
 class Plan(PydanticModel, frozen=True):
@@ -63,6 +64,8 @@ class Plan(PydanticModel, frozen=True):
     effective_from: t.Optional[TimeLike] = None
     execution_time: t.Optional[TimeLike] = None
 
+    user_provided_flags: t.Optional[t.Dict[str, UserProvidedFlags]] = None
+
     @cached_property
     def start(self) -> TimeLike:
         if self.provided_start is not None:
@@ -262,6 +265,7 @@ def to_evaluatable(self) -> EvaluatablePlan:
                 if s.is_model and s.model.disable_restatement
             },
             environment_statements=self.context_diff.environment_statements,
+            user_provided_flags=self.user_provided_flags,
         )
 
     @cached_property
@@ -294,6 +298,7 @@ class EvaluatablePlan(PydanticModel):
     execution_time: t.Optional[TimeLike] = None
     disabled_restatement_models: t.Set[str]
     environment_statements: t.Optional[t.List[EnvironmentStatements]] = None
+    user_provided_flags: t.Optional[t.Dict[str, UserProvidedFlags]] = 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

@@ -3105,3 +3105,50 @@ def test_set_choice_for_forward_only_model(make_snapshot):
         plan.snapshots[updated_snapshot.snapshot_id].change_category
         == SnapshotChangeCategory.FORWARD_ONLY
     )
+
+
+def test_user_provided_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_provided_flags == expected_flags
+    evaluatable_plan = plan_a.to_evaluatable()
+    assert evaluatable_plan.user_provided_flags == expected_flags
+
+    plan_b = sushi_context.plan()
+    assert plan_b.user_provided_flags == {}
+    evaluatable_plan_b = plan_b.to_evaluatable()
+    assert evaluatable_plan_b.user_provided_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_provided_flags={"forward_only": True},
+    ).build()
+    assert plan_builder.user_provided_flags == {"forward_only": True}
+    plan_builder = PlanBuilder(
+        context_diff,
+        DuckDBEngineAdapter.SCHEMA_DIFFER,
+    ).build()
+    assert plan_builder.user_provided_flags == None