Merge branch 'main' into fix/snowflake-iceberg-column-comments

Author: sineline

sqlmesh/core/engine_adapter/base.py

@@ -1098,6 +1098,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.
@@ -1107,6 +1109,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)}")
@@ -1115,7 +1121,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(
@@ -1218,9 +1224,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 [

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,

sqlmesh/core/engine_adapter/clickhouse.py

@@ -699,6 +699,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.

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,
         )
 

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,

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
@@ -674,6 +675,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
@@ -682,14 +685,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,

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")

tests/core/engine_adapter/integration/docker/_common-hive.yaml

@@ -12,22 +12,22 @@ services:
 
   # S3-style object storage
   minio:
-    image: 'minio/minio:RELEASE.2022-05-26T05-48-41Z'
+    image: 'cgr.dev/chainguard/minio:latest@sha256:039800e64ec7247d2fde7cff3697e964f6fe20b6d7d2c46aa7d82cc63355d512'
     ports:
       - '9000:9000'
       - '9001:9001'
     environment:
-      MINIO_ACCESS_KEY: minio
-      MINIO_SECRET_KEY: minio123
+      MINIO_ROOT_USER: minio
+      MINIO_ROOT_PASSWORD: minio123
     command: server /data --console-address ":9001"
 
   # Set up minio with default buckets / paths
   mc-job:
-    image: 'minio/mc:RELEASE.2022-05-09T04-08-26Z'
+    image: 'cgr.dev/chainguard/minio-client:latest-dev@sha256:fb635b967f5f32424391150dac4985cc1aea78fb0a1ce031c01952cf661a13ec'
     entrypoint: |
       /bin/bash -c "
       sleep 5;
-      /usr/bin/mc config --quiet host add myminio http://minio:9000 minio minio123;
+      /usr/bin/mc alias --quiet set myminio http://minio:9000 minio minio123;
       /usr/bin/mc mb --quiet myminio/trino/datalake;
       /usr/bin/mc mb --quiet myminio/trino/datalake_iceberg;
       /usr/bin/mc mb --quiet myminio/trino/datalake_delta;
@@ -39,4 +39,4 @@ services:
       /usr/bin/mc mb --quiet myminio/nessie/warehouse;
       "
     depends_on:
-      - minio
+      - minio

tests/core/engine_adapter/test_snowflake.py

@@ -1039,6 +1039,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,
 ):

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(
         [