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
19 changes: 17 additions & 2 deletions sqlmesh/core/engine_adapter/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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():
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 @@ -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)
Expand Down
1 change: 1 addition & 0 deletions sqlmesh/core/engine_adapter/mysql.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
7 changes: 7 additions & 0 deletions sqlmesh/core/engine_adapter/snowflake.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down
32 changes: 32 additions & 0 deletions tests/core/engine_adapter/test_snowflake.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down