move more of creation into the scheduler

Author: izeigerman

sqlmesh/core/plan/evaluator.py

@@ -254,6 +254,7 @@ def visit_backfill_stage(self, stage: stages.BackfillStage, plan: EvaluatablePla
             start=plan.start,
             end=plan.end,
             allow_destructive_snapshots=plan.allow_destructive_models,
+            selected_snapshot_ids=stage.selected_snapshot_ids,
         )
         if errors:
             raise PlanError("Plan application failed.")

sqlmesh/core/plan/stages.py

@@ -116,12 +116,14 @@ class BackfillStage:
     Args:
         snapshot_to_intervals: Intervals to backfill. This collection can be empty in which case no backfill is needed.
             This can be useful to report the lack of backfills back to the user.
+        selected_snapshot_ids: The snapshots to include in the run DAG.
         all_snapshots: All snapshots in the plan by name.
         deployability_index: Deployability index for this stage.
         before_promote: Whether this stage is before the promotion stage.
     """
 
     snapshot_to_intervals: SnapshotToIntervals
+    selected_snapshot_ids: t.Set[SnapshotId]
     all_snapshots: t.Dict[str, Snapshot]
     deployability_index: DeployabilityIndex
     before_promote: bool = True
@@ -298,26 +300,13 @@ def build(self, plan: EvaluatablePlan) -> t.List[PlanStage]:
             stages.append(CreateSnapshotRecordsStage(snapshots=plan.new_snapshots))
 
         snapshots_to_create = self._get_snapshots_to_create(plan, snapshots)
-        stages.append(
-            PhysicalLayerSchemaCreationStage(
-                snapshots=snapshots_to_create, deployability_index=deployability_index
-            )
-        )
-        if not plan.skip_backfill and not plan.empty_backfill:
-            # If the snapshot is selected for backfill and is not representative, then we assume
-            # this is a paused forward-only snapshot and we need to make sure a clone has been
-            # created for it in dev.
-            filtered_snapshots_to_create = []
-            for snapshot in snapshots_to_create:
-                if (
-                    plan.is_selected_for_backfill(snapshot.name)
-                    and snapshot not in snapshots_to_intervals
-                    and snapshot.is_materialized
-                    and not deployability_index.is_representative(snapshot)
-                ):
-                    filtered_snapshots_to_create.append(snapshot)
-            snapshots_to_create = filtered_snapshots_to_create
         if snapshots_to_create:
+            stages.append(
+                PhysicalLayerSchemaCreationStage(
+                    snapshots=snapshots_to_create, deployability_index=deployability_index
+                )
+            )
+        if not needs_backfill:
             stages.append(
                 self._get_physical_layer_update_stage(
                     plan,
@@ -340,6 +329,11 @@ def build(self, plan: EvaluatablePlan) -> t.List[PlanStage]:
             stages.append(
                 BackfillStage(
                     snapshot_to_intervals=missing_intervals_before_promote,
+                    selected_snapshot_ids={
+                        s_id
+                        for s_id in before_promote_snapshots
+                        if plan.is_selected_for_backfill(s_id.name)
+                    },
                     all_snapshots=snapshots_by_name,
                     deployability_index=deployability_index,
                 )
@@ -349,6 +343,7 @@ def build(self, plan: EvaluatablePlan) -> t.List[PlanStage]:
             stages.append(
                 BackfillStage(
                     snapshot_to_intervals={},
+                    selected_snapshot_ids=set(),
                     all_snapshots=snapshots_by_name,
                     deployability_index=deployability_index,
                 )
@@ -379,6 +374,11 @@ def build(self, plan: EvaluatablePlan) -> t.List[PlanStage]:
             stages.append(
                 BackfillStage(
                     snapshot_to_intervals=missing_intervals_after_promote,
+                    selected_snapshot_ids={
+                        s_id
+                        for s_id in after_promote_snapshots
+                        if plan.is_selected_for_backfill(s_id.name)
+                    },
                     all_snapshots=snapshots_by_name,
                     deployability_index=deployability_index,
                 )

sqlmesh/core/scheduler.py

@@ -332,8 +332,9 @@ def batch_intervals(
         merged_intervals: SnapshotToIntervals,
         deployability_index: t.Optional[DeployabilityIndex],
         environment_naming_info: EnvironmentNamingInfo,
+        dag: t.Optional[DAG[SnapshotId]] = None,
     ) -> t.Dict[Snapshot, Intervals]:
-        dag = snapshots_to_dag(merged_intervals)
+        dag = dag or snapshots_to_dag(merged_intervals)
 
         snapshot_intervals: t.Dict[SnapshotId, t.Tuple[Snapshot, t.List[Interval]]] = {
             snapshot.snapshot_id: (
@@ -413,6 +414,7 @@ def run_merged_intervals(
         start: t.Optional[TimeLike] = None,
         end: t.Optional[TimeLike] = None,
         allow_destructive_snapshots: t.Optional[t.Set[str]] = None,
+        selected_snapshot_ids: t.Optional[t.Set[SnapshotId]] = None,
         run_environment_statements: bool = False,
         audit_only: bool = False,
     ) -> t.Tuple[t.List[NodeExecutionFailedError[SchedulingUnit]], t.List[SchedulingUnit]]:
@@ -427,14 +429,21 @@ def run_merged_intervals(
             start: The start of the run.
             end: The end of the run.
             allow_destructive_snapshots: Snapshots for which destructive schema changes are allowed.
+            selected_snapshot_ids: The snapshots to include in the run DAG. If None, all snapshots with missing intervals will be included.
 
         Returns:
             A tuple of errors and skipped intervals.
         """
         execution_time = execution_time or now_timestamp()
 
