From 9596b85d1cf781a22341deeb4479e1dbba653fd1 Mon Sep 17 00:00:00 2001 From: Nick Muoh Date: Sun, 1 Nov 2026 18:11:32 -0800 Subject: [PATCH] fix(snowflake): correct alter order when dropping a clustered column and replacing the key MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 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. Co-authored-by: Cursor Signed-off-by: Nick Muoh --- sqlmesh/core/engine_adapter/mixins.py | 31 +- sqlmesh/core/engine_adapter/snowflake.py | 62 ++- tests/core/engine_adapter/test_bigquery.py | 109 ++++- tests/core/engine_adapter/test_redshift.py | 44 ++ tests/core/engine_adapter/test_snowflake.py | 433 +++++++++++++++++++- tests/core/engine_adapter/test_starrocks.py | 61 ++- 6 files changed, 729 insertions(+), 11 deletions(-) diff --git a/sqlmesh/core/engine_adapter/mixins.py b/sqlmesh/core/engine_adapter/mixins.py index bf4bb970a2..040680d085 100644 --- a/sqlmesh/core/engine_adapter/mixins.py +++ b/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, diff --git a/sqlmesh/core/engine_adapter/snowflake.py b/sqlmesh/core/engine_adapter/snowflake.py index 758e4847c1..8a503833e2 100644 --- a/sqlmesh/core/engine_adapter/snowflake.py +++ b/sqlmesh/core/engine_adapter/snowflake.py @@ -16,7 +16,11 @@ ClusteredByMixin, RowDiffMixin, GrantsFromInfoSchemaMixin, + TableAlterChangeClusterKeyOperation, + TableAlterClusterByOperation, + TableAlterDropClusterKeyOperation, ) +from sqlmesh.core.schema_diff import TableAlterDropColumnOperation, TableAlterOperation from sqlmesh.core.engine_adapter.shared import ( CatalogSupport, DataObject, @@ -24,7 +28,6 @@ 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]], diff --git a/tests/core/engine_adapter/test_bigquery.py b/tests/core/engine_adapter/test_bigquery.py index 134f144df1..2d8103bc08 100644 --- a/tests/core/engine_adapter/test_bigquery.py +++ b/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) diff --git a/tests/core/engine_adapter/test_redshift.py b/tests/core/engine_adapter/test_redshift.py index 7211244e91..a3260ed4cd 100644 --- a/tests/core/engine_adapter/test_redshift.py +++ b/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) diff --git a/tests/core/engine_adapter/test_snowflake.py b/tests/core/engine_adapter/test_snowflake.py index 89f19fc588..7a2c6140a6 100644 --- a/tests/core/engine_adapter/test_snowflake.py +++ b/tests/core/engine_adapter/test_snowflake.py @@ -11,7 +11,12 @@ 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.mixins import ( + TableAlterChangeClusterKeyOperation, + TableAlterDropClusterKeyOperation, +) +from sqlmesh.core.engine_adapter.shared import DataObject, DataObjectType +from sqlmesh.core.schema_diff import TableAlterAddColumnOperation, TableAlterDropColumnOperation from sqlmesh.core.model import load_sql_based_model from sqlmesh.core.model.definition import SqlModel from sqlmesh.core.node import IntervalUnit @@ -1187,3 +1192,429 @@ def test_drop_catalog(snowflake_mocked_engine_adapter: SnowflakeEngineAdapter) - """SELECT 1 FROM "INFORMATION_SCHEMA"."DATABASES" WHERE "DATABASE_NAME" = 'foo' AND "COMMENT" = 'sqlmesh_managed'""", 'DROP DATABASE IF EXISTS "foo"', ] + + +def _make_snowflake_clustering_adapter( + make_mocked_engine_adapter: t.Callable, + mocker: MockerFixture, + current_columns: t.Dict[str, str], + target_columns: t.Dict[str, str], + current_clustering_key: t.Optional[str] = None, + target_clustering_key: t.Optional[str] = None, +) -> SnowflakeEngineAdapter: + """Create a mocked SnowflakeEngineAdapter wired for clustering-key tests.""" + adapter = make_mocked_engine_adapter( + SnowflakeEngineAdapter, default_catalog="catalog", patch_get_data_objects=False + ) + mocker.patch("sqlmesh.core.engine_adapter.snowflake.SnowflakeEngineAdapter.set_current_catalog") + + adapter.columns = lambda table_name, **kwargs: { + k: exp.DataType.build(v) + for k, v in ( + current_columns if table_name == "catalog.db.current_table" else target_columns + ).items() + } + + current_obj = DataObject( + catalog="catalog", + schema="db", + name="current_table", + type=DataObjectType.TABLE, + clustering_key=current_clustering_key, + ) + target_obj = DataObject( + catalog="catalog", + schema="db", + name="target_table", + type=DataObjectType.TABLE, + clustering_key=target_clustering_key, + ) + + def _mock_get_data_objects( + 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=_mock_get_data_objects) + return adapter + + +def test_clustering_key_drop_before_column_drop( + make_mocked_engine_adapter: t.Callable, mocker: MockerFixture +) -> None: + """Case 1: Both clustered columns dropped; clustering key removed. + + Current: TENANT_ID, PARTITION_MONTH — clustered on LINEAR("TENANT_ID", "PARTITION_MONTH"). + Target: TENANT_ID only, no clustering. + Expected order: DROP CLUSTERING KEY, then DROP COLUMN PARTITION_MONTH. + """ + adapter = _make_snowflake_clustering_adapter( + make_mocked_engine_adapter, + mocker, + current_columns={"TENANT_ID": "INT", "PARTITION_MONTH": "INT"}, + target_columns={"TENANT_ID": "INT"}, + current_clustering_key='LINEAR("TENANT_ID", "PARTITION_MONTH")', + target_clustering_key=None, + ) + + 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], TableAlterDropClusterKeyOperation) + assert isinstance(operations[1], TableAlterDropColumnOperation) + + sql_calls = to_sql_calls(adapter) + assert len(sql_calls) == 2 + assert "DROP CLUSTERING KEY" in sql_calls[0] + assert "DROP" in sql_calls[1] and "PARTITION_MONTH" in sql_calls[1] + + +def test_clustering_key_cluster_by_before_column_drop( + make_mocked_engine_adapter: t.Callable, mocker: MockerFixture +) -> None: + """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: 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, + mocker, + current_columns={"ID": "INT", "PARTITION_MONTH": "INT"}, + target_columns={"ID": "INT"}, + current_clustering_key='LINEAR("ID", "PARTITION_MONTH")', + target_clustering_key='LINEAR("ID")', + ) + + operations = adapter.get_alter_operations("catalog.db.current_table", "catalog.db.target_table") + adapter.alter_table(operations) + + assert len(operations) == 3 + assert isinstance(operations[0], TableAlterDropClusterKeyOperation) + assert isinstance(operations[1], TableAlterDropColumnOperation) + assert isinstance(operations[2], TableAlterChangeClusterKeyOperation) + + sql_calls = to_sql_calls(adapter) + 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( + make_mocked_engine_adapter: t.Callable, mocker: MockerFixture +) -> None: + """Case 3: Dropped column is not referenced by the clustering key. + + Current: TENANT_ID, EXTRA — clustered on LINEAR("TENANT_ID"). + Target: TENANT_ID, clustered on LINEAR("TENANT_ID"). + Expected: only DROP COLUMN EXTRA; no clustering statements emitted. + """ + adapter = _make_snowflake_clustering_adapter( + make_mocked_engine_adapter, + mocker, + current_columns={"TENANT_ID": "INT", "EXTRA": "INT"}, + target_columns={"TENANT_ID": "INT"}, + current_clustering_key='LINEAR("TENANT_ID")', + target_clustering_key='LINEAR("TENANT_ID")', + ) + + operations = adapter.get_alter_operations("catalog.db.current_table", "catalog.db.target_table") + + assert len(operations) == 1 + assert isinstance(operations[0], TableAlterDropColumnOperation) + assert not any( + isinstance(op, (TableAlterChangeClusterKeyOperation, TableAlterDropClusterKeyOperation)) + for op in operations + ) + + +def test_clustering_key_add_column_before_cluster_by( + make_mocked_engine_adapter: t.Callable, mocker: MockerFixture +) -> None: + """Case 4: New clustering key references a column being added. + + Current: ID, not clustered. + Target: ID, C — clustered on LINEAR("C"). + Expected order: ADD COLUMN C, then CLUSTER BY ("C"). + """ + adapter = _make_snowflake_clustering_adapter( + make_mocked_engine_adapter, + mocker, + current_columns={"ID": "INT"}, + target_columns={"ID": "INT", "C": "INT"}, + current_clustering_key=None, + 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) == 2 + assert isinstance(operations[0], TableAlterAddColumnOperation) + assert isinstance(operations[1], TableAlterChangeClusterKeyOperation) + + sql_calls = to_sql_calls(adapter) + assert len(sql_calls) == 2 + assert "ADD" in sql_calls[0] and "C" in sql_calls[0] + 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: + """Case 5: No column changes, clustering key removed. + + Current: ID, C — clustered on LINEAR("ID"). + Target: ID, C — not clustered. + Expected: exactly one DROP CLUSTERING KEY, nothing else. + """ + adapter = _make_snowflake_clustering_adapter( + make_mocked_engine_adapter, + mocker, + current_columns={"ID": "INT", "C": "INT"}, + target_columns={"ID": "INT", "C": "INT"}, + current_clustering_key='LINEAR("ID")', + target_clustering_key=None, + ) + + operations = adapter.get_alter_operations("catalog.db.current_table", "catalog.db.target_table") + adapter.alter_table(operations) + + assert len(operations) == 1 + assert isinstance(operations[0], TableAlterDropClusterKeyOperation) + + sql_calls = to_sql_calls(adapter) + assert len(sql_calls) == 1 + assert "DROP CLUSTERING KEY" in sql_calls[0] + + +def test_clustering_key_inserted_between_multiple_drops( + make_mocked_engine_adapter: t.Callable, mocker: MockerFixture +) -> None: + """Case 6: Multiple columns dropped; only one belongs to the clustering key. + + Current: EXTRA, PARTITION_MONTH, TENANT_ID — clustered on LINEAR("PARTITION_MONTH"). + Target: TENANT_ID, not clustered. + Schema diff proposes: DROP EXTRA, DROP PARTITION_MONTH. + Expected order: DROP COLUMN EXTRA, DROP CLUSTERING KEY, DROP COLUMN PARTITION_MONTH. + """ + adapter = _make_snowflake_clustering_adapter( + make_mocked_engine_adapter, + mocker, + current_columns={"EXTRA": "INT", "PARTITION_MONTH": "INT", "TENANT_ID": "INT"}, + target_columns={"TENANT_ID": "INT"}, + current_clustering_key='LINEAR("PARTITION_MONTH")', + target_clustering_key=None, + ) + + operations = adapter.get_alter_operations("catalog.db.current_table", "catalog.db.target_table") + adapter.alter_table(operations) + + assert len(operations) == 3 + assert isinstance(operations[0], TableAlterDropColumnOperation) + assert isinstance(operations[1], TableAlterDropClusterKeyOperation) + assert isinstance(operations[2], TableAlterDropColumnOperation) + + sql_calls = to_sql_calls(adapter) + assert len(sql_calls) == 3 + 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] + + +def test_clustering_key_cluster_by_change_inserted_between_drops( + make_mocked_engine_adapter: t.Callable, mocker: MockerFixture +) -> None: + """Clustering key changes (not just dropped) with multiple column drops. + + Current: EXTRA, PARTITION_MONTH, ID — clustered on LINEAR("PARTITION_MONTH"). + Target: ID, clustered on LINEAR("ID"). + 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, + mocker, + current_columns={"EXTRA": "INT", "PARTITION_MONTH": "INT", "ID": "INT"}, + target_columns={"ID": "INT"}, + current_clustering_key='LINEAR("PARTITION_MONTH")', + target_clustering_key='LINEAR("ID")', + ) + + 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], 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) + + 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( + make_mocked_engine_adapter: t.Callable, mocker: MockerFixture +) -> None: + """All operations must target the current (source) table, not the target table.""" + adapter = _make_snowflake_clustering_adapter( + make_mocked_engine_adapter, + mocker, + current_columns={"TENANT_ID": "INT", "PARTITION_MONTH": "INT"}, + target_columns={"TENANT_ID": "INT"}, + current_clustering_key='LINEAR("TENANT_ID", "PARTITION_MONTH")', + target_clustering_key=None, + ) + + operations = adapter.get_alter_operations("catalog.db.current_table", "catalog.db.target_table") + + for op in operations: + assert "current_table" in op.target_table.sql(dialect="snowflake").lower() + + +def test_clustering_key_normalization_quoted_upper_matches_upper_column( + make_mocked_engine_adapter: t.Callable, mocker: MockerFixture +) -> None: + """LINEAR("PARTITION_MONTH") references the column string PARTITION_MONTH. + + Both the quoted key identifier and the column name normalize to PARTITION_MONTH + in Snowflake, so the reorder must fire. + """ + adapter = _make_snowflake_clustering_adapter( + make_mocked_engine_adapter, + mocker, + current_columns={"TENANT_ID": "INT", "PARTITION_MONTH": "INT"}, + target_columns={"TENANT_ID": "INT"}, + current_clustering_key='LINEAR("PARTITION_MONTH")', + target_clustering_key=None, + ) + + operations = adapter.get_alter_operations("catalog.db.current_table", "catalog.db.target_table") + + assert isinstance(operations[0], TableAlterDropClusterKeyOperation) + assert isinstance(operations[1], TableAlterDropColumnOperation) + + +def test_clustering_key_normalization_quoted_upper_matches_lower_column( + make_mocked_engine_adapter: t.Callable, mocker: MockerFixture +) -> None: + """LINEAR("PARTITION_MONTH") also references the column string partition_month. + + The unquoted column string partition_month normalizes to PARTITION_MONTH in + Snowflake, matching the quoted-upper key identifier. + """ + adapter = _make_snowflake_clustering_adapter( + make_mocked_engine_adapter, + mocker, + current_columns={"TENANT_ID": "INT", "partition_month": "INT"}, + target_columns={"TENANT_ID": "INT"}, + current_clustering_key='LINEAR("PARTITION_MONTH")', + target_clustering_key=None, + ) + + operations = adapter.get_alter_operations("catalog.db.current_table", "catalog.db.target_table") + + assert isinstance(operations[0], TableAlterDropClusterKeyOperation) + assert isinstance(operations[1], TableAlterDropColumnOperation) + + +def test_clustering_key_normalization_quoted_lower_no_match_upper_column( + make_mocked_engine_adapter: t.Callable, mocker: MockerFixture +) -> None: + """LINEAR("partition_month") does NOT reference the column PARTITION_MONTH. + + The quoted-lowercase identifier stays lowercase, while the column string + PARTITION_MONTH normalizes to PARTITION_MONTH. No match → no reorder. + """ + adapter = _make_snowflake_clustering_adapter( + make_mocked_engine_adapter, + mocker, + current_columns={"TENANT_ID": "INT", "PARTITION_MONTH": "INT"}, + target_columns={"TENANT_ID": "INT"}, + current_clustering_key='LINEAR("partition_month")', + target_clustering_key='LINEAR("partition_month")', + ) + + operations = adapter.get_alter_operations("catalog.db.current_table", "catalog.db.target_table") + + assert len(operations) == 1 + assert isinstance(operations[0], TableAlterDropColumnOperation) + assert not any( + isinstance(op, (TableAlterChangeClusterKeyOperation, TableAlterDropClusterKeyOperation)) + for op in operations + ) + + +def test_clustering_key_normalization_quoted_lower_no_match_lower_column( + make_mocked_engine_adapter: t.Callable, mocker: MockerFixture +) -> None: + """LINEAR("partition_month") does NOT reference the column string partition_month. + + The column string partition_month normalizes to PARTITION_MONTH in Snowflake, + but the quoted-lowercase key identifier stays lowercase, so they do not match. + """ + adapter = _make_snowflake_clustering_adapter( + make_mocked_engine_adapter, + mocker, + current_columns={"TENANT_ID": "INT", "partition_month": "INT"}, + target_columns={"TENANT_ID": "INT"}, + current_clustering_key='LINEAR("partition_month")', + target_clustering_key='LINEAR("partition_month")', + ) + + operations = adapter.get_alter_operations("catalog.db.current_table", "catalog.db.target_table") + + assert len(operations) == 1 + assert isinstance(operations[0], TableAlterDropColumnOperation) diff --git a/tests/core/engine_adapter/test_starrocks.py b/tests/core/engine_adapter/test_starrocks.py index db0b1cc4ae..76f7a8b695 100644 --- a/tests/core/engine_adapter/test_starrocks.py +++ b/tests/core/engine_adapter/test_starrocks.py @@ -23,7 +23,9 @@ from sqlglot import expressions as exp from sqlglot import parse_one from pytest_mock.plugin import MockerFixture -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.utils.errors import SQLMeshError from tests.core.engine_adapter import to_sql_calls @@ -2486,3 +2488,60 @@ def test_non_starrocks_engine_is_not_affected(self) -> None: "excluded_refresh_tables" in StarRocksEngineAdapter.RESOLVE_TABLE_REFS_IN_PHYSICAL_PROPERTIES ) + + +class TestClusteringKeyOrdering: + """StarRocks uses ClusteredByMixin but must KEEP the historical alter order. + + Dropping a column that belongs to the current clustering key must NOT trigger + a reorder. StarRocks must emit DROP COLUMN first, then the clustering change. + """ + + def test_column_drop_before_clustering_key_change( + self, + make_mocked_engine_adapter: t.Callable[..., StarRocksEngineAdapter], + mocker: MockerFixture, + ) -> None: + """DROP COLUMN is emitted before DROP CLUSTERING KEY (historical order).""" + adapter = make_mocked_engine_adapter(StarRocksEngineAdapter, patch_get_data_objects=False) + + current_columns = {"id": "INT", "partition_month": "INT"} + target_columns = {"id": "INT"} + + def _columns(table_name: t.Any, **kwargs: t.Any) -> t.Dict[str, exp.DataType]: + mapping = current_columns if table_name == "db.current_table" else target_columns + return {k: exp.DataType.build(v) for k, v in mapping.items()} + + mocker.patch.object(adapter, "columns", side_effect=_columns) + + current_obj = DataObject( + catalog=None, + schema="db", + name="current_table", + type=DataObjectType.TABLE, + clustering_key='LINEAR("partition_month")', + ) + target_obj = DataObject( + catalog=None, + schema="db", + name="target_table", + type=DataObjectType.TABLE, + clustering_key=None, + ) + + def _mock_get_data_objects( + 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=_mock_get_data_objects) + + operations = adapter.get_alter_operations("db.current_table", "db.target_table") + + assert len(operations) == 2 + assert isinstance(operations[0], TableAlterDropColumnOperation) + assert isinstance(operations[1], TableAlterDropClusterKeyOperation)