diff --git a/sqlmesh/core/engine_adapter/base.py b/sqlmesh/core/engine_adapter/base.py index bd435db76f..930cdf7cd4 100644 --- a/sqlmesh/core/engine_adapter/base.py +++ b/sqlmesh/core/engine_adapter/base.py @@ -1094,6 +1094,8 @@ def clone_table( replace: bool = False, exists: bool = True, clone_kwargs: t.Optional[t.Dict[str, t.Any]] = None, + table_format: t.Optional[str] = None, + table_kind: t.Optional[str] = None, **kwargs: t.Any, ) -> None: """Creates a table with the target name by cloning the source table. @@ -1103,6 +1105,10 @@ def clone_table( source_table_name: The name of the source table that should be cloned. replace: Whether or not to replace an existing table. exists: Indicates whether to include the IF NOT EXISTS check. + clone_kwargs: Additional arguments for the CLONE clause. + table_format: The table format of the source table, if any. Engines that require + format-specific DDL to clone a table use it to derive `table_kind`. + table_kind: The kind of table to create. Defaults to `TABLE`. """ if not self.SUPPORTS_CLONING: raise NotImplementedError(f"Engine does not support cloning: {type(self)}") @@ -1111,7 +1117,7 @@ def clone_table( self.execute( exp.Create( this=exp.to_table(target_table_name), - kind="TABLE", + kind=table_kind or "TABLE", replace=replace, exists=exists, clone=exp.Clone( @@ -1214,9 +1220,15 @@ def get_alter_operations( def alter_table( self, alter_expressions: t.Union[t.List[exp.Alter], t.List[TableAlterOperation]], + table_format: t.Optional[str] = None, ) -> None: """ Performs the alter statements to change the current table into the structure of the target table. + + Args: + alter_expressions: The alter operations to apply. + table_format: The table format of the target table, if any. Engines that require + format-specific DDL to alter a table use it to adjust the generated statements. """ with self.transaction(): for alter_expression in [ diff --git a/sqlmesh/core/engine_adapter/bigquery.py b/sqlmesh/core/engine_adapter/bigquery.py index d136445114..7917aaa149 100644 --- a/sqlmesh/core/engine_adapter/bigquery.py +++ b/sqlmesh/core/engine_adapter/bigquery.py @@ -405,6 +405,7 @@ def create_mapping_schema( def alter_table( self, alter_expressions: t.Union[t.List[exp.Alter], t.List[TableAlterOperation]], + table_format: t.Optional[str] = None, ) -> None: """ Performs the alter statements to change the current table into the structure of the target table, diff --git a/sqlmesh/core/engine_adapter/clickhouse.py b/sqlmesh/core/engine_adapter/clickhouse.py index d1f67e0564..0a811dba83 100644 --- a/sqlmesh/core/engine_adapter/clickhouse.py +++ b/sqlmesh/core/engine_adapter/clickhouse.py @@ -697,6 +697,7 @@ def delete_from(self, table_name: TableName, where: t.Union[str, exp.Expr]) -> N def alter_table( self, alter_expressions: t.Union[t.List[exp.Alter], t.List[TableAlterOperation]], + table_format: t.Optional[str] = None, ) -> None: """ Performs the alter statements to change the current table into the structure of the target table. diff --git a/sqlmesh/core/engine_adapter/databricks.py b/sqlmesh/core/engine_adapter/databricks.py index 098825bc2d..8d47cdfd61 100644 --- a/sqlmesh/core/engine_adapter/databricks.py +++ b/sqlmesh/core/engine_adapter/databricks.py @@ -386,6 +386,8 @@ def clone_table( replace: bool = False, exists: bool = True, clone_kwargs: t.Optional[t.Dict[str, t.Any]] = None, + table_format: t.Optional[str] = None, + table_kind: t.Optional[str] = None, **kwargs: t.Any, ) -> None: clone_kwargs = clone_kwargs or {} @@ -395,6 +397,8 @@ def clone_table( source_table_name, replace=replace, clone_kwargs=clone_kwargs, + table_format=table_format, + table_kind=table_kind, **kwargs, ) diff --git a/sqlmesh/core/engine_adapter/fabric.py b/sqlmesh/core/engine_adapter/fabric.py index 7b2f1acd73..427a04d087 100644 --- a/sqlmesh/core/engine_adapter/fabric.py +++ b/sqlmesh/core/engine_adapter/fabric.py @@ -225,7 +225,9 @@ def set_current_catalog(self, catalog_name: t.Optional[str]) -> None: self._target_catalog = target_catalog def alter_table( - self, alter_expressions: t.Union[t.List[exp.Alter], t.List[TableAlterOperation]] + self, + alter_expressions: t.Union[t.List[exp.Alter], t.List[TableAlterOperation]], + table_format: t.Optional[str] = None, ) -> None: """ Applies alter expressions to a table. Fabric has limited support for ALTER TABLE, diff --git a/sqlmesh/core/engine_adapter/snowflake.py b/sqlmesh/core/engine_adapter/snowflake.py index d589b5d15b..37ca6eee22 100644 --- a/sqlmesh/core/engine_adapter/snowflake.py +++ b/sqlmesh/core/engine_adapter/snowflake.py @@ -24,6 +24,7 @@ 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 @@ -667,6 +668,8 @@ def clone_table( replace: bool = False, exists: bool = True, clone_kwargs: t.Optional[t.Dict[str, t.Any]] = None, + table_format: t.Optional[str] = None, + table_kind: t.Optional[str] = None, **kwargs: t.Any, ) -> None: # The Snowflake adapter should use the transient property to clone transient tables @@ -675,14 +678,43 @@ def clone_table( if isinstance(table_type, exp.TransientProperty): kwargs["properties"] = exp.Properties(expressions=[table_type]) + # Snowflake rejects `CREATE TABLE ... CLONE` for Iceberg tables, it requires + # `CREATE ICEBERG TABLE ... CLONE` instead + if table_format and not table_kind: + table_kind = f"{table_format.upper()} TABLE" + super().clone_table( target_table_name, source_table_name, replace=replace, clone_kwargs=clone_kwargs, + table_kind=table_kind, **kwargs, ) + def alter_table( + self, + alter_expressions: t.Union[t.List[exp.Alter], t.List[TableAlterOperation]], + table_format: t.Optional[str] = None, + ) -> None: + # Snowflake rejects `ALTER TABLE` for Iceberg tables, it requires + # `ALTER ICEBERG TABLE` instead + if table_format: + table_kind = f"{table_format.upper()} TABLE" + resolved_expressions = [] + for alter_expression in alter_expressions: + resolved_expression = ( + alter_expression.expression + if isinstance(alter_expression, TableAlterOperation) + else alter_expression.copy() + ) + resolved_expression.set("kind", table_kind) + resolved_expressions.append(resolved_expression) + + super().alter_table(resolved_expressions) + else: + super().alter_table(alter_expressions) + @t.overload def _columns_to_types( self, diff --git a/sqlmesh/core/snapshot/evaluator.py b/sqlmesh/core/snapshot/evaluator.py index 11b3fd1f33..ad935310f4 100644 --- a/sqlmesh/core/snapshot/evaluator.py +++ b/sqlmesh/core/snapshot/evaluator.py @@ -1104,6 +1104,7 @@ def _clone_snapshot_in_dev( target_table_name, snapshot.table_name(), rendered_physical_properties=rendered_physical_properties, + table_format=snapshot.model.table_format, ) self._migrate_target_table( target_table_name=target_table_name, @@ -2161,7 +2162,7 @@ def migrate( _check_additive_schema_change( snapshot, alter_operations, kwargs["allow_additive_snapshots"] ) - self.adapter.alter_table(alter_operations) + self.adapter.alter_table(alter_operations, table_format=snapshot.model.table_format) # Apply grants after schema migration deployability_index = kwargs.get("deployability_index") diff --git a/tests/core/engine_adapter/test_snowflake.py b/tests/core/engine_adapter/test_snowflake.py index 085c51098b..05f0133ec8 100644 --- a/tests/core/engine_adapter/test_snowflake.py +++ b/tests/core/engine_adapter/test_snowflake.py @@ -1007,6 +1007,48 @@ def test_table_format_iceberg(snowflake_mocked_engine_adapter: SnowflakeEngineAd ] +def test_clone_table_iceberg(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") + + # Snowflake rejects `CREATE TABLE ... CLONE` for Iceberg tables + adapter.clone_table("target_table", "source_table", table_format="iceberg") + adapter.cursor.execute.assert_called_once_with( + 'CREATE ICEBERG TABLE IF NOT EXISTS "target_table" CLONE "source_table"' + ) + + # Engines that don't need format-specific DDL are unaffected + adapter = make_mocked_engine_adapter(EngineAdapter, default_catalog="test_catalog") + adapter.SUPPORTS_CLONING = True + adapter.clone_table("target_table", "source_table", table_format="iceberg") + adapter.cursor.execute.assert_called_once_with( + 'CREATE TABLE IF NOT EXISTS "target_table" CLONE "source_table"' + ) + + +def test_alter_table_iceberg(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"} + target_table = {"a": "INT", "b": "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() + } + + alter_operations = adapter.get_alter_operations("test_table", "target_table") + + # Snowflake rejects `ALTER TABLE` for Iceberg tables + adapter.alter_table(alter_operations, table_format="iceberg") + assert to_sql_calls(adapter) == ['ALTER ICEBERG TABLE "test_table" ADD "b" INT'] + + # Without a table format the regular `ALTER TABLE` is used + adapter = make_mocked_engine_adapter(SnowflakeEngineAdapter, default_catalog="test_catalog") + adapter.alter_table(alter_operations) + assert to_sql_calls(adapter) == ['ALTER TABLE "test_table" ADD "b" INT'] + + def test_create_view_with_schema_and_grants( snowflake_mocked_engine_adapter: SnowflakeEngineAdapter, ): diff --git a/tests/core/test_snapshot_evaluator.py b/tests/core/test_snapshot_evaluator.py index 27bcbe05ae..3950e0d53f 100644 --- a/tests/core/test_snapshot_evaluator.py +++ b/tests/core/test_snapshot_evaluator.py @@ -1943,6 +1943,7 @@ def test_create_clone_in_dev(mocker: MockerFixture, adapter_mock, make_snapshot) f"sqlmesh__test_schema.test_schema__test_model__{snapshot.dev_version}__dev", f"sqlmesh__test_schema.test_schema__test_model__{snapshot.version}", rendered_physical_properties={}, + table_format=None, ) 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) ignore_additive=False, ) - adapter_mock.alter_table.assert_called_once_with([]) + adapter_mock.alter_table.assert_called_once_with([], table_format=None) adapter_mock.drop_table.assert_called_once_with( 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 f"sqlmesh__test_schema.test_schema__test_model__{snapshot.version}__dev", f"sqlmesh__test_schema.test_schema__test_model__{snapshot.version}", rendered_physical_properties={}, + table_format=None, ) 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 ignore_additive=False, ) - adapter_mock.alter_table.assert_called_once_with([]) + adapter_mock.alter_table.assert_called_once_with([], table_format=None) adapter_mock.drop_table.assert_has_calls( [