Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 14 additions & 7 deletions sqlmesh/core/engine_adapter/mixins.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
41 changes: 40 additions & 1 deletion tests/core/engine_adapter/test_snowflake.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
):
Expand Down