fix model kind change edge case

Author: izeigerman

sqlmesh/core/engine_adapter/base.py

@@ -371,9 +371,10 @@ def replace_query(
         """
         target_table = exp.to_table(table_name)
 
-        target_data_object = self._get_data_object(target_table)
+        target_data_object = self.get_data_object(target_table)
         table_exists = target_data_object is not None
-        self._drop_data_object_on_type_mismatch(target_data_object, DataObjectType.TABLE)
+        if self.drop_data_object_on_type_mismatch(target_data_object, DataObjectType.TABLE):
+            table_exists = False
 
         source_queries, columns_to_types = self._get_source_queries_and_columns_to_types(
             query_or_df, columns_to_types, target_table=target_table
@@ -1147,8 +1148,8 @@ def create_view(
             create_kwargs["properties"] = properties
 
         if replace:
-            self._drop_data_object_on_type_mismatch(
-                self._get_data_object(view_name),
+            self.drop_data_object_on_type_mismatch(
+                self.get_data_object(view_name),
                 DataObjectType.VIEW if not materialized else DataObjectType.MATERIALIZED_VIEW,
             )
 
@@ -2056,6 +2057,15 @@ def rename_table(
                 )
         self._rename_table(old_table_name, new_table_name)
 
+    def get_data_object(self, target_name: TableName) -> t.Optional[DataObject]:
+        target_table = exp.to_table(target_name)
+        existing_data_objects = self.get_data_objects(
+            schema_(target_table.db, target_table.catalog), {target_table.name}
+        )
+        if existing_data_objects:
+            return existing_data_objects[0]
+        return None
+
     def get_data_objects(
         self, schema_name: SchemaName, object_names: t.Optional[t.Set[str]] = None
     ) -> t.List[DataObject]:
@@ -2517,26 +2527,20 @@ def _truncate_table(self, table_name: TableName) -> None:
         table = exp.to_table(table_name)
         self.execute(f"TRUNCATE TABLE {table.sql(dialect=self.dialect, identify=True)}")
 
-    def _get_data_object(self, target_name: TableName) -> t.Optional[DataObject]:
-        target_table = exp.to_table(target_name)
-        existing_data_objects = self.get_data_objects(
-            schema_(target_table.db, target_table.catalog), {target_table.name}
-        )
-        if existing_data_objects:
-            return existing_data_objects[0]
-        return None
-
-    def _drop_data_object_on_type_mismatch(
+    def drop_data_object_on_type_mismatch(
         self, data_object: t.Optional[DataObject], expected_type: DataObjectType
-    ) -> None:
+    ) -> bool:
         """Drops a data object if it exists and is not of the expected type.
 
         Args:
             data_object: The data object to check.
             expected_type: The expected type of the data object.
+
+        Returns:
+            True if the data object was dropped, False otherwise.
         """
         if data_object is None or data_object.type == expected_type:
-            return
+            return False
 
         logger.warning(
             "Target data object '%s' is a %s and not a %s, dropping it",
@@ -2545,6 +2549,7 @@ def _drop_data_object_on_type_mismatch(
             expected_type.value,
         )
         self.drop_data_object(data_object)
+        return True
 
     def _replace_by_key(
         self,

sqlmesh/core/engine_adapter/redshift.py

@@ -262,9 +262,10 @@ def replace_query(
         """
         import pandas as pd
 
-        target_data_object = self._get_data_object(table_name)
+        target_data_object = self.get_data_object(table_name)
         table_exists = target_data_object is not None
-        self._drop_data_object_on_type_mismatch(target_data_object, DataObjectType.TABLE)
+        if self.drop_data_object_on_type_mismatch(target_data_object, DataObjectType.TABLE):
+            table_exists = False
 
         if not isinstance(query_or_df, pd.DataFrame) or not table_exists:
             return super().replace_query(

sqlmesh/core/snapshot/evaluator.py

@@ -38,7 +38,7 @@
 from sqlmesh.core.audit import Audit, StandaloneAudit
 from sqlmesh.core.dialect import schema_
 from sqlmesh.core.engine_adapter import EngineAdapter
-from sqlmesh.core.engine_adapter.shared import InsertOverwriteStrategy
+from sqlmesh.core.engine_adapter.shared import InsertOverwriteStrategy, DataObjectType
 from sqlmesh.core.macros import RuntimeStage
 from sqlmesh.core.model import (
     AuditResult,
@@ -934,7 +934,14 @@ def _migrate_snapshot(
             adapter.transaction(),
             adapter.session(snapshot.model.render_session_properties(**render_kwargs)),
         ):
-            if adapter.table_exists(target_table_name):
+            target_data_object = adapter.get_data_object(target_table_name)
+            table_exists = target_data_object is not None
+            if adapter.drop_data_object_on_type_mismatch(
+                target_data_object, _snapshot_to_data_object_type(snapshot)
+            ):
+                table_exists = False
+
+            if table_exists:
                 evaluation_strategy = _evaluation_strategy(snapshot, adapter)
                 tmp_table_name = snapshot.table_name(is_deployable=False)
                 logger.info(
@@ -2307,3 +2314,15 @@ def _check_table_db_is_physical_schema(table_name: str, physical_schema: str) ->
         raise SQLMeshError(
             f"Table '{table_name}' is not a part of the physical schema '{physical_schema}' and so can't be dropped."
         )
+
+
+def _snapshot_to_data_object_type(snapshot: Snapshot) -> DataObjectType:
+    if snapshot.is_managed:
+        return DataObjectType.MANAGED_TABLE
+    if snapshot.is_materialized_view:
+        return DataObjectType.MATERIALIZED_VIEW
+    if snapshot.is_view:
+        return DataObjectType.VIEW
+    if snapshot.is_materialized:
+        return DataObjectType.TABLE
+    return DataObjectType.UNKNOWN

tests/core/test_integration.py

@@ -2603,6 +2603,48 @@ def test_virtual_environment_mode_dev_only_model_kind_change(init_and_plan_conte
     assert data_objects[0].type == "table"
 
 
+@time_machine.travel("2023-01-08 15:00:00 UTC")
+def test_virtual_environment_mode_dev_only_model_kind_change_with_follow_up_changes_in_dev(
+    init_and_plan_context: t.Callable,
+):
+    context, plan = init_and_plan_context(
+        "examples/sushi", config="test_config_virtual_environment_mode_dev_only"
+    )
+    context.apply(plan)
+
+    # Make sure the initial state is a view
+    data_objects = context.engine_adapter.get_data_objects("sushi", {"top_waiters"})
+    assert len(data_objects) == 1
+    assert data_objects[0].type == "view"
+
+    # Change to incremental unmanaged kind
+    model = context.get_model("sushi.top_waiters")
+    model = model.copy(update={"kind": IncrementalUnmanagedKind()})
+    context.upsert_model(model)
+    dev_plan = context.plan_builder("dev", skip_tests=True).build()
+    assert dev_plan.missing_intervals
+    assert dev_plan.requires_backfill
+    context.apply(dev_plan)
+
+    # Make a follow-up forward-only change
+    model = add_projection_to_model(t.cast(SqlModel, model))
+    context.upsert_model(model)
+    dev_plan = context.plan_builder("dev", skip_tests=True, forward_only=True).build()
+    context.apply(dev_plan)
+
+    # Deploy to prod
+    prod_plan = context.plan_builder("prod", skip_tests=True).build()
+    assert prod_plan.requires_backfill
+    assert prod_plan.missing_intervals
+    assert not prod_plan.context_diff.snapshots[
+        context.get_snapshot(model.name).snapshot_id
+    ].intervals
+    context.apply(prod_plan)
+    data_objects = context.engine_adapter.get_data_objects("sushi", {"top_waiters"})
+    assert len(data_objects) == 1
+    assert data_objects[0].type == "table"
+
+
 @time_machine.travel("2023-01-08 15:00:00 UTC")
 def test_virtual_environment_mode_dev_only_model_kind_change_manual_categorization(
     init_and_plan_context: t.Callable,