+        selected_snapshots = [self.snapshots[sid] for sid in (selected_snapshot_ids or set())]
+        if not selected_snapshots:
+            selected_snapshots = list(merged_intervals)
+
+        snapshot_dag = snapshots_to_dag(selected_snapshots)
+
         batched_intervals = self.batch_intervals(
-            merged_intervals, deployability_index, environment_naming_info
+            merged_intervals, deployability_index, environment_naming_info, dag=snapshot_dag
         )
 
         self.console.start_evaluation_progress(
@@ -447,11 +456,13 @@ def run_merged_intervals(
         snapshots_to_create = {
             s.snapshot_id
             for s in self.snapshot_evaluator.get_snapshots_to_create(
-                merged_intervals.keys(), deployability_index
+                selected_snapshots, deployability_index
             )
         }
 
-        dag = self._dag(batched_intervals, snapshots_to_create=snapshots_to_create)
+        dag = self._dag(
+            batched_intervals, snapshot_dag=snapshot_dag, snapshots_to_create=snapshots_to_create
+        )
 
         if run_environment_statements:
             environment_statements = self.state_sync.get_environment_statements(
@@ -575,12 +586,14 @@ def evaluate_node(node: SchedulingUnit) -> None:
     def _dag(
         self,
         batches: SnapshotToIntervals,
+        snapshot_dag: t.Optional[DAG[SnapshotId]] = None,
         snapshots_to_create: t.Optional[t.Set[SnapshotId]] = None,
     ) -> DAG[SchedulingUnit]:
         """Builds a DAG of snapshot intervals to be evaluated.
 
         Args:
             batches: The batches of snapshots and intervals to evaluate.
+            snapshot_dag: The DAG of all snapshots.
             snapshots_to_create: The snapshots with missing physical tables.
 
         Returns:
@@ -591,20 +604,24 @@ def _dag(
             snapshot.name: intervals for snapshot, intervals in batches.items()
         }
         snapshots_to_create = snapshots_to_create or set()
+        original_snapshots_to_create = snapshots_to_create.copy()
 
+        snapshot_dag = snapshot_dag or snapshots_to_dag(batches)
         dag = DAG[SchedulingUnit]()
 
-        for snapshot, intervals in batches.items():
-            if not intervals:
-                continue
+        for snapshot_id in snapshot_dag:
+            snapshot = self.snapshots_by_name[snapshot_id.name]
+            intervals = intervals_per_snapshot.get(snapshot.name, [])
 
             upstream_dependencies: t.List[SchedulingUnit] = []
 
             for p_sid in snapshot.parents:
                 if p_sid in self.snapshots:
                     p_intervals = intervals_per_snapshot.get(p_sid.name, [])
 
-                    if len(p_intervals) > 1:
+                    if not p_intervals and p_sid in original_snapshots_to_create:
+                        upstream_dependencies.append(CreateNode(snapshot_name=p_sid.name))
+                    elif len(p_intervals) > 1:
                         upstream_dependencies.append(DummyNode(snapshot_name=p_sid.name))
                     else:
                         for i, interval in enumerate(p_intervals):
@@ -620,14 +637,16 @@ def _dag(
                 batch_concurrency = 1
 
             create_node: t.Optional[CreateNode] = None
-            if snapshot.snapshot_id in snapshots_to_create and (
+            if snapshot.snapshot_id in original_snapshots_to_create and (
                 snapshot.is_incremental_by_time_range
                 or ((not batch_concurrency or batch_concurrency > 1) and batch_size)
+                or not intervals
             ):
                 # Add a separate node for table creation in case when there multiple concurrent
-                # evaluation nodes.
+                # evaluation nodes or when there are no intervals to evaluate.
                 create_node = CreateNode(snapshot_name=snapshot.name)
                 dag.add(create_node, upstream_dependencies)
+                snapshots_to_create.remove(snapshot.snapshot_id)
 
             for i, interval in enumerate(intervals):
                 node = EvaluateNode(snapshot_name=snapshot.name, interval=interval, batch_index=i)

sqlmesh/core/snapshot/definition.py

@@ -1043,8 +1043,15 @@ def categorize_as(self, category: SnapshotChangeCategory, forward_only: bool = F
             # If the model has a pinned version then use that.
             self.version = self.model.physical_version
         elif is_no_rebuild and self.previous_version:
+            self.version = self.previous_version.data_version.version
+        elif self.is_model and self.model.forward_only and not self.previous_version:
+            # If this is a new model then use a deterministic version, independent of the fingerprint.
+            self.version = hash_data([self.name, *self.model.kind.data_hash_values])
+        else:
+            self.version = self.fingerprint.to_version()
+
+        if is_no_rebuild and self.previous_version:
             previous_version = self.previous_version
-            self.version = previous_version.data_version.version
             self.physical_schema_ = previous_version.physical_schema
             self.table_naming_convention = previous_version.table_naming_convention
             if self.is_materialized and (category.is_indirect_non_breaking or category.is_metadata):
@@ -1054,11 +1061,6 @@ def categorize_as(self, category: SnapshotChangeCategory, forward_only: bool = F
                     or previous_version.fingerprint.to_version()
                 )
                 self.dev_table_suffix = previous_version.data_version.dev_table_suffix
-        elif self.is_model and self.model.forward_only and not self.previous_version:
-            # If this is a new model then use a deterministic version, independent of the fingerprint.
-            self.version = hash_data([self.name, *self.model.kind.data_hash_values])
-        else:
-            self.version = self.fingerprint.to_version()
 
         self.change_category = category
         self.forward_only = forward_only
@@ -1603,9 +1605,7 @@ def create(
                 )
             else:
                 children_deployable = False
-                if not snapshots[node].is_paused or (
-                    snapshot.is_indirect_non_breaking and snapshot.intervals
-                ):
+                if not snapshots[node].is_paused:
                     representative_shared_version_ids.add(node)
 
             deployability_mapping[node] = this_deployable

sqlmesh/core/snapshot/evaluator.py

@@ -695,9 +695,11 @@ def _evaluate_snapshot(
         target_table_name = snapshot.table_name(is_deployable=is_snapshot_deployable)
         # https://github.com/TobikoData/sqlmesh/issues/2609
         # If there are no existing intervals yet; only consider this a first insert for the first snapshot in the batch
-        is_first_insert = not _intervals(snapshot, deployability_index) and batch_index == 0
         if target_table_exists is None:
             target_table_exists = adapter.table_exists(target_table_name)
+        is_first_insert = (
+            not _intervals(snapshot, deployability_index) or not target_table_exists
+        ) and batch_index == 0
 
         common_render_kwargs = dict(
             start=start,
@@ -749,7 +751,6 @@ def _evaluate_snapshot(
                         allow_destructive_snapshots=allow_destructive_snapshots,
                     )
                 else:
-                    is_first_insert = True
                     if model.annotated or model.is_seed or model.kind.is_scd_type_2:
                         self._execute_create(
                             snapshot=snapshot,

tests/core/engine_adapter/integration/test_integration.py

@@ -1914,6 +1914,11 @@ def test_sushi(ctx: TestContext, tmp_path_factory: pytest.TempPathFactory):
         ],
         personal_paths=[pathlib.Path("~/.sqlmesh/config.yaml").expanduser()],
     )
+    config.before_all = [
+        f"CREATE SCHEMA IF NOT EXISTS {raw_test_schema}",
+        f"DROP VIEW IF EXISTS {raw_test_schema}.demographics",
+        f"CREATE VIEW {raw_test_schema}.demographics AS (SELECT 1 AS customer_id, '00000' AS zip)",
+    ]
 
     # To enable parallelism in integration tests
     config.gateways = {ctx.gateway: config.gateways[ctx.gateway]}
@@ -2132,6 +2137,8 @@ def validate_comments(
             }
 
             for model_name, comment in comments.items():
+                if not model_name in layer_models:
+                    continue
                 layer_table_name = layer_models[model_name]["table_name"]
                 table_kind = "VIEW" if layer_models[model_name]["is_view"] else "BASE TABLE"
 

tests/core/test_integration.py

@@ -5890,7 +5890,7 @@ def get_default_catalog_and_non_tables(
     ) = get_default_catalog_and_non_tables(metadata, context.default_catalog)
     assert len(prod_views) == 16
     assert len(dev_views) == 16
-    assert len(user_default_tables) == 15
+    assert len(user_default_tables) == 16
     assert len(non_default_tables) == 0
     assert state_metadata.schemas == ["sqlmesh"]
     assert {x.sql() for x in state_metadata.qualified_tables}.issuperset(
@@ -5910,7 +5910,7 @@ def get_default_catalog_and_non_tables(
     ) = get_default_catalog_and_non_tables(metadata, context.default_catalog)
     assert len(prod_views) == 16
     assert len(dev_views) == 32
-    assert len(user_default_tables) == 15
+    assert len(user_default_tables) == 16
     assert len(non_default_tables) == 0
     assert state_metadata.schemas == ["sqlmesh"]
     assert {x.sql() for x in state_metadata.qualified_tables}.issuperset(
@@ -5931,7 +5931,7 @@ def get_default_catalog_and_non_tables(
     ) = get_default_catalog_and_non_tables(metadata, context.default_catalog)
     assert len(prod_views) == 16
     assert len(dev_views) == 16
-    assert len(user_default_tables) == 15
+    assert len(user_default_tables) == 16
     assert len(non_default_tables) == 0
     assert state_metadata.schemas == ["sqlmesh"]
     assert {x.sql() for x in state_metadata.qualified_tables}.issuperset(
@@ -6902,17 +6902,7 @@ def plan_with_output(ctx: Context, environment: str):
     assert "New environment `dev` will be created from `prod`" in output.stdout
     assert "Differences from the `prod` environment" in output.stdout
 
-    assert (
-        """MODEL (                                                                        
-   name test.a,                                                                 
-+  owner test,                                                                  
-   kind FULL                                                                    
- )                                                                              
- SELECT                                                                         
--  5 AS col                                                                     
-+  10 AS col"""
-        in output.stdout
-    )
+    assert "Directly Modified: test__dev.a" in output.stdout
 
     # Case 6: Ensure that target environment and create_from environment are not the same
     output = plan_with_output(ctx, "prod")

tests/utils/test_helpers.py

@@ -82,7 +82,9 @@ def use_terminal_console(func):
     def test_wrapper(*args, **kwargs):
         orig_console = get_console()
         try:
-            set_console(TerminalConsole())
+            new_console = TerminalConsole()
+            new_console.console.no_color = True
+            set_console(new_console)
             func(*args, **kwargs)
         finally:
             set_console(orig_console)