From 9664fca51cce47640095fb3396ba8faa98dc623a Mon Sep 17 00:00:00 2001 From: Srujan Kumar Gandla Date: Thu, 24 Sep 2026 12:52:22 -0400 Subject: [PATCH] fix: drop clustering key before dropping columns it references Signed-off-by: Srujan Kumar Gandla --- sqlmesh/core/engine_adapter/mixins.py | 21 +++++++---- tests/core/engine_adapter/test_snowflake.py | 41 ++++++++++++++++++++- 2 files changed, 54 insertions(+), 8 deletions(-) diff --git a/sqlmesh/core/engine_adapter/mixins.py b/sqlmesh/core/engine_adapter/mixins.py index bf4bb970a2..281d612684 100644 --- a/sqlmesh/core/engine_adapter/mixins.py +++ b/sqlmesh/core/engine_adapter/mixins.py @@ -383,20 +383,27 @@ def get_alter_operations( self.get_data_objects(target_table_schema, {target_table.name}), 0 ) + cluster_key_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, - ) + cluster_key_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)) + cluster_key_operation = TableAlterDropClusterKeyOperation( + target_table=current_table + ) + + if cluster_key_operation: + # The clustering key must be dropped (or replaced) before any of the + # columns it references can be dropped: engines like Snowflake reject + # dropping a column that belongs to a clustering key. + operations.insert(0, cluster_key_operation) return operations diff --git a/tests/core/engine_adapter/test_snowflake.py b/tests/core/engine_adapter/test_snowflake.py index 05f0133ec8..ac68a02062 100644 --- a/tests/core/engine_adapter/test_snowflake.py +++ b/tests/core/engine_adapter/test_snowflake.py @@ -10,7 +10,7 @@ from sqlmesh.core.dialect import normalize_model_name from sqlmesh.core.engine_adapter import SnowflakeEngineAdapter from sqlmesh.core.engine_adapter.base import EngineAdapter -from sqlmesh.core.engine_adapter.shared import DataObjectType +from sqlmesh.core.engine_adapter.shared import DataObject, DataObjectType from sqlmesh.core.model import load_sql_based_model from sqlmesh.core.model.definition import SqlModel from sqlmesh.core.node import IntervalUnit @@ -1049,6 +1049,45 @@ def test_alter_table_iceberg(mocker: MockerFixture, make_mocked_engine_adapter: assert to_sql_calls(adapter) == ['ALTER TABLE "test_table" ADD "b" INT'] +def test_alter_table_drops_clustering_key_before_drop_column( + mocker: MockerFixture, make_mocked_engine_adapter: t.Callable +): + mocker.patch("sqlmesh.core.engine_adapter.snowflake.SnowflakeEngineAdapter.set_current_catalog") + adapter = make_mocked_engine_adapter(SnowflakeEngineAdapter, default_catalog="test_catalog") + + current_table = {"a": "INT", "b": "INT"} + target_table = {"a": "INT"} + adapter.columns = lambda table_name, **kwargs: { + k: exp.DataType.build(v) + for k, v in (current_table if table_name == "test_table" else target_table).items() + } + + def _get_data_objects(schema_name, object_names=None, **kwargs): + table_name = next(iter(object_names or [])) + return [ + DataObject( + catalog="test_catalog", + schema="test_schema", + name=table_name, + type=DataObjectType.TABLE, + # the current table is clustered by the column being dropped + clustering_key='"b"' if table_name == "test_table" else None, + ) + ] + + mocker.patch.object(adapter, "get_data_objects", side_effect=_get_data_objects) + + alter_operations = adapter.get_alter_operations("test_table", "target_table") + adapter.alter_table(alter_operations) + + # Snowflake rejects dropping a column that belongs to a clustering key, so the + # clustering key must be dropped before the column is dropped. + assert to_sql_calls(adapter) == [ + 'ALTER TABLE "test_table" DROP CLUSTERING KEY', + 'ALTER TABLE "test_table" DROP COLUMN "b"', + ] + + def test_create_view_with_schema_and_grants( snowflake_mocked_engine_adapter: SnowflakeEngineAdapter, ):