diff --git a/sqlmesh/core/engine_adapter/base.py b/sqlmesh/core/engine_adapter/base.py index bd435db76f..fd3d9a5c5c 100644 --- a/sqlmesh/core/engine_adapter/base.py +++ b/sqlmesh/core/engine_adapter/base.py @@ -833,7 +833,9 @@ def _create_table_from_columns( and self.COMMENT_CREATION_TABLE.is_comment_command_only and self.comments_enabled ): - self._create_column_comments(table_name, column_descriptions) + self._create_column_comments( + table_name, column_descriptions, table_format=kwargs.get("table_format") + ) def _build_schema_exp( self, @@ -990,7 +992,9 @@ def _create_table_from_source_queries( ): self._create_table_comment(table_name, table_description) if column_descriptions and schema is None and self.comments_enabled: - self._create_column_comments(table_name, column_descriptions) + self._create_column_comments( + table_name, column_descriptions, table_format=kwargs.get("table_format") + ) def _create_table( self, @@ -3058,7 +3062,18 @@ def _create_column_comments( column_comments: t.Dict[str, str], table_kind: str = "TABLE", materialized_view: bool = False, + table_format: t.Optional[str] = None, ) -> None: + """Registers column comments with a post-creation command. + + Args: + table_name: The name of the table or view. + column_comments: Mapping between the column name and its comment. + table_kind: The kind of object being commented on, `TABLE` or `VIEW`. + materialized_view: Whether the view is materialized. + table_format: The table format of the table, if any. Engines that require + format-specific DDL to alter a table use it to derive `table_kind`. + """ table = exp.to_table(table_name) for col, comment in column_comments.items(): diff --git a/sqlmesh/core/engine_adapter/bigquery.py b/sqlmesh/core/engine_adapter/bigquery.py index d136445114..8139e799a7 100644 --- a/sqlmesh/core/engine_adapter/bigquery.py +++ b/sqlmesh/core/engine_adapter/bigquery.py @@ -818,6 +818,7 @@ def _create_column_comments( column_comments: t.Dict[str, str], table_kind: str = "TABLE", materialized_view: bool = False, + table_format: t.Optional[str] = None, ) -> None: if not (table_kind == "VIEW" and materialized_view): table = self._get_table(table_name) diff --git a/sqlmesh/core/engine_adapter/mysql.py b/sqlmesh/core/engine_adapter/mysql.py index 6918cdec49..a20ba2af26 100644 --- a/sqlmesh/core/engine_adapter/mysql.py +++ b/sqlmesh/core/engine_adapter/mysql.py @@ -131,6 +131,7 @@ def _create_column_comments( column_comments: t.Dict[str, str], table_kind: str = "TABLE", materialized_view: bool = False, + table_format: t.Optional[str] = None, ) -> None: table = exp.to_table(table_name) table_sql = table.sql(dialect=self.dialect, identify=True) diff --git a/sqlmesh/core/engine_adapter/snowflake.py b/sqlmesh/core/engine_adapter/snowflake.py index d589b5d15b..885e1e7c1f 100644 --- a/sqlmesh/core/engine_adapter/snowflake.py +++ b/sqlmesh/core/engine_adapter/snowflake.py @@ -632,13 +632,20 @@ def _create_column_comments( column_comments: t.Dict[str, str], table_kind: str = "TABLE", materialized_view: bool = False, + table_format: t.Optional[str] = None, ) -> None: """ Reference: https://docs.snowflake.com/en/sql-reference/sql/alter-table-column#syntax + Reference: https://docs.snowflake.com/en/sql-reference/sql/alter-iceberg-table#syntax """ if not column_comments: return + # Snowflake rejects `ALTER TABLE` for Iceberg tables, it requires + # `ALTER ICEBERG TABLE` instead + if table_format and table_kind == "TABLE": + table_kind = f"{table_format.upper()} TABLE" + table = exp.to_table(table_name) table_sql = self._to_sql(table) diff --git a/tests/core/engine_adapter/test_snowflake.py b/tests/core/engine_adapter/test_snowflake.py index 085c51098b..1e772d8d29 100644 --- a/tests/core/engine_adapter/test_snowflake.py +++ b/tests/core/engine_adapter/test_snowflake.py @@ -246,6 +246,38 @@ def test_multiple_column_comments(make_mocked_engine_adapter: t.Callable, mocker ] +def test_column_comments_iceberg(make_mocked_engine_adapter: t.Callable): + adapter = make_mocked_engine_adapter(SnowflakeEngineAdapter) + + adapter._create_column_comments( + "test_table", + {"a": "a column description", "b": "b column description"}, + table_format="iceberg", + ) + + assert to_sql_calls(adapter) == [ + """ALTER ICEBERG TABLE "test_table" ALTER COLUMN "a" COMMENT 'a column description', COLUMN "b" COMMENT 'b column description'""", + ] + + +def test_ctas_column_comments_iceberg(make_mocked_engine_adapter: t.Callable): + adapter = make_mocked_engine_adapter(SnowflakeEngineAdapter) + + # The column types are unknown, so the comments can't be inlined into the CTAS + # schema definition and are registered with a post-creation ALTER instead + adapter.ctas( + "test_table", + parse_one("SELECT a, b FROM source_table"), + table_format="iceberg", + column_descriptions={"a": "a column description"}, + ) + + assert to_sql_calls(adapter) == [ + """CREATE ICEBERG TABLE IF NOT EXISTS "test_table" AS SELECT "a", "b" FROM "source_table\"""", + """ALTER ICEBERG TABLE "test_table" ALTER COLUMN "a" COMMENT 'a column description'""", + ] + + def test_sync_grants_config(make_mocked_engine_adapter: t.Callable, mocker: MockerFixture): adapter = make_mocked_engine_adapter(SnowflakeEngineAdapter) relation = normalize_identifiers(