fix(snowflake): change the clustering key before dropping a clustered column

sqlmesh/core/engine_adapter/mixins.py

@@ -383,23 +383,40 @@ def get_alter_operations(
             self.get_data_objects(target_table_schema, {target_table.name}), 0
         )
 
+        clustering_operation: t.Optional[TableAlterClusterByOperation] = None
         if current_table_info and target_table_info:
             if target_table_info.is_clustered:
                 if target_table_info.clustering_key and (
                     current_table_info.clustering_key != target_table_info.clustering_key
                 ):
-                    operations.append(
-                        TableAlterChangeClusterKeyOperation(
-                            target_table=current_table,
-                            clustering_key=target_table_info.clustering_key,
-                            dialect=self.dialect,
-                        )
+                    clustering_operation = TableAlterChangeClusterKeyOperation(
+                        target_table=current_table,
+                        clustering_key=target_table_info.clustering_key,
+                        dialect=self.dialect,
                     )
             elif current_table_info.is_clustered:
-                operations.append(TableAlterDropClusterKeyOperation(target_table=current_table))
+                clustering_operation = TableAlterDropClusterKeyOperation(target_table=current_table)
+
+        if clustering_operation:
+            self._apply_clustering_operation(operations, clustering_operation, current_table_info)
 
         return operations
 
+    def _apply_clustering_operation(
+        self,
+        operations: t.List[TableAlterOperation],
+        clustering_operation: "TableAlterClusterByOperation",
+        current_table_info: t.Any,
+    ) -> None:
+        """Insert the clustering operation into the alter operations list.
+
+        The default implementation appends the clustering change after all column
+        changes (historical order). Snowflake overrides this to handle the case
+        where the old key must be dropped before a column it names is removed,
+        and the new key must be set after new columns are added.
+        """
+        operations.append(clustering_operation)
+
 
 def logical_merge(
     engine_adapter: EngineAdapter,

sqlmesh/core/engine_adapter/snowflake.py

@@ -16,15 +16,18 @@
     ClusteredByMixin,
     RowDiffMixin,
     GrantsFromInfoSchemaMixin,
+    TableAlterChangeClusterKeyOperation,
+    TableAlterClusterByOperation,
+    TableAlterDropClusterKeyOperation,
 )
+from sqlmesh.core.schema_diff import TableAlterDropColumnOperation, TableAlterOperation
 from sqlmesh.core.engine_adapter.shared import (
     CatalogSupport,
     DataObject,
     DataObjectType,
     SourceQuery,
     set_catalog,
 )
-from sqlmesh.core.schema_diff import TableAlterOperation
 from sqlmesh.utils import optional_import, get_source_columns_to_types
 from sqlmesh.utils.errors import SQLMeshError
 from sqlmesh.utils.pandas import columns_to_types_from_dtypes
@@ -692,6 +695,63 @@ def clone_table(
             **kwargs,
         )
 
+    def _apply_clustering_operation(
+        self,
+        operations: t.List[TableAlterOperation],
+        clustering_operation: TableAlterClusterByOperation,
+        current_table_info: t.Any,
+    ) -> None:
+        """Apply clustering DDL in the order Snowflake requires.
+
+        Snowflake rejects DROP COLUMN while the column is still named in the clustering
+        key, and also rejects CLUSTER BY on a column that has not yet been added.  The
+        required sequence when the old key names a column being dropped is:
+
+        1. DROP CLUSTERING KEY — before any drop of a column the old key still names.
+        2. Column changes (drops and adds) — in their existing schema-diff order.
+        3. CLUSTER BY — after all column changes, if the target table is clustered.
+
+        When no column being dropped belongs to the current key the historical order is
+        kept: column changes first, then the single clustering change at the end.
+        """
+        clustering_key = getattr(current_table_info, "clustering_key", None)
+        if clustering_key:
+            referenced = self._clustering_key_column_names(clustering_key)
+            drop_index = next(
+                (
+                    i
+                    for i, op in enumerate(operations)
+                    if isinstance(op, TableAlterDropColumnOperation)
+                    and op.column_parts
+                    and normalize_identifiers(
+                        exp.to_identifier(op.column_parts[0].name), dialect=self.dialect
+                    ).name
+                    in referenced
+                ),
+                None,
+            )
+            if drop_index is not None:
+                drop_key_op = TableAlterDropClusterKeyOperation(
+                    target_table=clustering_operation.target_table
+                )
+                operations.insert(drop_index, drop_key_op)
+                if isinstance(clustering_operation, TableAlterChangeClusterKeyOperation):
+                    operations.append(clustering_operation)
+                return
+        # No key column is being dropped: append at end (historical order).
+        operations.append(clustering_operation)
+
+    def _clustering_key_column_names(self, clustering_key: str) -> t.Set[str]:
+        expressions = TableAlterChangeClusterKeyOperation(
+            target_table=exp.to_table("_"),
+            clustering_key=clustering_key,
+            dialect=self.dialect,
+        ).cluster_key_expressions
+        return {
+            normalize_identifiers(expression, dialect=self.dialect).name
+            for expression in expressions
+        }
+
     def alter_table(
         self,
         alter_expressions: t.Union[t.List[exp.Alter], t.List[TableAlterOperation]],

tests/core/engine_adapter/test_bigquery.py

@@ -13,7 +13,9 @@
 import sqlmesh.core.dialect as d
 from sqlmesh.core.engine_adapter import BigQueryEngineAdapter
 from sqlmesh.core.engine_adapter.bigquery import select_partitions_expr
-from sqlmesh.core.engine_adapter.shared import DataObjectType
+from sqlmesh.core.engine_adapter.mixins import TableAlterDropClusterKeyOperation
+from sqlmesh.core.engine_adapter.shared import DataObject, DataObjectType
+from sqlmesh.core.schema_diff import TableAlterDropColumnOperation
 from sqlmesh.core.node import IntervalUnit
 from sqlmesh.utils import AttributeDict
 from sqlmesh.utils.errors import SQLMeshError
@@ -1380,3 +1382,108 @@ def test_sync_grants_config_no_schema(
 
     with pytest.raises(ValueError, match="Table test_table does not have a schema \\(dataset\\)"):
         adapter.sync_grants_config(relation, new_grants_config)
+
+
+def _mock_bq_data_objects(
+    mocker: MockerFixture,
+    adapter: BigQueryEngineAdapter,
+    current_obj: DataObject,
+    target_obj: DataObject,
+) -> None:
+    """Patch _get_data_objects on *adapter* to return the supplied DataObjects."""
+
+    def _side_effect(
+        schema_name: t.Any, object_names: t.Optional[t.Set[str]] = None
+    ) -> t.List[DataObject]:
+        results = []
+        for obj in [current_obj, target_obj]:
+            if object_names is None or obj.name in object_names:
+                results.append(obj)
+        return results
+
+    mocker.patch.object(adapter, "_get_data_objects", side_effect=_side_effect)
+
+
+def test_bigquery_clustering_column_drop_before_clustering_change(
+    make_mocked_engine_adapter: t.Callable, mocker: MockerFixture
+) -> None:
+    """BigQuery keeps the historical order: DROP COLUMN before clustering key change.
+
+    Dropping a column that belongs to the current clustering key must NOT trigger
+    a reorder.  BigQuery must emit DROP COLUMN first, then DROP CLUSTERING KEY.
+    """
+    adapter = make_mocked_engine_adapter(BigQueryEngineAdapter, patch_get_data_objects=False)
+
+    current_columns = {"tenant_id": "INT64", "partition_month": "INT64"}
+    target_columns = {"tenant_id": "INT64"}
+    adapter.columns = lambda table_name, **kwargs: {
+        k: exp.DataType.build(v, dialect="bigquery")
+        for k, v in (
+            current_columns if table_name == "project.dataset.current_table" else target_columns
+        ).items()
+    }
+
+    _mock_bq_data_objects(
+        mocker,
+        adapter,
+        current_obj=DataObject(
+            catalog="project",
+            schema="dataset",
+            name="current_table",
+            type=DataObjectType.TABLE,
+            clustering_key="(partition_month)",
+        ),
+        target_obj=DataObject(
+            catalog="project",
+            schema="dataset",
+            name="target_table",
+            type=DataObjectType.TABLE,
+            clustering_key=None,
+        ),
+    )
+
+    operations = adapter.get_alter_operations(
+        "project.dataset.current_table", "project.dataset.target_table"
+    )
+
+    assert len(operations) == 2
+    assert isinstance(operations[0], TableAlterDropColumnOperation)
+    assert isinstance(operations[1], TableAlterDropClusterKeyOperation)
+
+
+def test_bigquery_no_column_change_clustering_drop_only(
+    make_mocked_engine_adapter: t.Callable, mocker: MockerFixture
+) -> None:
+    """BigQuery: same columns, only clustering key removed — one DROP CLUSTERING KEY."""
+    adapter = make_mocked_engine_adapter(BigQueryEngineAdapter, patch_get_data_objects=False)
+
+    columns = {"tenant_id": "INT64", "partition_month": "INT64"}
+    adapter.columns = lambda table_name, **kwargs: {
+        k: exp.DataType.build(v, dialect="bigquery") for k, v in columns.items()
+    }
+
+    _mock_bq_data_objects(
+        mocker,
+        adapter,
+        current_obj=DataObject(
+            catalog="project",
+            schema="dataset",
+            name="current_table",
+            type=DataObjectType.TABLE,
+            clustering_key="(partition_month)",
+        ),
+        target_obj=DataObject(
+            catalog="project",
+            schema="dataset",
+            name="target_table",
+            type=DataObjectType.TABLE,
+            clustering_key=None,
+        ),
+    )
+
+    operations = adapter.get_alter_operations(
+        "project.dataset.current_table", "project.dataset.target_table"
+    )
+
+    assert len(operations) == 1
+    assert isinstance(operations[0], TableAlterDropClusterKeyOperation)

tests/core/engine_adapter/test_redshift.py

@@ -11,6 +11,7 @@
 import sqlmesh.core.dialect as d
 from sqlmesh.core.engine_adapter import RedshiftEngineAdapter
 from sqlmesh.core.engine_adapter.shared import DataObject, DataObjectType
+from sqlmesh.core.schema_diff import TableAlterDropColumnOperation
 from sqlmesh.core.model import load_sql_based_model
 from sqlmesh.core.model.definition import SqlModel
 from sqlmesh.utils.errors import SQLMeshError
@@ -912,3 +913,46 @@ def test_merge_logical(
         'INSERT INTO "target" ("ID", "ts") SELECT "ID", "ts" FROM (SELECT "ID" AS "ID", "ts" AS "ts", ROW_NUMBER() OVER (PARTITION BY "ID" ORDER BY "ID") AS _row_number FROM "__temp_test_abcdefgh") AS _t WHERE _row_number = 1',
         'DROP TABLE IF EXISTS "__temp_test_abcdefgh"',
     ]
+
+
+def test_redshift_no_clustering_operations_on_column_drop(
+    make_mocked_engine_adapter: t.Callable,
+) -> None:
+    """Dropping a column on Redshift never emits clustering statements."""
+    adapter = make_mocked_engine_adapter(RedshiftEngineAdapter)
+    adapter.cursor.fetchall.return_value = []
+
+    current_columns = {"TENANT_ID": "INT", "PARTITION_MONTH": "INT"}
+    target_columns = {"TENANT_ID": "INT"}
+    adapter.columns = lambda table_name, **kwargs: {
+        k: exp.DataType.build(v)
+        for k, v in (current_columns if "current" in str(table_name) else target_columns).items()
+    }
+
+    operations = adapter.get_alter_operations("current_table", "target_table")
+
+    assert len(operations) == 1
+    assert all(isinstance(op, TableAlterDropColumnOperation) for op in operations)
+
+
+def test_redshift_multiple_column_drops_no_clustering_operations(
+    make_mocked_engine_adapter: t.Callable,
+) -> None:
+    """Dropping multiple columns on Redshift preserves the schema-diff order.
+
+    No clustering operations must be injected between or after the column drops.
+    """
+    adapter = make_mocked_engine_adapter(RedshiftEngineAdapter)
+    adapter.cursor.fetchall.return_value = []
+
+    current_columns = {"EXTRA": "INT", "PARTITION_MONTH": "INT", "TENANT_ID": "INT"}
+    target_columns = {"TENANT_ID": "INT"}
+    adapter.columns = lambda table_name, **kwargs: {
+        k: exp.DataType.build(v)
+        for k, v in (current_columns if "current" in str(table_name) else target_columns).items()
+    }
+
+    operations = adapter.get_alter_operations("current_table", "target_table")
+
+    assert len(operations) == 2
+    assert all(isinstance(op, TableAlterDropColumnOperation) for op in operations)