type alias, remove optionals, rename to user_provided_flags

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
@@ -1335,47 +1336,33 @@ 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]],
-            ],
-        ] = {
+        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": restate_models,
+            "restate_models": list(restate_models) if restate_models else None,
             "no_gaps": no_gaps,
             "skip_backfill": skip_backfill,
             "empty_backfill": empty_backfill,
             "forward_only": forward_only,
-            "allow_destructive_models": allow_destructive_models,
+            "allow_destructive_models": list(allow_destructive_models)
+            if allow_destructive_models
+            else None,
             "no_auto_categorization": no_auto_categorization,
             "effective_from": effective_from,
             "include_unmodified": include_unmodified,
-            "select_models": select_models,
-            "backfill_models": backfill_models,
+            "select_models": list(select_models) if select_models else None,
+            "backfill_models": list(backfill_models) if backfill_models else None,
             "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}
+        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
@@ -1519,7 +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_defined_flags=user_defined_flags,
+            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,18 +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_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,
+        user_provided_flags: t.Optional[t.Dict[str, UserProvidedFlags]] = None,
     ):
         self._context_diff = context_diff
         self._no_gaps = no_gaps
@@ -146,7 +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_defined_flags = user_defined_flags
+        self._user_provided_flags = user_provided_flags
 
         self._start = start
         if not self._start and (
@@ -293,7 +287,7 @@ def build(self) -> Plan:
             execution_time=self._execution_time,
             end_bounded=self._end_bounded,
             ensure_finalized_snapshots=self._ensure_finalized_snapshots,
-            user_defined_flags=self._user_defined_flags,
+            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,18 +64,7 @@ class Plan(PydanticModel, frozen=True):
     effective_from: t.Optional[TimeLike] = None
     execution_time: t.Optional[TimeLike] = 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
+    user_provided_flags: t.Optional[t.Dict[str, UserProvidedFlags]] = None
 
     @cached_property
     def start(self) -> TimeLike:
@@ -275,7 +265,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,
+            user_provided_flags=self.user_provided_flags,
         )
 
     @cached_property
@@ -308,18 +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_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
+    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

@@ -3107,20 +3107,20 @@ def test_set_choice_for_forward_only_model(make_snapshot):
     )
 
 
-def test_user_defined_flags(sushi_context: Context):
+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_defined_flags == expected_flags
+    assert plan_a.user_provided_flags == expected_flags
     evaluatable_plan = plan_a.to_evaluatable()
-    assert evaluatable_plan.user_defined_flags == expected_flags
+    assert evaluatable_plan.user_provided_flags == expected_flags
 
     plan_b = sushi_context.plan()
-    assert plan_b.user_defined_flags == {}
+    assert plan_b.user_provided_flags == {}
     evaluatable_plan_b = plan_b.to_evaluatable()
-    assert evaluatable_plan_b.user_defined_flags == {}
+    assert evaluatable_plan_b.user_provided_flags == {}
 
     context_diff = ContextDiff(
         environment="test_environment",
@@ -3144,11 +3144,11 @@ def test_user_defined_flags(sushi_context: Context):
         context_diff,
         DuckDBEngineAdapter.SCHEMA_DIFFER,
         forward_only=True,
-        user_defined_flags={"forward_only": True},
+        user_provided_flags={"forward_only": True},
     ).build()
-    assert plan_builder.user_defined_flags == {"forward_only": True}
+    assert plan_builder.user_provided_flags == {"forward_only": True}
     plan_builder = PlanBuilder(
         context_diff,
         DuckDBEngineAdapter.SCHEMA_DIFFER,
     ).build()
-    assert plan_builder.user_defined_flags == None
+    assert plan_builder.user_provided_flags == None