drop data objects of different types

Author: izeigerman

sqlmesh/core/engine_adapter/base.py

@@ -32,6 +32,7 @@
     CommentCreationTable,
     CommentCreationView,
     DataObject,
+    DataObjectType,
     EngineRunMode,
     InsertOverwriteStrategy,
     SourceQuery,
@@ -369,6 +370,9 @@ def replace_query(
             kwargs: Optional create table properties.
         """
         target_table = exp.to_table(table_name)
+
+        table_exists = self._drop_data_object_on_type_mismatch(target_table, DataObjectType.TABLE)
+
         source_queries, columns_to_types = self._get_source_queries_and_columns_to_types(
             query_or_df, columns_to_types, target_table=target_table
         )
@@ -390,7 +394,7 @@ def replace_query(
             )
         # All engines support `CREATE TABLE AS` so we use that if the table doesn't already exist and we
         # use `CREATE OR REPLACE TABLE AS` if the engine supports it
-        if self.SUPPORTS_REPLACE_TABLE or not self.table_exists(target_table):
+        if self.SUPPORTS_REPLACE_TABLE or not table_exists:
             return self._create_table_from_source_queries(
                 target_table,
                 source_queries,
@@ -930,6 +934,28 @@ def clone_table(
             )
         )
 
+    def drop_data_object(self, data_object: DataObject, ignore_if_not_exists: bool = True) -> None:
+        """Drops a data object of arbitrary type.
+
+        Args:
+            data_object: The data object to drop.
+            ignore_if_not_exists: If True, no error will be raised if the data object does not exist.
+        """
+        if data_object.type.is_view:
+            self.drop_view(data_object.to_table(), ignore_if_not_exists=ignore_if_not_exists)
+        elif data_object.type.is_materialized_view:
+            self.drop_view(
+                data_object.to_table(), ignore_if_not_exists=ignore_if_not_exists, materialized=True
+            )
+        elif data_object.type.is_table:
+            self.drop_table(data_object.to_table(), exists=ignore_if_not_exists)
+        elif data_object.type.is_managed_table:
+            self.drop_managed_table(data_object.to_table(), exists=ignore_if_not_exists)
+        else:
+            raise SQLMeshError(
+                f"Can't drop data object '{data_object.to_table().sql(dialect=self.dialect)}' of type '{data_object.type.value}'"
+            )
+
     def drop_table(self, table_name: TableName, exists: bool = True) -> None:
         """Drops a table.
 
@@ -1118,6 +1144,12 @@ def create_view(
         if properties.expressions:
             create_kwargs["properties"] = properties
 
+        if replace:
+            self._drop_data_object_on_type_mismatch(
+                view_name,
+                DataObjectType.VIEW if not materialized else DataObjectType.MATERIALIZED_VIEW,
+            )
+
         with source_queries[0] as query:
             self.execute(
                 exp.Create(
@@ -2483,6 +2515,35 @@ 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 _drop_data_object_on_type_mismatch(
+        self, target_name: TableName, expected_type: DataObjectType
+    ) -> bool:
+        """Drops a data object if it exists and is not of the expected type.
+
+        Args:
+            target_name: The name of the data object to check.
+            expected_type: The expected type of the data object.
+
+        Returns:
+            True if the data object exists and is of the expected type, False otherwise.
+        """
+        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:
+            if existing_data_objects[0].type == expected_type:
+                return True
+
+            logger.warning(
+                "Target data object '%s' is a %s and not a %s, dropping it",
+                target_table.sql(dialect=self.dialect),
+                existing_data_objects[0].type.value,
+                expected_type.value,
+            )
+            self.drop_data_object(existing_data_objects[0])
+        return False
+
     def _replace_by_key(
         self,
         target_table: TableName,

sqlmesh/core/engine_adapter/redshift.py

@@ -262,7 +262,9 @@ def replace_query(
         """
         import pandas as pd
 
-        if not isinstance(query_or_df, pd.DataFrame) or not self.table_exists(table_name):
+        table_exists = self._drop_data_object_on_type_mismatch(table_name, DataObjectType.TABLE)
+
+        if not isinstance(query_or_df, pd.DataFrame) or not table_exists:
             return super().replace_query(
                 table_name,
                 query_or_df,

sqlmesh/core/engine_adapter/shared.py

@@ -171,6 +171,9 @@ class DataObject(PydanticModel):
     def is_clustered(self) -> bool:
         return bool(self.clustering_key)
 
+    def to_table(self) -> exp.Table:
+        return exp.table_(self.name, db=self.schema_name, catalog=self.catalog, quoted=True)
+
 
 class CatalogSupport(Enum):
     # The engine has no concept of catalogs

sqlmesh/core/plan/builder.py

@@ -896,15 +896,15 @@ def _modified_and_added_snapshots(self) -> t.List[Snapshot]:
 
 def _should_force_rebuild(old: Snapshot, new: Snapshot) -> bool:
     if old.virtual_environment_mode != new.virtual_environment_mode:
-        # If the virtual environment mode has changed, then it's a breaking change
+        # If the virtual environment mode has changed, then we need to rebuild
         return True
     if old.model.kind.name == new.model.kind.name:
-        # If the kind hasn't changed, then it's not a breaking change
+        # If the kind hasn't changed, then we don't need to rebuild
         return False
     if not old.is_incremental or not new.is_incremental:
-        # If either is not incremental, then it's a breaking change
+        # If either is not incremental, then we need to rebuild
         return True
     if old.model.partitioned_by == new.model.partitioned_by:
-        # If the partitioning hasn't changed, then it's not a breaking change
+        # If the partitioning hasn't changed, then we don't need to rebuild
         return False
     return True

tests/core/test_integration.py

@@ -2525,6 +2525,52 @@ def test_virtual_environment_mode_dev_only(init_and_plan_context: t.Callable):
     assert "one" not in context.engine_adapter.columns("sushi.waiter_revenue_by_day")
 
 
+@time_machine.travel("2023-01-08 15:00:00 UTC")
+def test_virtual_environment_mode_dev_only_model_kind_change(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)
+
+    # Change to full kind
+    model = context.get_model("sushi.top_waiters")
+    model = model.copy(update={"kind": FullKind()})
+    context.upsert_model(model)
+    prod_plan = context.plan_builder("prod", skip_tests=True).build()
+    assert prod_plan.missing_intervals
+    assert prod_plan.requires_backfill
+    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"
+
+    context.state_sync.clear_cache()
+
+    # Change back to view
+    model = context.get_model("sushi.top_waiters")
+    model = model.copy(update={"kind": ViewKind()})
+    context.upsert_model(model)
+    prod_plan = context.plan_builder("prod", skip_tests=True).build()
+    assert prod_plan.requires_backfill
+    assert prod_plan.missing_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 == "view"
+
+    # Change to incremental
+    model = context.get_model("sushi.top_waiters")
+    model = model.copy(update={"kind": IncrementalUnmanagedKind()})
+    context.upsert_model(model)
+    prod_plan = context.plan_builder("prod", skip_tests=True).build()
+    assert prod_plan.requires_backfill
+    assert prod_plan.missing_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_restatement_plan_ignores_changes(init_and_plan_context: t.Callable):
     context, plan = init_and_plan_context("examples/sushi")