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
14 changes: 13 additions & 1 deletion sqlmesh/core/engine_adapter/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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)}")
Expand All @@ -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(
Expand Down Expand Up @@ -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 [
Expand Down
1 change: 1 addition & 0 deletions sqlmesh/core/engine_adapter/bigquery.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
1 change: 1 addition & 0 deletions sqlmesh/core/engine_adapter/clickhouse.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
4 changes: 4 additions & 0 deletions sqlmesh/core/engine_adapter/databricks.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 {}
Expand All @@ -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,
)

Expand Down
4 changes: 3 additions & 1 deletion sqlmesh/core/engine_adapter/fabric.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
32 changes: 32 additions & 0 deletions sqlmesh/core/engine_adapter/snowflake.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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,
Expand Down
3 changes: 2 additions & 1 deletion sqlmesh/core/snapshot/evaluator.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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")
Expand Down
42 changes: 42 additions & 0 deletions tests/core/engine_adapter/test_snowflake.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
):
Expand Down
6 changes: 4 additions & 2 deletions tests/core/test_snapshot_evaluator.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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"
Expand Down Expand Up @@ -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(
Expand All @@ -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(
[
Expand Down