Merge branch 'main' into indirect_change_to_materialized_view_is_breaking

Author: xardasos

pyproject.toml

@@ -225,7 +225,8 @@ module = [
     "pydantic_core.*",
     "dlt.*",
     "bigframes.*",
-    "json_stream.*"
+    "json_stream.*",
+    "duckdb.*"
 ]
 ignore_missing_imports = true
 

sqlmesh/core/context.py

@@ -1429,6 +1429,7 @@ def plan_builder(
         explain: t.Optional[bool] = None,
         ignore_cron: t.Optional[bool] = None,
         min_intervals: t.Optional[int] = None,
+        always_include_local_changes: t.Optional[bool] = None,
     ) -> PlanBuilder:
         """Creates a plan builder.
 
@@ -1467,6 +1468,8 @@ def plan_builder(
             diff_rendered: Whether the diff should compare raw vs rendered models
             min_intervals: Adjust the plan start date on a per-model basis in order to ensure at least this many intervals are covered
                 on every model when checking for missing intervals
+            always_include_local_changes: Usually when restatements are present, local changes in the filesystem are ignored.
+                However, it can be desirable to deploy changes + restatements in the same plan, so this flag overrides the default behaviour.
 
         Returns:
             The plan builder.
@@ -1583,13 +1586,20 @@ def plan_builder(
                 "Selector did not return any models. Please check your model selection and try again."
             )
 
+        if always_include_local_changes is None:
+            # default behaviour - if restatements are detected; we operate entirely out of state and ignore local changes
+            force_no_diff = restate_models is not None or (
+                backfill_models is not None and not backfill_models
+            )
+        else:
+            force_no_diff = not always_include_local_changes
+
         snapshots = self._snapshots(models_override)
         context_diff = self._context_diff(
             environment or c.PROD,
             snapshots=snapshots,
             create_from=create_from,
-            force_no_diff=restate_models is not None
-            or (backfill_models is not None and not backfill_models),
+            force_no_diff=force_no_diff,
             ensure_finalized_snapshots=self.config.plan.use_finalized_state,
             diff_rendered=diff_rendered,
             always_recreate_environment=self.config.plan.always_recreate_environment,
@@ -1644,13 +1654,22 @@ def plan_builder(
         elif forward_only is None:
             forward_only = self.config.plan.forward_only
 
+        # When handling prod restatements, only clear intervals from other model versions if we are using full virtual environments
+        # If we are not, then there is no point, because none of the data in dev environments can be promoted by definition
+        restate_all_snapshots = (
+            expanded_restate_models is not None
+            and not is_dev
+            and self.config.virtual_environment_mode.is_full
+        )
+
         return self.PLAN_BUILDER_TYPE(
             context_diff=context_diff,
             start=start,
             end=end,
             execution_time=execution_time,
             apply=self.apply,
             restate_models=expanded_restate_models,
+            restate_all_snapshots=restate_all_snapshots,
             backfill_models=backfill_models,
             no_gaps=no_gaps,
             skip_backfill=skip_backfill,

sqlmesh/core/plan/builder.py

@@ -65,6 +65,9 @@ class PlanBuilder:
         restate_models: A list of models for which the data should be restated for the time range
             specified in this plan. Note: models defined outside SQLMesh (external) won't be a part
             of the restatement.
+        restate_all_snapshots: If restatements are present, this flag indicates whether or not the intervals
+            being restated should be cleared from state for other versions of this model (typically, versions that are present in other environments).
+            If set to None, the default behaviour is to not clear anything unless the target environment is prod.
         backfill_models: A list of fully qualified model names for which the data should be backfilled as part of this plan.
         no_gaps:  Whether to ensure that new snapshots for nodes that are already a
             part of the target environment have no data gaps when compared against previous
@@ -103,6 +106,7 @@ def __init__(
         execution_time: t.Optional[TimeLike] = None,
         apply: t.Optional[t.Callable[[Plan], None]] = None,
         restate_models: t.Optional[t.Iterable[str]] = None,
+        restate_all_snapshots: bool = False,
         backfill_models: t.Optional[t.Iterable[str]] = None,
         no_gaps: bool = False,
         skip_backfill: bool = False,
@@ -154,6 +158,7 @@ def __init__(
         self._auto_categorization_enabled = auto_categorization_enabled
         self._include_unmodified = include_unmodified
         self._restate_models = set(restate_models) if restate_models is not None else None
+        self._restate_all_snapshots = restate_all_snapshots
         self._effective_from = effective_from
 
         # note: this deliberately doesnt default to now() here.
@@ -277,7 +282,6 @@ def build(self) -> Plan:
         if self._latest_plan:
             return self._latest_plan
 
-        self._ensure_no_new_snapshots_with_restatements()
         self._ensure_new_env_with_changes()
         self._ensure_valid_date_range()
         self._ensure_no_broken_references()
@@ -340,6 +344,7 @@ def build(self) -> Plan:
             deployability_index=deployability_index,
             selected_models_to_restate=self._restate_models,
             restatements=restatements,
+            restate_all_snapshots=self._restate_all_snapshots,
             start_override_per_model=self._start_override_per_model,
             end_override_per_model=end_override_per_model,
             selected_models_to_backfill=self._backfill_models,
@@ -867,15 +872,6 @@ def _ensure_no_broken_references(self) -> None:
                     f"""Removed {broken_references_msg} are referenced in '{snapshot.name}'. Please remove broken references before proceeding."""
                 )
 
-    def _ensure_no_new_snapshots_with_restatements(self) -> None:
-        if self._restate_models is not None and (
-            self._context_diff.new_snapshots or self._context_diff.modified_snapshots
-        ):
-            raise PlanError(
-                "Model changes and restatements can't be a part of the same plan. "
-                "Revert or apply changes before proceeding with restatements."
-            )
-
     def _ensure_new_env_with_changes(self) -> None:
         if (
             self._is_dev

sqlmesh/core/plan/definition.py

@@ -67,6 +67,8 @@ class Plan(PydanticModel, frozen=True):
     Note that dev previews are also considered restatements, so :selected_models_to_restate can be empty
     while :restatements is still populated with dev previews
     """
