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

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

Author: 3 people

sqlmesh/core/engine_adapter/mixins.py

@@ -18,6 +18,7 @@
 
 if t.TYPE_CHECKING:
     from sqlmesh.core._typing import TableName
+    from sqlmesh.core.engine_adapter.shared import DataObject
     from sqlmesh.core.engine_adapter._typing import (
         DCL,
         DF,
@@ -383,23 +384,53 @@ 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,
-                        )
-                    )
+            if target_table_info.is_clustered and target_table_info.clustering_key:
+                # Include an unchanged key. Snowflake may have to drop it and set it
+                # again when a key column is dropped and re-added. Other engines ignore
+                # that operation in _apply_clustering_operation.
+                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 _clustering_key_is_unchanged(
+        self,
+        clustering_operation: "TableAlterClusterByOperation",
+        current_table_info: t.Optional["DataObject"],
+    ) -> bool:
+        return (
+            isinstance(clustering_operation, TableAlterChangeClusterKeyOperation)
+            and current_table_info is not None
+            and current_table_info.clustering_key == clustering_operation.clustering_key
+        )
+
+    def _apply_clustering_operation(
+        self,
+        operations: t.List[TableAlterOperation],
+        clustering_operation: "TableAlterClusterByOperation",
+        current_table_info: t.Optional["DataObject"],
+    ) -> None:
+        """Insert the clustering operation into the alter operations list.
+
+        The default implementation appends a real clustering change after all column
+        changes. An unchanged key is ignored here. Snowflake overrides this so it can
+        drop that key before removing a column the key names, then set the same key
+        after the column comes back.
+        """
+        if self._clustering_key_is_unchanged(clustering_operation, current_table_info):
+            return
+        operations.append(clustering_operation)
+
 
 def logical_merge(
     engine_adapter: EngineAdapter,

sqlmesh/core/engine_adapter/snowflake.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import contextlib
+import itertools
 import logging
 import typing as t
 
@@ -16,15 +17,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 +696,62 @@ def clone_table(
             **kwargs,
         )
 
+    def _apply_clustering_operation(
+        self,
+        operations: t.List[TableAlterOperation],
+        clustering_operation: TableAlterClusterByOperation,
+        current_table_info: t.Optional[DataObject],
+    ) -> None:
+        """Apply clustering DDL in the order Snowflake requires.
+
+        DROP COLUMN fails while the column is still in the clustering key, including
+        inside an expression. CLUSTER BY fails on a column that does not exist yet. A
+        type change that becomes DROP plus ADD still needs the key dropped first even
+        when the key text is unchanged. An unchanged key with no matching column drop
+        emits nothing. DESCRIBE names are compared as stored.
+        """
+        clustering_key = current_table_info.clustering_key if current_table_info else None
+        operations_include_drop_column = any(
+            isinstance(op, TableAlterDropColumnOperation) for op in operations
+        )
+
+        if clustering_key and operations_include_drop_column:
+            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 op.column_parts[0].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
+        if self._clustering_key_is_unchanged(clustering_operation, current_table_info):
+            return
+        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
+        cluster_columns = itertools.chain.from_iterable(
+            expression.find_all(exp.Column) for expression in expressions
+        )
+        return {
+            normalize_identifiers(column, dialect=self.dialect).name for column in cluster_columns
+        }
+
     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,12 @@
 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 (
+    TableAlterClusterByOperation,
+    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 +1385,158 @@ 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)
+
+
+def test_unchanged_clustering_key_is_not_altered_when_column_added(
+    adapter: BigQueryEngineAdapter, mocker: MockerFixture
+):
+    # Adding a column to a clustered table must not emit a clustering change when the key
+    # itself is unchanged. BigQuery applies a new key by rewriting the clustered columns.
+    current_columns = {"c1": exp.DataType.build("INT64", dialect="bigquery")}
+    target_columns = {
+        **current_columns,
+        "c2": exp.DataType.build("INT64", dialect="bigquery"),
+    }
+
+    def columns(table_name: str, include_pseudo_columns: bool = False) -> t.Dict[str, exp.DataType]:
+        if exp.to_table(table_name).name == "current_table":
+            return current_columns
+        return target_columns
+
+    mocker.patch.object(adapter, "columns", side_effect=columns)
+
+    def get_data_objects(
+        schema_name: exp.Table,
+        object_names: t.Optional[t.Set[str]] = None,
+        safe_to_cache: bool = False,
+    ) -> t.List[DataObject]:
+        assert object_names is not None
+        return [
+            DataObject(
+                schema=schema_name.db,
+                name=next(iter(object_names)),
+                type=DataObjectType.TABLE,
+                clustering_key="(c1)",
+            )
+        ]
+
+    mocker.patch.object(adapter, "get_data_objects", side_effect=get_data_objects)
+    # Let a wrongly emitted clustering change surface as the column-rewrite UPDATE.
+    mocker.patch.object(adapter, "_get_table", return_value=mocker.Mock())
+    mocker.patch.object(adapter, "_db_call")
+
+    operations = adapter.get_alter_operations(
+        "test_schema.current_table", "test_schema.target_table"
+    )
+    assert not any(isinstance(operation, TableAlterClusterByOperation) for operation in operations)
+
+    adapter.alter_table(operations)
+
+    assert _to_sql_calls(adapter) == [
+        "ALTER TABLE `test_schema`.`current_table` ADD COLUMN `c2` INT64"
+    ]

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)