Comment registration support for MSSQL engine

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

Author: ujfalusis

sqlmesh/core/engine_adapter/mssql.py

@@ -459,63 +459,74 @@ 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"
+    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(column_comment).sql(dialect=self.dialect) if column_comment is not None else "NULL"};
+        template = dedent("""
+            DECLARE @comment sql_variant = {comment};
             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 @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, 'column', @column_name);
+            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, 'column', @column_name;
+                        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, 'column', @column_name;
+                        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, 'column', @column_name;
+                        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_table_exp(
-        self, table: exp.Table, table_comment: str, table_kind: str
+    def _build_create_comment_column_exp(
+        self, table: exp.Table, column_name: str, column_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"};
+        template = dedent("""
+            DECLARE @comment sql_variant = {comment};
             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 @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, DEFAULT, DEFAULT);
+            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;
+                        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;
+                        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;
+                        EXEC sp_updateextendedproperty @property_name, @comment, 'schema', @schema_name, @object_kind, @object_name, 'column', @column_name;
                 END
         """)
-        return tsql_text
+        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