Fix!: Support snapshots of added models with forward-only parents (#1372)

izeigerman · web-flow · commit ff753e5b3ca4 · 2023-09-01T10:33:49.000-07:00
diff --git a/sqlmesh/core/plan/definition.py b/sqlmesh/core/plan/definition.py
@@ -545,14 +545,16 @@ def _categorize_snapshots(self) -> None:
         returns a list of added and directly modified snapshots as well as the mapping of
         indirectly modified snapshots.
         """
-        for model_name, snapshot in self._snapshot_mapping.items():
-            if snapshot.change_category:
+
+        # Iterating in DAG order since a category for a snapshot may depend on the categories
+        # assigned to its upstream dependencies.
+        for model_name in self._dag:
+            snapshot = self._snapshot_mapping.get(model_name)
+            if not snapshot or snapshot.change_category:
                 continue
-            upstream_model_names = self._dag.upstream(model_name)
 
             if model_name in self.context_diff.modified_snapshots:
                 is_directly_modified = self.context_diff.directly_modified(model_name)
-
                 if self.is_new_snapshot(snapshot):
                     if self.forward_only:
                         # In case of the forward only plan any modifications result in reuse of the
@@ -597,13 +599,17 @@ def _categorize_snapshots(self) -> None:
                     and not any(
                         self.context_diff.directly_modified(upstream)
                         and not self._snapshot_mapping[upstream].version
-                        for upstream in upstream_model_names
+                        for upstream in self._dag.upstream(model_name)
                     )
                 ):
                     snapshot.categorize_as(SnapshotChangeCategory.INDIRECT_BREAKING)
 
             elif model_name in self.context_diff.added and self.is_new_snapshot(snapshot):
-                snapshot.categorize_as(SnapshotChangeCategory.BREAKING)
+                snapshot.categorize_as(
+                    SnapshotChangeCategory.FORWARD_ONLY
+                    if self._is_forward_only_model(model_name)
+                    else SnapshotChangeCategory.BREAKING
+                )
 
     def _ensure_valid_date_range(
         self, start: t.Optional[TimeLike], end: t.Optional[TimeLike]
@@ -636,12 +642,6 @@ def _ensure_no_forward_only_revert(self) -> None:
                     "To proceed with the change, restamp this model's definition to produce a new version."
                 )
 
-    def _ensure_no_forward_only_new_models(self) -> None:
-        if self.forward_only and self.context_diff.added_materialized_models:
-            raise PlanError(
-                "New models that require materialization can't be added as part of the forward-only plan."
-            )
-
     def _ensure_no_broken_references(self) -> None:
         for snapshot in self.context_diff.snapshots.values():
             broken_references = self.context_diff.removed & snapshot.model.depends_on
@@ -683,7 +683,6 @@ def _refresh_dag_and_ignored_snapshots(self) -> None:
         self._ensure_new_env_with_changes()
         self._ensure_valid_date_range(self._start, self._end)
         self._ensure_no_forward_only_revert()
-        self._ensure_no_forward_only_new_models()
         self._ensure_no_broken_references()
 
         (
@@ -732,16 +731,29 @@ def _build_snapshots_and_dag(
         )
 
     def _is_forward_only_model(self, model_name: str) -> bool:
-        snapshot = self._snapshot_mapping.get(model_name)
-        if snapshot and snapshot.model.forward_only:
+        def _is_forward_only_expected(snapshot: Snapshot) -> bool:
+            # Returns True if the snapshot is not categorized yet but is expected
+            # to be categorized as forward-only. Checking the previous versions to make
+            # sure that the snapshot doesn't represent a newly added model.
+            return (
+                snapshot.model.forward_only
+                and not snapshot.change_category
+                and bool(snapshot.previous_versions)
+            )
+
+        snapshot = self._snapshot_mapping[model_name]
+        if _is_forward_only_expected(snapshot):
             return True
 
         for upstream in self._dag.upstream(model_name):
             upstream_snapshot = self._snapshot_mapping.get(upstream)
             if (
                 upstream_snapshot
                 and upstream_snapshot.is_paused
-                and (upstream_snapshot.is_forward_only or upstream_snapshot.model.forward_only)
+                and (
+                    upstream_snapshot.is_forward_only
+                    or _is_forward_only_expected(upstream_snapshot)
+                )
             ):
                 return True
         return False
diff --git a/sqlmesh/core/snapshot/evaluator.py b/sqlmesh/core/snapshot/evaluator.py
@@ -407,6 +407,19 @@ def _create_snapshot(
         evaluation_strategy = _evaluation_strategy(snapshot, self.adapter)
 
         with self.adapter.transaction(TransactionType.DDL), self.adapter.session():
+            if is_dev and not snapshot.previous_versions:
+                # This is a FORWARD_ONLY snapshot which represents an added model. This means
+                # that the physical table associated with it doesn't exist yet.
+                logger.info(
+                    f"Detected a forward-only snapshot without previous versions: {snapshot.snapshot_id}"
+                )
+                non_dev_render_kwargs: t.Dict[str, t.Any] = {**render_kwargs, "is_dev": False}
+                self.adapter.execute(snapshot.model.render_pre_statements(**non_dev_render_kwargs))
+                evaluation_strategy.create(
+                    snapshot.model, snapshot.table_name(), **non_dev_render_kwargs
+                )
+                self.adapter.execute(snapshot.model.render_post_statements(**non_dev_render_kwargs))
+
             self.adapter.execute(snapshot.model.render_pre_statements(**render_kwargs))
 
             if is_dev and snapshot.is_materialized and self.adapter.SUPPORTS_CLONING:
@@ -480,9 +493,13 @@ def _demote_snapshot(
 
     def _cleanup_snapshot(self, snapshot: SnapshotInfoLike) -> None:
         snapshot = snapshot.table_info
-        table_names = [snapshot.table_name()]
-        if snapshot.version != snapshot.fingerprint.to_version():
-            table_names.append(snapshot.table_name(is_dev=True))
+
+        table_name = snapshot.table_name()
+        dev_table_name = snapshot.table_name(is_dev=True)
+
+        table_names = [table_name]
+        if table_name != dev_table_name:
+            table_names.append(dev_table_name)
 
         evaluation_strategy = _evaluation_strategy(snapshot, self.adapter)
 
diff --git a/sqlmesh/migrations/v0023_fix_added_models_with_forward_only_parents.py b/sqlmesh/migrations/v0023_fix_added_models_with_forward_only_parents.py
@@ -0,0 +1,60 @@
+"""Fix snapshots of added models with forward only parents."""
+import json
+import typing as t
+
+from sqlglot import exp
+
+from sqlmesh.utils.dag import DAG
+
+
+def migrate(state_sync: t.Any) -> None:
+    engine_adapter = state_sync.engine_adapter
+    schema = state_sync.schema
+    snapshots_table = "_snapshots"
+    environments_table = "_environments"
+    if schema:
+        snapshots_table = f"{schema}.{snapshots_table}"
+        environments_table = f"{schema}.{environments_table}"
+
+    dag: DAG[t.Tuple[str, str]] = DAG()
+    snapshot_mapping: t.Dict[t.Tuple[str, str], t.Dict[str, t.Any]] = {}
+
+    for identifier, snapshot in engine_adapter.fetchall(
+        exp.select("identifier", "snapshot").from_(snapshots_table),
+        quote_identifiers=True,
+    ):
+        parsed_snapshot = json.loads(snapshot)
+
+        snapshot_id = (parsed_snapshot["name"], identifier)
+        snapshot_mapping[snapshot_id] = parsed_snapshot
+
+        parent_ids = [
+            (parent["name"], parent["identifier"]) for parent in parsed_snapshot["parents"]
+        ]
+        dag.add(snapshot_id, parent_ids)
+
+    snapshots_to_delete = set()
+
+    for snapshot_id in dag:
+        parsed_snapshot = snapshot_mapping[snapshot_id]
+        is_breaking = parsed_snapshot.get("change_category") == 1
+        has_previous_versions = bool(parsed_snapshot.get("previous_versions", []))
+
+        has_paused_forward_only_parent = False
+        if is_breaking and not has_previous_versions:
+            for upstream_id in dag.upstream(snapshot_id):
+                upstream_snapshot = snapshot_mapping[upstream_id]
+                upstream_change_category = upstream_snapshot.get("change_category")
+                is_forward_only_upstream = upstream_change_category == 3
+                if is_forward_only_upstream and not upstream_snapshot.get("unpaused_ts"):
+                    has_paused_forward_only_parent = True
+                    break
+
+        if has_paused_forward_only_parent:
+            snapshots_to_delete.add(snapshot_id)
+
+    if snapshots_to_delete:
+        where = t.cast(exp.Tuple, exp.convert((exp.column("name"), exp.column("identifier")))).isin(
+            *snapshots_to_delete
+        )
+        engine_adapter.delete_from(snapshots_table, where)
diff --git a/tests/core/test_plan.py b/tests/core/test_plan.py
@@ -90,26 +90,27 @@ def test_forward_only_dev(make_snapshot, mocker: MockerFixture):
     now_ds_mock.call_count == 2
 
 
-def test_forward_only_plan_new_models_not_allowed(make_snapshot, mocker: MockerFixture):
-    snapshot_a = make_snapshot(SqlModel(name="a", query=parse_one("select 1, ds")))
-    snapshot_a.categorize_as(SnapshotChangeCategory.BREAKING)
+def test_forward_only_plan_added_models(make_snapshot, mocker: MockerFixture):
+    snapshot_a = make_snapshot(SqlModel(name="a", query=parse_one("select 1 as a, ds")))
+
+    snapshot_b = make_snapshot(
+        SqlModel(name="b", query=parse_one("select a, ds from a")), nodes={"a": snapshot_a.node}
+    )
 
     context_diff_mock = mocker.Mock()
-    context_diff_mock.snapshots = {"a": snapshot_a}
-    context_diff_mock.added = {"a"}
+    context_diff_mock.snapshots = {"a": snapshot_a, "b": snapshot_b}
+    context_diff_mock.added = {"b"}
     context_diff_mock.removed = set()
-    context_diff_mock.modified_snapshots = {}
-    context_diff_mock.new_snapshots = {}
-    context_diff_mock.added_materialized_models = {"a"}
-
-    with pytest.raises(
-        PlanError,
-        match="New models that require materialization can't be added as part of the forward-only plan.",
-    ):
-        Plan(context_diff_mock, forward_only=True)
+    context_diff_mock.modified_snapshots = {"a": (snapshot_a, snapshot_a)}
+    context_diff_mock.new_snapshots = {
+        snapshot_a.snapshot_id: snapshot_a,
+        snapshot_b.snapshot_id: snapshot_b,
+    }
+    context_diff_mock.added_materialized_models = {"b"}
 
-    context_diff_mock.added_materialized_models = set()
     Plan(context_diff_mock, forward_only=True)
+    assert snapshot_a.change_category == SnapshotChangeCategory.FORWARD_ONLY
+    assert snapshot_b.change_category == SnapshotChangeCategory.FORWARD_ONLY
 
 
 def test_paused_forward_only_parent(make_snapshot, mocker: MockerFixture):
@@ -577,6 +578,7 @@ def test_forward_only_models(make_snapshot, mocker: MockerFixture):
             kind=IncrementalByTimeRangeKind(time_column="ds", forward_only=True),
         )
     )
+    updated_snapshot.previous_versions = snapshot.all_versions
 
     context_diff_mock = mocker.Mock()
     context_diff_mock.snapshots = {"a": updated_snapshot}
@@ -607,6 +609,7 @@ def test_indirectly_modified_forward_only_model(make_snapshot, mocker: MockerFix
     snapshot_a = make_snapshot(SqlModel(name="a", query=parse_one("select 1 as a, ds")))
     snapshot_a.categorize_as(SnapshotChangeCategory.BREAKING)
     updated_snapshot_a = make_snapshot(SqlModel(name="a", query=parse_one("select 2 as a, ds")))
+    updated_snapshot_a.previous_versions = snapshot_a.all_versions
 
     snapshot_b = make_snapshot(
         SqlModel(
@@ -618,12 +621,14 @@ def test_indirectly_modified_forward_only_model(make_snapshot, mocker: MockerFix
     )
     snapshot_b.categorize_as(SnapshotChangeCategory.FORWARD_ONLY)
     updated_snapshot_b = make_snapshot(snapshot_b.model, nodes={"a": updated_snapshot_a.model})
+    updated_snapshot_b.previous_versions = snapshot_b.all_versions
 
     snapshot_c = make_snapshot(
         SqlModel(name="c", query=parse_one("select a, ds from b")), nodes={"b": snapshot_b.model}
     )
     snapshot_c.categorize_as(SnapshotChangeCategory.BREAKING)
     updated_snapshot_c = make_snapshot(snapshot_c.model, nodes={"b": updated_snapshot_b.model})
+    updated_snapshot_c.previous_versions = snapshot_c.all_versions
 
     context_diff_mock = mocker.Mock()
     context_diff_mock.snapshots = {
@@ -660,6 +665,65 @@ def test_indirectly_modified_forward_only_model(make_snapshot, mocker: MockerFix
     assert updated_snapshot_c.change_category == SnapshotChangeCategory.FORWARD_ONLY
 
 
+def test_added_model_with_forward_only_parent(make_snapshot, mocker: MockerFixture):
+    snapshot_a = make_snapshot(SqlModel(name="a", query=parse_one("select 1 as a, ds")))
+    snapshot_a.categorize_as(SnapshotChangeCategory.FORWARD_ONLY)
+
+    snapshot_b = make_snapshot(SqlModel(name="b", query=parse_one("select a, ds from a")))
+
+    context_diff_mock = mocker.Mock()
+    context_diff_mock.snapshots = {
+        "a": snapshot_a,
+        "b": snapshot_b,
+    }
+    context_diff_mock.removed = set()
+    context_diff_mock.added = {"b"}
+    context_diff_mock.added_materialized_models = set()
+    context_diff_mock.modified_snapshots = {}
+    context_diff_mock.new_snapshots = {
+        snapshot_b.snapshot_id: snapshot_b,
+    }
+    context_diff_mock.has_snapshot_changes = True
+    context_diff_mock.environment = "test_dev"
+    context_diff_mock.previous_plan_id = "previous_plan_id"
+
+    Plan(context_diff_mock)
+    assert snapshot_b.change_category == SnapshotChangeCategory.FORWARD_ONLY
+
+
+def test_added_forward_only_model(make_snapshot, mocker: MockerFixture):
+    snapshot_a = make_snapshot(
+        SqlModel(
+            name="a",
+            query=parse_one("select 1 as a, ds"),
+            kind=IncrementalByTimeRangeKind(time_column="ds", forward_only=True),
+        )
+    )
+
+    snapshot_b = make_snapshot(SqlModel(name="b", query=parse_one("select a, ds from a")))
+
+    context_diff_mock = mocker.Mock()
+    context_diff_mock.snapshots = {
+        "a": snapshot_a,
+        "b": snapshot_b,
+    }
+    context_diff_mock.removed = set()
+    context_diff_mock.added = {"a", "b"}
+    context_diff_mock.added_materialized_models = set()
+    context_diff_mock.modified_snapshots = {}
+    context_diff_mock.new_snapshots = {
+        snapshot_a.snapshot_id: snapshot_a,
+        snapshot_b.snapshot_id: snapshot_b,
+    }
+    context_diff_mock.has_snapshot_changes = True
+    context_diff_mock.environment = "test_dev"
+    context_diff_mock.previous_plan_id = "previous_plan_id"
+
+    Plan(context_diff_mock)
+    assert snapshot_a.change_category == SnapshotChangeCategory.BREAKING
+    assert snapshot_b.change_category == SnapshotChangeCategory.BREAKING
+
+
 def test_disable_restatement(make_snapshot, mocker: MockerFixture):
     snapshot = make_snapshot(
         SqlModel(
diff --git a/tests/core/test_snapshot_evaluator.py b/tests/core/test_snapshot_evaluator.py