+    restate_all_snapshots: bool
+    """Whether or not to clear intervals from state for other versions of the models listed in :restatements"""
 
     start_override_per_model: t.Optional[t.Dict[str, datetime]]
     end_override_per_model: t.Optional[t.Dict[str, datetime]]
@@ -268,6 +270,7 @@ def to_evaluatable(self) -> EvaluatablePlan:
             skip_backfill=self.skip_backfill,
             empty_backfill=self.empty_backfill,
             restatements={s.name: i for s, i in self.restatements.items()},
+            restate_all_snapshots=self.restate_all_snapshots,
             is_dev=self.is_dev,
             allow_destructive_models=self.allow_destructive_models,
             allow_additive_models=self.allow_additive_models,
@@ -312,6 +315,7 @@ class EvaluatablePlan(PydanticModel):
     skip_backfill: bool
     empty_backfill: bool
     restatements: t.Dict[str, Interval]
+    restate_all_snapshots: bool
     is_dev: bool
     allow_destructive_models: t.Set[str]
     allow_additive_models: t.Set[str]

sqlmesh/core/plan/stages.py

@@ -12,7 +12,9 @@
     Snapshot,
     SnapshotTableInfo,
     SnapshotId,
+    snapshots_to_dag,
 )
+from sqlmesh.utils.errors import PlanError
 
 
 @dataclass
@@ -248,6 +250,7 @@ def build(self, plan: EvaluatablePlan) -> t.List[PlanStage]:
         stored_snapshots = self.state_reader.get_snapshots(plan.environment.snapshots)
         snapshots = {**new_snapshots, **stored_snapshots}
         snapshots_by_name = {s.name: s for s in snapshots.values()}
