Skip to content
2 changes: 1 addition & 1 deletion docs/concepts/models/overview.md
Original file line number Diff line number Diff line change
Expand Up @@ -184,7 +184,7 @@ This table lists each engine's support for `TABLE` and `VIEW` object comments:
| DuckDB <=0.9 | N | N |
| DuckDB >=0.10 | Y | Y |
| MySQL | Y | Y |
| MSSQL | N | N |
| MSSQL | Y | Y |
| Postgres | Y | Y |
| GCP Postgres | Y | Y |
| Redshift | Y | N |
Expand Down
77 changes: 75 additions & 2 deletions sqlmesh/core/engine_adapter/mssql.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

from __future__ import annotations

from textwrap import dedent
import typing as t
import logging

Expand Down Expand Up @@ -53,8 +54,8 @@ class MSSQLEngineAdapter(
SUPPORTS_TUPLE_IN = False
SUPPORTS_MATERIALIZED_VIEWS = False
CURRENT_CATALOG_EXPRESSION = exp.func("db_name")
COMMENT_CREATION_TABLE = CommentCreationTable.UNSUPPORTED
COMMENT_CREATION_VIEW = CommentCreationView.UNSUPPORTED
COMMENT_CREATION_TABLE = CommentCreationTable.COMMENT_COMMAND_ONLY
COMMENT_CREATION_VIEW = CommentCreationView.COMMENT_COMMAND_ONLY
SUPPORTS_REPLACE_TABLE = False
MAX_IDENTIFIER_LENGTH = 128
SUPPORTS_QUERY_EXECUTION_TRACKING = True
Expand Down Expand Up @@ -457,3 +458,75 @@ def delete_from(self, table_name: TableName, where: t.Union[str, exp.Expr]) -> N
)

return super().delete_from(table_name, where)

def _build_create_comment_table_exp(
self, table: exp.Table, table_comment: str, table_kind: str
) -> exp.Comment | str:
template = dedent("""
DECLARE @comment sql_variant = {comment};
DECLARE @property_name VARCHAR(128) = 'MS_Description';
DECLARE @schema_name VARCHAR(128) = '{schema_name}';
DECLARE @object_name VARCHAR(128) = '{object_name}';
DECLARE @object_kind VARCHAR(128) = '{object_kind}';
DECLARE @existing sql_variant;

SELECT TOP 1 @existing = CAST(VALUE AS NVARCHAR) FROM fn_listextendedproperty(@property_name, 'schema', @schema_name, @object_kind, @object_name, DEFAULT, DEFAULT);

IF @comment IS NULL
BEGIN
IF @existing IS NOT NULL
EXEC sp_dropextendedproperty @property_name, 'schema', @schema_name, @object_kind, @object_name;
END
ELSE
BEGIN
IF @existing IS NULL
EXEC sp_addextendedproperty @property_name,@comment, 'schema', @schema_name, @object_kind, @object_name;
ELSE IF @existing != @comment
EXEC sp_updateextendedproperty @property_name, @comment, 'schema', @schema_name, @object_kind, @object_name;
END
""")
tsql_text = template.format(
comment = exp.Literal.string(table_comment).sql(dialect=self.dialect) if table_comment is not None else "NULL",
schema_name = table.db if table.db else "dbo",
object_name = table.name,
object_kind = table_kind,
)
return tsql_text

def _build_create_comment_column_exp(
self, table: exp.Table, column_name: str, column_comment: str, table_kind: str
) -> exp.Comment | str:

template = dedent("""
DECLARE @comment sql_variant = {comment};
DECLARE @property_name VARCHAR(128) = 'MS_Description';
DECLARE @schema_name VARCHAR(128) = '{schema_name}';
DECLARE @object_name VARCHAR(128) = '{object_name}';
DECLARE @object_kind VARCHAR(128) = '{object_kind}';
DECLARE @column_name VARCHAR(128) = '{column_name}';
DECLARE @existing sql_variant;

SELECT TOP 1 @existing = CAST(VALUE AS NVARCHAR) FROM fn_listextendedproperty(@property_name, 'schema', @schema_name, @object_kind, @object_name, 'column', @column_name);

IF @comment IS NULL
BEGIN
IF @existing IS NOT NULL
EXEC sp_dropextendedproperty @property_name, 'schema', @schema_name, @object_kind, @object_name, 'column', @column_name;
END
ELSE
BEGIN
IF @existing IS NULL
EXEC sp_addextendedproperty @property_name,@comment, 'schema', @schema_name, @object_kind, @object_name, 'column', @column_name;
ELSE IF @existing != @comment
EXEC sp_updateextendedproperty @property_name, @comment, 'schema', @schema_name, @object_kind, @object_name, 'column', @column_name;
END
""")
tsql_text = template.format(
comment = exp.Literal.string(column_comment).sql(dialect=self.dialect) if column_comment is not None else "NULL",
schema_name = table.db if table.db else "dbo",
object_name = table.name,
object_kind = table_kind,
column_name = column_name,
)

return tsql_text
18 changes: 18 additions & 0 deletions tests/core/engine_adapter/integration/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -526,6 +526,14 @@ def get_table_comment(
AND c.relkind = '{"v" if table_kind == "VIEW" else "r"}'
;
"""
elif self.dialect == "tsql":
kind = "table" if table_kind == "BASE TABLE" else "view"
query = f"""
SELECT
ep.name,
CAST(ep.value AS NVARCHAR(MAX)) comment
FROM fn_listextendedproperty('MS_Description', 'schema', '{schema_name}', '{kind}', '{table_name}', DEFAULT, DEFAULT) ep
"""

result = self.engine_adapter.fetchall(query)

Expand Down Expand Up @@ -636,6 +644,16 @@ def get_column_comments(
AND c.relkind = '{"v" if table_kind == "VIEW" else "r"}'
;
"""
elif self.dialect == "tsql":
kind = "table" if table_kind == "BASE TABLE" else "view"
query = f"""
SELECT
col.COLUMN_NAME column_name,
CAST(ep.value AS NVARCHAR(MAX)) comment
FROM INFORMATION_SCHEMA.COLUMNS col
CROSS APPLY fn_listextendedproperty('MS_Description', 'schema', col.TABLE_SCHEMA, '{kind}', col.TABLE_NAME, 'column', col.COLUMN_NAME) ep
WHERE col.TABLE_SCHEMA = '{schema_name}' AND col.TABLE_NAME = '{table_name}'
"""

result = self.engine_adapter.fetchall(query)

Expand Down