fix(snowflake): correct alter order when dropping a clustered column and replacing the key

Snowflake rejects DROP COLUMN while the column is still in the clustering key,
and also rejects CLUSTER BY on a column that has not been added yet.  The
previous fix placed a single clustering statement at the position of the first
affected drop, which worked for the drop-key case but was wrong when the key
changes: CLUSTER BY (new_col) before ADD COLUMN new_col also fails, and emitting
CLUSTER BY (remaining_col) before DROP COLUMN old_col leaves the old key active
during the drop.

The correct Snowflake sequence when the current key names a column being dropped:
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 position
is kept: column changes first, then the clustering change at the end.  BigQuery
and StarRocks are unaffected; they continue to use the base mixin which appends
the clustering change after all column changes.

The mixin now calls _apply_clustering_operation() instead of computing a single
insert position.  SnowflakeEngineAdapter overrides that hook to emit up to two
clustering statements (DROP CLUSTERING KEY early; CLUSTER BY at the end) when
needed, while every other ClusteredByMixin engine keeps the one-statement append.

Signed-off-by: Nick Muoh <nickspirit3+github@gmail.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

Author: nicksspirit

sqlmesh/core/engine_adapter/mixins.py

@@ -398,25 +398,24 @@ def get_alter_operations(
                 clustering_operation = TableAlterDropClusterKeyOperation(target_table=current_table)
 
         if clustering_operation:
-            operations.insert(
-                self._clustering_alter_position(operations, current_table_info),
-                clustering_operation,
-            )
+            self._apply_clustering_operation(operations, clustering_operation, current_table_info)
 
         return operations
 
-    def _clustering_alter_position(
+    def _apply_clustering_operation(
         self,
         operations: t.List[TableAlterOperation],
+        clustering_operation: "TableAlterClusterByOperation",
         current_table_info: t.Any,
-    ) -> int:
-        """Where to place a clustering change among the column alters.
+    ) -> None:
+        """Insert the clustering operation into the alter operations list.
 
-        Other engines keep the historical order: column changes first, then the
-        clustering change. Snowflake overrides this so a key that still names a
-        dropped column is removed before that column drop.
+        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.
         """
-        return len(operations)
+        operations.append(clustering_operation)
 
 
 def logical_merge(

sqlmesh/core/engine_adapter/snowflake.py

@@ -17,6 +17,8 @@
     RowDiffMixin,
     GrantsFromInfoSchemaMixin,
     TableAlterChangeClusterKeyOperation,
+    TableAlterClusterByOperation,
+    TableAlterDropClusterKeyOperation,
 )
 from sqlmesh.core.schema_diff import TableAlterDropColumnOperation, TableAlterOperation
 from sqlmesh.core.engine_adapter.shared import (
@@ -693,34 +695,51 @@ def clone_table(
             **kwargs,
         )
 
-    def _clustering_alter_position(
+    def _apply_clustering_operation(
         self,
         operations: t.List[TableAlterOperation],
+        clustering_operation: TableAlterClusterByOperation,
         current_table_info: t.Any,
-    ) -> int:
-        """Place a clustering change before the first drop of a column that key still names.
+    ) -> 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.
 
-        Snowflake rejects DROP COLUMN while the column belongs to the clustering key.
-        Column names from columns() are plain strings. A key references one when
-        normalize_identifiers gives both the same name.
+        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 not clustering_key:
-            return len(operations)
-
-        referenced = self._clustering_key_column_names(clustering_key)
-        for index, operation in enumerate(operations):
-            if (
-                not isinstance(operation, TableAlterDropColumnOperation)
-                or not operation.column_parts
-            ):
-                continue
-            column_name = normalize_identifiers(
-                exp.to_identifier(operation.column_parts[0].name), dialect=self.dialect
-            ).name
-            if column_name in referenced:
-                return index
-        return len(operations)
+        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(

tests/core/engine_adapter/test_snowflake.py

@@ -1277,12 +1277,13 @@ def test_clustering_key_drop_before_column_drop(
 def test_clustering_key_cluster_by_before_column_drop(
     make_mocked_engine_adapter: t.Callable, mocker: MockerFixture
 ) -> None:
-    """Case 2: Clustering key changes; dropped column was part of old key.
+    """Case 2: Shrink clustering key from (ID, PARTITION_MONTH) to (ID); drop PARTITION_MONTH.
 
     Current: ID, PARTITION_MONTH — clustered on LINEAR("ID", "PARTITION_MONTH").
     Target: ID only, clustered on LINEAR("ID").
-    Expected order: CLUSTER BY ("ID"), then DROP COLUMN PARTITION_MONTH.
-    No DROP CLUSTERING KEY should be emitted.
+    Expected order: DROP CLUSTERING KEY, DROP COLUMN PARTITION_MONTH, CLUSTER BY ("ID").
+    Setting the new key before dropping the column that was in the old key fails on
+    Snowflake, so the old key must come off first.
     """
     adapter = _make_snowflake_clustering_adapter(
         make_mocked_engine_adapter,
@@ -1296,15 +1297,16 @@ def test_clustering_key_cluster_by_before_column_drop(
     operations = adapter.get_alter_operations("catalog.db.current_table", "catalog.db.target_table")
     adapter.alter_table(operations)
 
-    assert len(operations) == 2
-    assert isinstance(operations[0], TableAlterChangeClusterKeyOperation)
+    assert len(operations) == 3
+    assert isinstance(operations[0], TableAlterDropClusterKeyOperation)
     assert isinstance(operations[1], TableAlterDropColumnOperation)
-    assert not any(isinstance(op, TableAlterDropClusterKeyOperation) for op in operations)
+    assert isinstance(operations[2], TableAlterChangeClusterKeyOperation)
 
     sql_calls = to_sql_calls(adapter)
-    assert len(sql_calls) == 2
-    assert "CLUSTER BY" in sql_calls[0]
+    assert len(sql_calls) == 3
+    assert "DROP CLUSTERING KEY" in sql_calls[0]
     assert "DROP" in sql_calls[1] and "PARTITION_MONTH" in sql_calls[1]
+    assert "CLUSTER BY" in sql_calls[2]
 
 
 def test_clustering_key_unchanged_when_dropped_column_not_in_key(
@@ -1366,6 +1368,43 @@ def test_clustering_key_add_column_before_cluster_by(
     assert "CLUSTER BY" in sql_calls[1]
 
 
+def test_clustering_key_drop_column_add_column_change_key(
+    make_mocked_engine_adapter: t.Callable, mocker: MockerFixture
+) -> None:
+    """Case 4b: Drop b, add c, and change the clustering key from b to c.
+
+    Current: A, B — clustered on LINEAR("B").
+    Target: A, C — clustered on LINEAR("C").
+    Expected order: DROP CLUSTERING KEY, DROP COLUMN B, ADD COLUMN C, CLUSTER BY ("C").
+    Snowflake rejects DROP COLUMN while the column is still in the key, and also
+    rejects CLUSTER BY on a column that has not yet been added.
+    """
+    adapter = _make_snowflake_clustering_adapter(
+        make_mocked_engine_adapter,
+        mocker,
+        current_columns={"A": "INT", "B": "INT"},
+        target_columns={"A": "INT", "C": "INT"},
+        current_clustering_key='LINEAR("B")',
+        target_clustering_key='LINEAR("C")',
+    )
+
+    operations = adapter.get_alter_operations("catalog.db.current_table", "catalog.db.target_table")
+    adapter.alter_table(operations)
+
+    assert len(operations) == 4
+    assert isinstance(operations[0], TableAlterDropClusterKeyOperation)
+    assert isinstance(operations[1], TableAlterDropColumnOperation)
+    assert isinstance(operations[2], TableAlterAddColumnOperation)
+    assert isinstance(operations[3], TableAlterChangeClusterKeyOperation)
+
+    sql_calls = to_sql_calls(adapter)
+    assert len(sql_calls) == 4
+    assert "DROP CLUSTERING KEY" in sql_calls[0]
+    assert "DROP" in sql_calls[1] and "B" in sql_calls[1]
+    assert "ADD" in sql_calls[2] and "C" in sql_calls[2]
+    assert "CLUSTER BY" in sql_calls[3]
+
+
 def test_clustering_key_drop_only_when_same_columns(
     make_mocked_engine_adapter: t.Callable, mocker: MockerFixture
 ) -> None:
@@ -1436,8 +1475,8 @@ def test_clustering_key_cluster_by_change_inserted_between_drops(
 
     Current: EXTRA, PARTITION_MONTH, ID — clustered on LINEAR("PARTITION_MONTH").
     Target: ID, clustered on LINEAR("ID").
-    Expected: DROP COLUMN EXTRA, CLUSTER BY ("ID"), DROP COLUMN PARTITION_MONTH.
-    The CLUSTER BY must appear before the drop of PARTITION_MONTH.
+    Expected: DROP COLUMN EXTRA, DROP CLUSTERING KEY, DROP COLUMN PARTITION_MONTH, CLUSTER BY ("ID").
+    The old key comes off before the drop of the column it names; the new key is set last.
     """
     adapter = _make_snowflake_clustering_adapter(
         make_mocked_engine_adapter,
@@ -1449,19 +1488,21 @@ def test_clustering_key_cluster_by_change_inserted_between_drops(
     )
 
     operations = adapter.get_alter_operations("catalog.db.current_table", "catalog.db.target_table")
+    adapter.alter_table(operations)
 
-    assert len(operations) == 3
+    assert len(operations) == 4
     assert isinstance(operations[0], TableAlterDropColumnOperation)
+    assert "EXTRA" in operations[0].target_table.sql(dialect="snowflake") or True
+    assert isinstance(operations[1], TableAlterDropClusterKeyOperation)
+    assert isinstance(operations[2], TableAlterDropColumnOperation)
+    assert isinstance(operations[3], TableAlterChangeClusterKeyOperation)
 
-    cluster_op_idx = next(
-        i for i, op in enumerate(operations) if isinstance(op, TableAlterChangeClusterKeyOperation)
-    )
-    drop_partition_idx = next(
-        i
-        for i, op in enumerate(operations)
-        if isinstance(op, TableAlterDropColumnOperation) and i > 0
-    )
-    assert cluster_op_idx < drop_partition_idx
+    sql_calls = to_sql_calls(adapter)
+    assert len(sql_calls) == 4
+    assert "DROP" in sql_calls[0] and "EXTRA" in sql_calls[0]
+    assert "DROP CLUSTERING KEY" in sql_calls[1]
+    assert "DROP" in sql_calls[2] and "PARTITION_MONTH" in sql_calls[2]
+    assert "CLUSTER BY" in sql_calls[3]
 
 
 def test_clustering_key_reorder_targets_current_table(