Fix: Add CASCADE modifier when dropping columns in Postgres / Redshift (#3427)

Author: izeigerman

sqlmesh/core/engine_adapter/postgres.py

@@ -61,6 +61,7 @@ class PostgresEngineAdapter(
                 exp.DataType.build("BPCHAR", dialect=DIALECT).this
             },
         },
+        drop_cascade=True,
     )
 
     def _fetch_native_df(

sqlmesh/core/engine_adapter/redshift.py

@@ -58,6 +58,7 @@ class RedshiftEngineAdapter(
             exp.DataType.build("CHAR", dialect=DIALECT).this: 4096,
             exp.DataType.build("VARCHAR", dialect=DIALECT).this: 65535,
         },
+        drop_cascade=True,
     )
     VARIABLE_LENGTH_DATA_TYPES = {
         "char",

sqlmesh/core/schema_diff.py

@@ -169,6 +169,7 @@ class TableAlterOperation(PydanticModel):
     expected_table_struct: exp.DataType
     add_position: t.Optional[TableAlterColumnPosition] = None
     current_type: t.Optional[exp.DataType] = None
+    cascade: bool = False
 
     @classmethod
     def add(
@@ -192,13 +193,15 @@ def drop(
         columns: t.Union[TableAlterColumn, t.List[TableAlterColumn]],
         expected_table_struct: t.Union[str, exp.DataType],
         column_type: t.Optional[t.Union[str, exp.DataType]] = None,
+        cascade: bool = False,
     ) -> TableAlterOperation:
         column_type = exp.DataType.build(column_type) if column_type else exp.DataType.build("INT")
         return cls(
             op=TableAlterOperationType.DROP,
             columns=ensure_list(columns),
             column_type=column_type,
             expected_table_struct=exp.DataType.build(expected_table_struct),
+            cascade=cascade,
         )
 
     @classmethod
@@ -274,7 +277,9 @@ def expression(
             return alter_table
         elif self.is_drop:
             alter_table = exp.Alter(this=exp.to_table(table_name), kind="TABLE")
-            drop_column = exp.Drop(this=self.column(array_element_selector), kind="COLUMN")
+            drop_column = exp.Drop(
+                this=self.column(array_element_selector), kind="COLUMN", cascade=self.cascade
+            )
             alter_table.set("actions", [drop_column])
             return alter_table
         else:
@@ -313,6 +318,7 @@ class SchemaDiffer(PydanticModel):
             into a FLOAT64 column just fine.
         support_coercing_compatible_types: Whether or not the engine for which the diff is being computed supports direct
             coercion of compatible types.
+        drop_cascade: Whether to add CASCADE modifier when dropping a column.
         parameterized_type_defaults: Default values for parameterized data types. Dict key is a sqlglot exp.DataType.Type,
             but in the engine adapter specification we build it from the dialect string instead of specifying it directly.
             Example: `exp.DataType.build("STRING", dialect=DIALECT).this` instead of the underlying `exp.DataType.Type.TEXT`
@@ -336,6 +342,7 @@ class SchemaDiffer(PydanticModel):
         default_factory=dict, alias="coerceable_types"
     )
     support_coercing_compatible_types: bool = False
+    drop_cascade: bool = False
     parameterized_type_defaults: t.Dict[
         exp.DataType.Type, t.List[t.Tuple[t.Union[int, float], ...]]
     ] = {}
@@ -458,7 +465,9 @@ def _drop_operation(
         assert column_kwarg
         struct.expressions.pop(column_pos)
         operations.append(
-            TableAlterOperation.drop(columns, root_struct.copy(), column_kwarg.args["kind"])
+            TableAlterOperation.drop(
+                columns, root_struct.copy(), column_kwarg.args["kind"], cascade=self.drop_cascade
+            )
         )
         return operations
 

tests/core/engine_adapter/test_postgres.py

@@ -141,3 +141,23 @@ def test_merge_version_lt_15(
         'INSERT INTO "target" ("ID", "ts", "val") SELECT DISTINCT ON ("ID") "ID", "ts", "val" FROM "__temp_test_abcdefgh"',
         'DROP TABLE IF EXISTS "__temp_test_abcdefgh"',
     ]
+
+
+def test_alter_table_drop_column_cascade(make_mocked_engine_adapter: t.Callable):
+    adapter = make_mocked_engine_adapter(PostgresEngineAdapter)
+
+    current_table_name = "test_table"
+    target_table_name = "target_table"
+
+    def table_columns(table_name: str) -> t.Dict[str, exp.DataType]:
+        if table_name == current_table_name:
+            return {"id": exp.DataType.build("int"), "test_column": exp.DataType.build("int")}
+        else:
+            return {"id": exp.DataType.build("int")}
+
+    adapter.columns = table_columns
+
+    adapter.alter_table(adapter.get_alter_expressions(current_table_name, target_table_name))
+    assert to_sql_calls(adapter) == [
+        'ALTER TABLE "test_table" DROP COLUMN "test_column" CASCADE',
+    ]

tests/core/engine_adapter/test_redshift.py

@@ -304,3 +304,21 @@ def test_create_view(adapter: t.Callable):
         'DROP VIEW IF EXISTS "test_view" CASCADE',
         'CREATE VIEW "test_view" ("a", "b") AS SELECT "cola" FROM "table" WITH NO SCHEMA BINDING',
     ]
+
+
+def test_alter_table_drop_column_cascade(adapter: t.Callable):
+    current_table_name = "test_table"
+    target_table_name = "target_table"
+
+    def table_columns(table_name: str) -> t.Dict[str, exp.DataType]:
+        if table_name == current_table_name:
+            return {"id": exp.DataType.build("int"), "test_column": exp.DataType.build("int")}
+        else:
+            return {"id": exp.DataType.build("int")}
+
+    adapter.columns = table_columns
+
+    adapter.alter_table(adapter.get_alter_expressions(current_table_name, target_table_name))
+    assert to_sql_calls(adapter) == [
+        'ALTER TABLE "test_table" DROP COLUMN "test_column" CASCADE',
+    ]

tests/core/test_schema_diff.py

@@ -45,6 +45,32 @@ def test_schema_diff_calculate():
     ]
 
 
+def test_schema_diff_drop_cascade():
+    alter_expressions = SchemaDiffer(
+        **{
+            "support_positional_add": False,
+            "support_nested_operations": False,
+            "array_element_selector": "",
+            "drop_cascade": True,
+        }
+    ).compare_columns(
+        "apply_to_table",
+        {
+            "id": exp.DataType.build("INT"),
+            "name": exp.DataType.build("STRING"),
+            "price": exp.DataType.build("DOUBLE"),
+        },
+        {
+            "id": exp.DataType.build("INT"),
+            "name": exp.DataType.build("STRING"),
+        },
+    )
+
+    assert [x.sql() for x in alter_expressions] == [
+        """ALTER TABLE apply_to_table DROP COLUMN price CASCADE"""
+    ]
+
+
 def test_schema_diff_calculate_type_transitions():
     alter_expressions = SchemaDiffer(
         **{