fix!: allow restatement of external models (#3901)

Author: tobymao

sqlmesh/core/model/kind.py

@@ -813,7 +813,6 @@ def supports_python_models(self) -> bool:
 
 class EmbeddedKind(_ModelKind):
     name: t.Literal[ModelKindName.EMBEDDED] = ModelKindName.EMBEDDED
-    disable_restatement: t.Literal[True] = True
 
     @property
     def supports_python_models(self) -> bool:
@@ -822,7 +821,6 @@ def supports_python_models(self) -> bool:
 
 class ExternalKind(_ModelKind):
     name: t.Literal[ModelKindName.EXTERNAL] = ModelKindName.EXTERNAL
-    disable_restatement: t.Literal[True] = True
 
 
 class CustomKind(_ModelKind):

sqlmesh/core/plan/builder.py

@@ -349,7 +349,7 @@ def _build_restatements(
                         "If you want to restate this model, change the model's `disable_restatement` setting to `false`."
                     )
                     continue
-                elif snapshot.is_symbolic or snapshot.is_seed:
+                elif snapshot.is_seed:
                     logger.info("Skipping restatement for model '%s'", snapshot.name)
                     continue
 

sqlmesh/migrations/v0073_remove_symbolic_disable_restatement.py

@@ -0,0 +1,72 @@
+"""Remove disable restatement from external and embedded models."""
+
+import json
+import pandas as pd
+
+from sqlglot import exp
+from sqlmesh.utils.migration import index_text_type, blob_text_type
+
+
+def migrate(state_sync, **kwargs):  # type: ignore
+    engine_adapter = state_sync.engine_adapter
+    schema = state_sync.schema
+    snapshots_table = "_snapshots"
+    if schema:
+        snapshots_table = f"{schema}.{snapshots_table}"
+
+    index_type = index_text_type(engine_adapter.dialect)
+    blob_type = blob_text_type(engine_adapter.dialect)
+    snapshots_columns_to_types = {
+        "name": exp.DataType.build(index_type),
+        "identifier": exp.DataType.build(index_type),
+        "version": exp.DataType.build(index_type),
+        "snapshot": exp.DataType.build(blob_type),
+        "kind_name": exp.DataType.build(index_type),
+        "updated_ts": exp.DataType.build("bigint"),
+        "unpaused_ts": exp.DataType.build("bigint"),
+        "ttl_ms": exp.DataType.build("bigint"),
+        "unrestorable": exp.DataType.build("boolean"),
+    }
+
+    new_snapshots = []
+    for (
+        name,
+        identifier,
+        version,
+        snapshot,
+        kind_name,
+        updated_ts,
+        unpaused_ts,
+        ttl_ms,
+        unrestorable,
+    ) in engine_adapter.fetchall(
+        exp.select(*snapshots_columns_to_types).from_(snapshots_table),
+        quote_identifiers=True,
+    ):
+        parsed_snapshot = json.loads(snapshot)
+        kind = parsed_snapshot["node"].get("kind")
+
+        if kind and kind_name in ("EMBEDDED", "EXTERNAL"):
+            kind.pop("disable_restatement", None)
+
+        new_snapshots.append(
+            {
+                "name": name,
+                "identifier": identifier,
+                "version": version,
+                "snapshot": json.dumps(parsed_snapshot),
+                "kind_name": kind_name,
+                "updated_ts": updated_ts,
+                "unpaused_ts": unpaused_ts,
+                "ttl_ms": ttl_ms,
+                "unrestorable": unrestorable,
+            }
+        )
+
+    if new_snapshots:
+        engine_adapter.delete_from(snapshots_table, "TRUE")
+        engine_adapter.insert_append(
+            snapshots_table,
+            pd.DataFrame(new_snapshots),
+            columns_to_types=snapshots_columns_to_types,
+        )

tests/core/test_plan.py

@@ -789,6 +789,21 @@ def test_restate_models(sushi_context_pre_scheduling: Context):
     assert not plan.restatements
     assert plan.models_to_backfill is None
 
+    plan = sushi_context_pre_scheduling.plan(restate_models=["raw.demographics"], no_prompts=True)
+    assert not plan.has_changes
+    assert plan.restatements
+    assert plan.models_to_backfill == {
+        '"memory"."raw"."demographics"',
+        '"memory"."sushi"."active_customers"',
+        '"memory"."sushi"."customers"',
+        '"memory"."sushi"."marketing"',
+        '"memory"."sushi"."orders"',
+        '"memory"."sushi"."raw_marketing"',
+        '"memory"."sushi"."waiter_as_customer_by_day"',
+        '"memory"."sushi"."waiter_names"',
+        '"memory"."sushi"."waiters"',
+    }
+
 
 @pytest.mark.slow
 @time_machine.travel(now(minute_floor=False), tick=False)
@@ -884,7 +899,7 @@ def test_restate_symbolic_model(make_snapshot, mocker: MockerFixture):
     plan = PlanBuilder(
         context_diff, DuckDBEngineAdapter.SCHEMA_DIFFER, restate_models=[snapshot_a.name]
     ).build()
-    assert not plan.restatements
+    assert plan.restatements
 
 
 def test_restate_seed_model(make_snapshot, mocker: MockerFixture):