+        dag = snapshots_to_dag(snapshots.values())
 
         all_selected_for_backfill_snapshots = {
             s.snapshot_id for s in snapshots.values() if plan.is_selected_for_backfill(s.name)
@@ -271,8 +274,15 @@ def build(self, plan: EvaluatablePlan) -> t.List[PlanStage]:
             after_promote_snapshots = all_selected_for_backfill_snapshots - before_promote_snapshots
             deployability_index = DeployabilityIndex.all_deployable()
 
+            snapshot_ids_with_schema_migration = [
+                s.snapshot_id for s in snapshots.values() if s.requires_schema_migration_in_prod
+            ]
+            # Include all upstream dependencies of snapshots that require schema migration to make sure
+            # the upstream tables are created before the schema updates are applied
             snapshots_with_schema_migration = [
-                s for s in snapshots.values() if s.requires_schema_migration_in_prod
+                snapshots[s_id]
+                for s_id in dag.subdag(*snapshot_ids_with_schema_migration)
+                if snapshots[s_id].supports_schema_migration_in_prod
             ]
 
         snapshots_to_intervals = self._missing_intervals(
@@ -452,13 +462,18 @@ def _get_after_all_stage(
     def _get_restatement_stage(
         self, plan: EvaluatablePlan, snapshots_by_name: t.Dict[str, Snapshot]
     ) -> t.Optional[RestatementStage]:
-        if not plan.restatements or plan.is_dev:
-            # The RestatementStage to clear intervals from state across all environments is not needed for plans against dev, only prod
-            return None
+        if plan.restate_all_snapshots:
+            if plan.is_dev:
+                raise PlanError(
+                    "Clearing intervals from state across dev model versions is only valid for prod plans"
+                )
 
-        return RestatementStage(
-            all_snapshots=snapshots_by_name,
-        )
+            if plan.restatements:
+                return RestatementStage(
+                    all_snapshots=snapshots_by_name,
+                )
+
+        return None
 
     def _get_physical_layer_update_stage(
         self,

sqlmesh/core/snapshot/definition.py

@@ -1477,19 +1477,19 @@ def expiration_ts(self) -> int:
             check_categorical_relative_expression=False,
         )
 
+    @property
+    def supports_schema_migration_in_prod(self) -> bool:
+        """Returns whether or not this snapshot supports schema migration when deployed to production."""
+        return self.is_paused and self.is_model and not self.is_symbolic
+
     @property
     def requires_schema_migration_in_prod(self) -> bool:
         """Returns whether or not this snapshot requires a schema migration when deployed to production."""
-        return (
-            self.is_paused
-            and self.is_model
-            and self.is_materialized
-            and (
-                (self.previous_version and self.previous_version.version == self.version)
-                or self.model.forward_only
-                or bool(self.model.physical_version)
-                or not self.virtual_environment_mode.is_full
-            )
+        return self.supports_schema_migration_in_prod and (
+            (self.previous_version and self.previous_version.version == self.version)
+            or self.model.forward_only
+            or bool(self.model.physical_version)
+            or not self.virtual_environment_mode.is_full
         )
 
     @property

sqlmesh/core/snapshot/evaluator.py

@@ -489,15 +489,14 @@ def migrate(
         allow_destructive_snapshots = allow_destructive_snapshots or set()
         allow_additive_snapshots = allow_additive_snapshots or set()
         snapshots_by_name = {s.name: s for s in snapshots.values()}
-        snapshots_with_data_objects = [snapshots[s_id] for s_id in target_data_objects]
         with self.concurrent_context():
             # Only migrate snapshots for which there's an existing data object
             concurrent_apply_to_snapshots(
-                snapshots_with_data_objects,
+                snapshots_by_name.values(),
                 lambda s: self._migrate_snapshot(
                     s,
                     snapshots_by_name,
-                    target_data_objects[s.snapshot_id],
+                    target_data_objects.get(s.snapshot_id),
                     allow_destructive_snapshots,
                     allow_additive_snapshots,
                     self.get_adapter(s.model_gateway),
@@ -747,11 +746,15 @@ def _evaluate_snapshot(
             adapter.execute(model.render_pre_statements(**render_statements_kwargs))
 
             if not target_table_exists or (model.is_seed and not snapshot.intervals):
-                columns_to_types_provided = (
+                # Only create the empty table if the columns were provided explicitly by the user
+                should_create_empty_table = (
                     model.kind.is_materialized
                     and model.columns_to_types_
                     and columns_to_types_all_known(model.columns_to_types_)
                 )
+                if not should_create_empty_table:
+                    # Or if the model is self-referential and its query is fully annotated with types
+                    should_create_empty_table = model.depends_on_self and model.annotated
                 if self._can_clone(snapshot, deployability_index):
                     self._clone_snapshot_in_dev(
                         snapshot=snapshot,
@@ -764,7 +767,7 @@ def _evaluate_snapshot(
                     )
                     runtime_stage = RuntimeStage.EVALUATING
                     target_table_exists = True
-                elif columns_to_types_provided or model.is_seed or model.kind.is_scd_type_2:
+                elif should_create_empty_table or model.is_seed or model.kind.is_scd_type_2:
                     self._execute_create(
                         snapshot=snapshot,
                         table_name=target_table_name,
@@ -1059,7 +1062,7 @@ def _migrate_snapshot(
         adapter: EngineAdapter,
         deployability_index: DeployabilityIndex,
     ) -> None:
-        if not snapshot.requires_schema_migration_in_prod:
+        if not snapshot.is_model or snapshot.is_symbolic:
             return
 
         deployability_index = DeployabilityIndex.all_deployable()
@@ -1081,20 +1084,32 @@ def _migrate_snapshot(
             ):
                 table_exists = False
 
+            rendered_physical_properties = snapshot.model.render_physical_properties(
+                **render_kwargs
+            )
+
             if table_exists:
                 self._migrate_target_table(
                     target_table_name=target_table_name,
                     snapshot=snapshot,
                     snapshots=snapshots,
                     deployability_index=deployability_index,
                     render_kwargs=render_kwargs,
-                    rendered_physical_properties=snapshot.model.render_physical_properties(
-                        **render_kwargs
-                    ),
+                    rendered_physical_properties=rendered_physical_properties,
                     allow_destructive_snapshots=allow_destructive_snapshots,
                     allow_additive_snapshots=allow_additive_snapshots,
                     run_pre_post_statements=True,
                 )
+            else:
+                self._execute_create(
+                    snapshot=snapshot,
+                    table_name=snapshot.table_name(is_deployable=True),
+                    is_table_deployable=True,
+                    deployability_index=deployability_index,
+                    create_render_kwargs=render_kwargs,
+                    rendered_physical_properties=rendered_physical_properties,
+                    dry_run=True,
+                )
 
     def _migrate_target_table(
         self,

sqlmesh_dbt/cli.py

@@ -90,7 +90,9 @@ def dbt(
 @click.option(
     "-f",
     "--full-refresh",
-    help="If specified, dbt will drop incremental models and fully-recalculate the incremental table from the model definition.",
+    is_flag=True,
+    default=False,
+    help="If specified, sqlmesh will drop incremental models and fully-recalculate the incremental table from the model definition.",
 )
 @click.option(
     "--env",