Skip to content

Commit 2f0404c

Browse files
committed
Fix: Use Iceberg-specific DDL to clone and alter Snowflake Iceberg tables
Snowflake rejects `CREATE TABLE ... CLONE` and `ALTER TABLE` for Iceberg tables, requiring `CREATE ICEBERG TABLE ... CLONE` and `ALTER ICEBERG TABLE` instead. The model's `table_format` was already honoured when creating tables but was never propagated to the clone and alter code paths, so both failed with a SQL compilation error during the virtual layer update and schema migration respectively. `clone_table` now accepts `table_format`/`table_kind` and `alter_table` accepts `table_format`, mirroring the existing `_create_table` convention. The Snowflake adapter derives the Iceberg-specific table kind from the format, and the evaluator passes the model's table format through both paths. Fixes #5721 Signed-off-by: Guillem G <guillem.gimenez@titanos.tv>
1 parent 24bb095 commit 2f0404c

9 files changed

Lines changed: 102 additions & 5 deletions

File tree

sqlmesh/core/engine_adapter/base.py

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1094,6 +1094,8 @@ def clone_table(
10941094
replace: bool = False,
10951095
exists: bool = True,
10961096
clone_kwargs: t.Optional[t.Dict[str, t.Any]] = None,
1097+
table_format: t.Optional[str] = None,
1098+
table_kind: t.Optional[str] = None,
10971099
**kwargs: t.Any,
10981100
) -> None:
10991101
"""Creates a table with the target name by cloning the source table.
@@ -1103,6 +1105,10 @@ def clone_table(
11031105
source_table_name: The name of the source table that should be cloned.
11041106
replace: Whether or not to replace an existing table.
11051107
exists: Indicates whether to include the IF NOT EXISTS check.
1108+
clone_kwargs: Additional arguments for the CLONE clause.
1109+
table_format: The table format of the source table, if any. Engines that require
1110+
format-specific DDL to clone a table use it to derive `table_kind`.
1111+
table_kind: The kind of table to create. Defaults to `TABLE`.
11061112
"""
11071113
if not self.SUPPORTS_CLONING:
11081114
raise NotImplementedError(f"Engine does not support cloning: {type(self)}")
@@ -1111,7 +1117,7 @@ def clone_table(
11111117
self.execute(
11121118
exp.Create(
11131119
this=exp.to_table(target_table_name),
1114-
kind="TABLE",
1120+
kind=table_kind or "TABLE",
11151121
replace=replace,
11161122
exists=exists,
11171123
clone=exp.Clone(
@@ -1214,9 +1220,15 @@ def get_alter_operations(
12141220
def alter_table(
12151221
self,
12161222
alter_expressions: t.Union[t.List[exp.Alter], t.List[TableAlterOperation]],
1223+
table_format: t.Optional[str] = None,
12171224
) -> None:
12181225
"""
12191226
Performs the alter statements to change the current table into the structure of the target table.
1227+
1228+
Args:
1229+
alter_expressions: The alter operations to apply.
1230+
table_format: The table format of the target table, if any. Engines that require
1231+
format-specific DDL to alter a table use it to adjust the generated statements.
12201232
"""
12211233
with self.transaction():
12221234
for alter_expression in [

sqlmesh/core/engine_adapter/bigquery.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -405,6 +405,7 @@ def create_mapping_schema(
405405
def alter_table(
406406
self,
407407
alter_expressions: t.Union[t.List[exp.Alter], t.List[TableAlterOperation]],
408+
table_format: t.Optional[str] = None,
408409
) -> None:
409410
"""
410411
Performs the alter statements to change the current table into the structure of the target table,

sqlmesh/core/engine_adapter/clickhouse.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -697,6 +697,7 @@ def delete_from(self, table_name: TableName, where: t.Union[str, exp.Expr]) -> N
697697
def alter_table(
698698
self,
699699
alter_expressions: t.Union[t.List[exp.Alter], t.List[TableAlterOperation]],
700+
table_format: t.Optional[str] = None,
700701
) -> None:
701702
"""
702703
Performs the alter statements to change the current table into the structure of the target table.

sqlmesh/core/engine_adapter/databricks.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -386,6 +386,8 @@ def clone_table(
386386
replace: bool = False,
387387
exists: bool = True,
388388
clone_kwargs: t.Optional[t.Dict[str, t.Any]] = None,
389+
table_format: t.Optional[str] = None,
390+
table_kind: t.Optional[str] = None,
389391
**kwargs: t.Any,
390392
) -> None:
391393
clone_kwargs = clone_kwargs or {}
@@ -395,6 +397,8 @@ def clone_table(
395397
source_table_name,
396398
replace=replace,
397399
clone_kwargs=clone_kwargs,
400+
table_format=table_format,
401+
table_kind=table_kind,
398402
**kwargs,
399403
)
400404

sqlmesh/core/engine_adapter/fabric.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -225,7 +225,9 @@ def set_current_catalog(self, catalog_name: t.Optional[str]) -> None:
225225
self._target_catalog = target_catalog
226226

227227
def alter_table(
228-
self, alter_expressions: t.Union[t.List[exp.Alter], t.List[TableAlterOperation]]
228+
self,
229+
alter_expressions: t.Union[t.List[exp.Alter], t.List[TableAlterOperation]],
230+
table_format: t.Optional[str] = None,
229231
) -> None:
230232
"""
231233
Applies alter expressions to a table. Fabric has limited support for ALTER TABLE,

sqlmesh/core/engine_adapter/snowflake.py

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
SourceQuery,
2525
set_catalog,
2626
)
27+
from sqlmesh.core.schema_diff import TableAlterOperation
2728
from sqlmesh.utils import optional_import, get_source_columns_to_types
2829
from sqlmesh.utils.errors import SQLMeshError
2930
from sqlmesh.utils.pandas import columns_to_types_from_dtypes
@@ -667,6 +668,8 @@ def clone_table(
667668
replace: bool = False,
668669
exists: bool = True,
669670
clone_kwargs: t.Optional[t.Dict[str, t.Any]] = None,
671+
table_format: t.Optional[str] = None,
672+
table_kind: t.Optional[str] = None,
670673
**kwargs: t.Any,
671674
) -> None:
672675
# The Snowflake adapter should use the transient property to clone transient tables
@@ -675,14 +678,43 @@ def clone_table(
675678
if isinstance(table_type, exp.TransientProperty):
676679
kwargs["properties"] = exp.Properties(expressions=[table_type])
677680

681+
# Snowflake rejects `CREATE TABLE ... CLONE` for Iceberg tables, it requires
682+
# `CREATE ICEBERG TABLE ... CLONE` instead
683+
if table_format and not table_kind:
684+
table_kind = f"{table_format.upper()} TABLE"
685+
678686
super().clone_table(
679687
target_table_name,
680688
source_table_name,
681689
replace=replace,
682690
clone_kwargs=clone_kwargs,
691+
table_kind=table_kind,
683692
**kwargs,
684693
)
685694

695+
def alter_table(
696+
self,
697+
alter_expressions: t.Union[t.List[exp.Alter], t.List[TableAlterOperation]],
698+
table_format: t.Optional[str] = None,
699+
) -> None:
700+
# Snowflake rejects `ALTER TABLE` for Iceberg tables, it requires
701+
# `ALTER ICEBERG TABLE` instead
702+
if table_format:
703+
table_kind = f"{table_format.upper()} TABLE"
704+
resolved_expressions = []
705+
for alter_expression in alter_expressions:
706+
resolved_expression = (
707+
alter_expression.expression
708+
if isinstance(alter_expression, TableAlterOperation)
709+
else alter_expression.copy()
710+
)
711+
resolved_expression.set("kind", table_kind)
712+
resolved_expressions.append(resolved_expression)
713+
714+
super().alter_table(resolved_expressions)
715+
else:
716+
super().alter_table(alter_expressions)
717+
686718
@t.overload
687719
def _columns_to_types(
688720
self,

sqlmesh/core/snapshot/evaluator.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1104,6 +1104,7 @@ def _clone_snapshot_in_dev(
11041104
target_table_name,
11051105
snapshot.table_name(),
11061106
rendered_physical_properties=rendered_physical_properties,
1107+
table_format=snapshot.model.table_format,
11071108
)
11081109
self._migrate_target_table(
11091110
target_table_name=target_table_name,
@@ -2161,7 +2162,7 @@ def migrate(
21612162
_check_additive_schema_change(
21622163
snapshot, alter_operations, kwargs["allow_additive_snapshots"]
21632164
)
2164-
self.adapter.alter_table(alter_operations)
2165+
self.adapter.alter_table(alter_operations, table_format=snapshot.model.table_format)
21652166

21662167
# Apply grants after schema migration
21672168
deployability_index = kwargs.get("deployability_index")

tests/core/engine_adapter/test_snowflake.py

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1007,6 +1007,48 @@ def test_table_format_iceberg(snowflake_mocked_engine_adapter: SnowflakeEngineAd
10071007
]
10081008

10091009

1010+
def test_clone_table_iceberg(mocker: MockerFixture, make_mocked_engine_adapter: t.Callable):
1011+
mocker.patch("sqlmesh.core.engine_adapter.snowflake.SnowflakeEngineAdapter.set_current_catalog")
1012+
adapter = make_mocked_engine_adapter(SnowflakeEngineAdapter, default_catalog="test_catalog")
1013+
1014+
# Snowflake rejects `CREATE TABLE ... CLONE` for Iceberg tables
1015+
adapter.clone_table("target_table", "source_table", table_format="iceberg")
1016+
adapter.cursor.execute.assert_called_once_with(
1017+
'CREATE ICEBERG TABLE IF NOT EXISTS "target_table" CLONE "source_table"'
1018+
)
1019+
1020+
# Engines that don't need format-specific DDL are unaffected
1021+
adapter = make_mocked_engine_adapter(EngineAdapter, default_catalog="test_catalog")
1022+
adapter.SUPPORTS_CLONING = True
1023+
adapter.clone_table("target_table", "source_table", table_format="iceberg")
1024+
adapter.cursor.execute.assert_called_once_with(
1025+
'CREATE TABLE IF NOT EXISTS "target_table" CLONE "source_table"'
1026+
)
1027+
1028+
1029+
def test_alter_table_iceberg(mocker: MockerFixture, make_mocked_engine_adapter: t.Callable):
1030+
mocker.patch("sqlmesh.core.engine_adapter.snowflake.SnowflakeEngineAdapter.set_current_catalog")
1031+
adapter = make_mocked_engine_adapter(SnowflakeEngineAdapter, default_catalog="test_catalog")
1032+
1033+
current_table = {"a": "INT"}
1034+
target_table = {"a": "INT", "b": "INT"}
1035+
adapter.columns = lambda table_name, **kwargs: { # type: ignore[assignment]
1036+
k: exp.DataType.build(v)
1037+
for k, v in (current_table if table_name == "test_table" else target_table).items()
1038+
}
1039+
1040+
alter_operations = adapter.get_alter_operations("test_table", "target_table")
1041+
1042+
# Snowflake rejects `ALTER TABLE` for Iceberg tables
1043+
adapter.alter_table(alter_operations, table_format="iceberg")
1044+
assert to_sql_calls(adapter) == ['ALTER ICEBERG TABLE "test_table" ADD "b" INT']
1045+
1046+
# Without a table format the regular `ALTER TABLE` is used
1047+
adapter = make_mocked_engine_adapter(SnowflakeEngineAdapter, default_catalog="test_catalog")
1048+
adapter.alter_table(alter_operations)
1049+
assert to_sql_calls(adapter) == ['ALTER TABLE "test_table" ADD "b" INT']
1050+
1051+
10101052
def test_create_view_with_schema_and_grants(
10111053
snowflake_mocked_engine_adapter: SnowflakeEngineAdapter,
10121054
):

tests/core/test_snapshot_evaluator.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1943,6 +1943,7 @@ def test_create_clone_in_dev(mocker: MockerFixture, adapter_mock, make_snapshot)
19431943
f"sqlmesh__test_schema.test_schema__test_model__{snapshot.dev_version}__dev",
19441944
f"sqlmesh__test_schema.test_schema__test_model__{snapshot.version}",
19451945
rendered_physical_properties={},
1946+
table_format=None,
19461947
)
19471948

19481949
adapter_mock.get_alter_operations.assert_called_once_with(
@@ -1952,7 +1953,7 @@ def test_create_clone_in_dev(mocker: MockerFixture, adapter_mock, make_snapshot)
19521953
ignore_additive=False,
19531954
)
19541955

1955-
adapter_mock.alter_table.assert_called_once_with([])
1956+
adapter_mock.alter_table.assert_called_once_with([], table_format=None)
19561957

19571958
adapter_mock.drop_table.assert_called_once_with(
19581959
f"sqlmesh__test_schema.test_schema__test_model__{snapshot.version}__dev_schema_tmp"
@@ -1992,6 +1993,7 @@ def test_drop_clone_in_dev_when_migration_fails(mocker: MockerFixture, adapter_m
19921993
f"sqlmesh__test_schema.test_schema__test_model__{snapshot.version}__dev",
19931994
f"sqlmesh__test_schema.test_schema__test_model__{snapshot.version}",
19941995
rendered_physical_properties={},
1996+
table_format=None,
19951997
)
19961998

19971999
adapter_mock.get_alter_operations.assert_called_once_with(
@@ -2001,7 +2003,7 @@ def test_drop_clone_in_dev_when_migration_fails(mocker: MockerFixture, adapter_m
20012003
ignore_additive=False,
20022004
)
20032005

2004-
adapter_mock.alter_table.assert_called_once_with([])
2006+
adapter_mock.alter_table.assert_called_once_with([], table_format=None)
20052007

20062008
adapter_mock.drop_table.assert_has_calls(
20072009
[

0 commit comments

Comments
 (0)