Comment registration support for MSSQL engine

Signed-off-by: Ujfalusi Sándor <ujfalusi.sandor@gmail.com>

Author: ujfalusis

docs/concepts/models/overview.md

@@ -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               |

sqlmesh/core/engine_adapter/mssql.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+from textwrap import dedent
 import typing as t
 import logging
 
@@ -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
@@ -457,3 +458,64 @@ 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_column_exp(
+        self, table: exp.Table, column_name: str, column_comment: str, table_kind: str = "TABLE"
+    ) -> exp.Comment | str:
+        tsql_text = dedent(f"""
+            SET NOCOUNT ON;
+
+            DECLARE @comment sql_variant = {exp.Literal.string(column_comment).sql(dialect=self.dialect) if column_comment is not None else "NULL"};
+            DECLARE @property_name VARCHAR(128) = 'MS_Description';
+            DECLARE @schema_name VARCHAR(128) = '{table.db if table.db else "dbo"}';
+            DECLARE @object_name VARCHAR(128) = '{table.name}';
+            DECLARE @object_kind VARCHAR(128) = '{table_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
+        """)
+        return tsql_text
+
+    def _build_create_comment_table_exp(
+        self, table: exp.Table, table_comment: str, table_kind: str
+    ) -> exp.Comment | str:
+        tsql_text = dedent(f"""
+            SET NOCOUNT ON;
+
+            DECLARE @comment sql_variant = {exp.Literal.string(table_comment).sql(dialect=self.dialect) if table_comment is not None else "NULL"};
+            DECLARE @property_name VARCHAR(128) = 'MS_Description';
+            DECLARE @schema_name VARCHAR(128) = '{table.db if table.db else "dbo"}';
+            DECLARE @object_name VARCHAR(128) = '{table.name}';
+            DECLARE @object_kind VARCHAR(128) = '{table_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
+        """)
+        return tsql_text

tests/core/engine_adapter/test_mssql.py

@@ -1002,3 +1002,26 @@ def python_scd2_model(context, **kwargs):
     snapshot: Snapshot = make_snapshot(m)
     assert snapshot.node.physical_properties == m.physical_properties
     assert snapshot.node.physical_properties.get("mssql_merge_exists")
+
+
+def test_comments(make_mocked_engine_adapter: t.Callable, mocker: MockerFixture):
+    adapter = make_mocked_engine_adapter(MSSQLEngineAdapter)
+
+    columns_to_types = {
+        "cola": exp.DataType.build("INT"),
+        "colb": exp.DataType.build("TEXT"),
+    }
+    adapter.create_table(
+        "test_table", columns_to_types, table_description="\\", column_descriptions={"cola": "\\"}
+    )
+
+    sql_calls = to_sql_calls(adapter)
+    assert sql_calls == [
+        """IF NOT EXISTS (SELECT * FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_NAME = 'test_table') EXEC('CREATE TABLE [test_table] ([cola] INTEGER, [colb] VARCHAR(MAX))');""",
+        adapter._build_create_comment_table_exp(
+            exp.table_("test_table", quoted=True), "\\", "TABLE"
+        ),
+        adapter._build_create_comment_column_exp(
+            exp.table_("test_table", quoted=True), "cola", "\\", "TABLE"
+        ),
+    ]