SNOW-1344775: Support adding comments to tables and views (#1489)

CHANGELOG.md

@@ -2,6 +2,14 @@
 
 ## 1.17.0 (TBD)
 
+### New Features
+
+- Added support to add a comment on tables and views using functions listed below:
+  - `DataFrameWriter.save_as_table`
+  - `DataFrame.create_or_replace_view`
+  - `DataFrame.create_or_replace_temp_view`
+  - `DataFrame.create_or_replace_dynamic_table`
+
 ### Improvements
 
 - Improved error message to remind users set `{"infer_schema": True}` when reading csv file without specifying its schema.

src/snowflake/snowpark/_internal/analyzer/analyzer.py

@@ -939,6 +939,7 @@ def do_resolve_with_resolved_children(
                     self.analyze(x, df_aliased_col_name_to_real_col_name)
                     for x in logical_plan.clustering_exprs
                 ],
+                logical_plan.comment,
                 resolved_children[logical_plan.children[0]],
             )
 
@@ -1064,14 +1065,18 @@ def do_resolve_with_resolved_children(
                 )
 
             return self.plan_builder.create_or_replace_view(
-                logical_plan.name, resolved_children[logical_plan.child], is_temp
+                logical_plan.name,
+                resolved_children[logical_plan.child],
+                is_temp,
+                logical_plan.comment,
             )
 
         if isinstance(logical_plan, CreateDynamicTableCommand):
             return self.plan_builder.create_or_replace_dynamic_table(
                 logical_plan.name,
                 logical_plan.warehouse,
                 logical_plan.lag,
+                logical_plan.comment,
                 resolved_children[logical_plan.child],
             )
 

src/snowflake/snowpark/_internal/analyzer/analyzer_utils.py

@@ -29,6 +29,7 @@
     EMPTY_STRING,
     TempObjectType,
     escape_quotes,
+    escape_single_quotes,
     get_temp_type_for_object,
     is_single_quoted,
     is_sql_select_statement,
@@ -154,6 +155,7 @@
 MATCHED = " MATCHED "
 LISTAGG = " LISTAGG "
 HEADER = " HEADER "
+COMMENT = " COMMENT "
 IGNORE_NULLS = " IGNORE NULLS "
 UNION = " UNION "
 UNION_ALL = " UNION ALL "
@@ -737,13 +739,22 @@ def join_statement(
     )
 
 
+def get_comment_sql(comment: Optional[str]) -> str:
+    return (
+        COMMENT + EQUALS + SINGLE_QUOTE + escape_single_quotes(comment) + SINGLE_QUOTE
+        if comment
+        else EMPTY_STRING
+    )
+
+
 def create_table_statement(
     table_name: str,
     schema: str,
     replace: bool = False,
     error: bool = True,
     table_type: str = EMPTY_STRING,
     clustering_key: Optional[Iterable[str]] = None,
+    comment: Optional[str] = None,
     *,
     use_scoped_temp_objects: bool = False,
     is_generated: bool = False,
@@ -753,12 +764,13 @@ def create_table_statement(
         if clustering_key
         else EMPTY_STRING
     )
+    comment_sql = get_comment_sql(comment)
     return (
         f"{CREATE}{(OR + REPLACE) if replace else EMPTY_STRING}"
         f" {(get_temp_type_for_object(use_scoped_temp_objects, is_generated) if table_type.lower() in TEMPORARY_STRING_SET else table_type).upper()} "
         f"{TABLE}{table_name}{(IF + NOT + EXISTS) if not replace and not error else EMPTY_STRING}"
         f"{LEFT_PARENTHESIS}{schema}{RIGHT_PARENTHESIS}"
-        f"{cluster_by_clause}"
+        f"{cluster_by_clause}{comment_sql}"
     )
 
 
@@ -791,17 +803,19 @@ def create_table_as_select_statement(
     error: bool = True,
     table_type: str = EMPTY_STRING,
     clustering_key: Optional[Iterable[str]] = None,
+    comment: Optional[str] = None,
 ) -> str:
     cluster_by_clause = (
         (CLUSTER_BY + LEFT_PARENTHESIS + COMMA.join(clustering_key) + RIGHT_PARENTHESIS)
         if clustering_key
         else EMPTY_STRING
     )
+    comment_sql = get_comment_sql(comment)
     return (
         f"{CREATE}{OR + REPLACE if replace else EMPTY_STRING} {table_type.upper()} {TABLE}"
         f"{IF + NOT + EXISTS if not replace and not error else EMPTY_STRING}"
         f" {table_name}{LEFT_PARENTHESIS}{column_definition}{RIGHT_PARENTHESIS}"
-        f"{cluster_by_clause} {AS}{project_statement([], child)}"
+        f"{cluster_by_clause} {comment_sql} {AS}{project_statement([], child)}"
     )
 
 
@@ -1006,22 +1020,27 @@ def order_expression(name: str, direction: str, null_ordering: str) -> str:
     return name + SPACE + direction + SPACE + null_ordering
 
 
-def create_or_replace_view_statement(name: str, child: str, is_temp: bool) -> str:
+def create_or_replace_view_statement(
+    name: str, child: str, is_temp: bool, comment: Optional[str]
+) -> str:
+    comment_sql = get_comment_sql(comment)
     return (
         CREATE
         + OR
         + REPLACE
         + f"{TEMPORARY if is_temp else EMPTY_STRING}"
         + VIEW
         + name
+        + comment_sql
         + AS
         + project_statement([], child)
     )
 
 
 def create_or_replace_dynamic_table_statement(
-    name: str, warehouse: str, lag: str, child: str
+    name: str, warehouse: str, lag: str, comment: Optional[str], child: str
 ) -> str:
+    comment_sql = get_comment_sql(comment)
     return (
         CREATE
         + OR
@@ -1031,6 +1050,7 @@ def create_or_replace_dynamic_table_statement(
         + name
         + f"{LAG + EQUALS + convert_value_to_sql_option(lag)}"
         + f"{WAREHOUSE + EQUALS + warehouse}"
+        + comment_sql
         + AS
         + project_statement([], child)
     )

src/snowflake/snowpark/_internal/analyzer/snowflake_plan.py

@@ -695,6 +695,7 @@ def save_as_table(
         mode: SaveMode,
         table_type: str,
         clustering_keys: Iterable[str],
+        comment: Optional[str],
         child: SnowflakePlan,
     ) -> SnowflakePlan:
         full_table_name = ".".join(table_name)
@@ -722,6 +723,7 @@ def get_create_and_insert_plan(child: SnowflakePlan, replace=False, error=True):
                 error=error,
                 table_type=table_type,
                 clustering_key=clustering_keys,
+                comment=comment,
             )
 
             # so that dataframes created from non-select statements,
@@ -779,6 +781,7 @@ def get_create_and_insert_plan(child: SnowflakePlan, replace=False, error=True):
                         replace=True,
                         table_type=table_type,
                         clustering_key=clustering_keys,
+                        comment=comment,
                     ),
                     child,
                     None,
@@ -792,6 +795,7 @@ def get_create_and_insert_plan(child: SnowflakePlan, replace=False, error=True):
                     replace=True,
                     table_type=table_type,
                     clustering_key=clustering_keys,
+                    comment=comment,
                 ),
                 child,
                 None,
@@ -805,6 +809,7 @@ def get_create_and_insert_plan(child: SnowflakePlan, replace=False, error=True):
                     error=False,
                     table_type=table_type,
                     clustering_key=clustering_keys,
+                    comment=comment,
                 ),
                 child,
                 None,
@@ -817,6 +822,7 @@ def get_create_and_insert_plan(child: SnowflakePlan, replace=False, error=True):
                     column_definition,
                     table_type=table_type,
                     clustering_key=clustering_keys,
+                    comment=comment,
                 ),
                 child,
                 None,
@@ -880,7 +886,7 @@ def rename(
         )
 
     def create_or_replace_view(
-        self, name: str, child: SnowflakePlan, is_temp: bool
+        self, name: str, child: SnowflakePlan, is_temp: bool, comment: Optional[str]
     ) -> SnowflakePlan:
         if len(child.queries) != 1:
             raise SnowparkClientExceptionMessages.PLAN_CREATE_VIEW_FROM_DDL_DML_OPERATIONS()
@@ -890,7 +896,7 @@ def create_or_replace_view(
 
         child = child.replace_repeated_subquery_with_cte()
         return self.build(
-            lambda x: create_or_replace_view_statement(name, x, is_temp),
+            lambda x: create_or_replace_view_statement(name, x, is_temp, comment),
             child,
             None,
         )
@@ -900,6 +906,7 @@ def create_or_replace_dynamic_table(
         name: str,
         warehouse: str,
         lag: str,
+        comment: Optional[str],
         child: SnowflakePlan,
     ) -> SnowflakePlan:
         if len(child.queries) != 1:
@@ -911,7 +918,7 @@ def create_or_replace_dynamic_table(
         child = child.replace_repeated_subquery_with_cte()
         return self.build(
             lambda x: create_or_replace_dynamic_table_statement(
-                name, warehouse, lag, x
+                name, warehouse, lag, comment, x
             ),
             child,
             None,

src/snowflake/snowpark/_internal/analyzer/snowflake_plan_node.py

@@ -76,6 +76,7 @@ def __init__(
         query: Optional[LogicalPlan],
         table_type: str = "",
         clustering_exprs: Optional[Iterable[Expression]] = None,
+        comment: Optional[str] = None,
     ) -> None:
         super().__init__()
         self.table_name = table_name
@@ -85,6 +86,7 @@ def __init__(
         self.table_type = table_type
         self.children.append(query)
         self.clustering_exprs = clustering_exprs or []
+        self.comment = comment
 
 
 class Limit(LogicalPlan):

src/snowflake/snowpark/_internal/analyzer/unary_plan_node.py

@@ -120,11 +120,13 @@ def __init__(
         self,
         name: str,
         view_type: ViewType,
+        comment: Optional[str],
         child: LogicalPlan,
     ) -> None:
         super().__init__(child)
         self.name = name
         self.view_type = view_type
+        self.comment = comment
 
 
 class CreateDynamicTableCommand(UnaryNode):
@@ -133,9 +135,11 @@ def __init__(
         name: str,
         warehouse: str,
         lag: str,
+        comment: Optional[str],
         child: LogicalPlan,
     ) -> None:
         super().__init__(child)
         self.name = name
         self.warehouse = warehouse
         self.lag = lag
+        self.comment = comment

src/snowflake/snowpark/_internal/udf_utils.py

@@ -41,6 +41,7 @@
 from snowflake.snowpark._internal.utils import (
     STAGE_PREFIX,
     TempObjectType,
+    escape_single_quotes,
     get_udf_upload_prefix,
     is_single_quoted,
     normalize_remote_file_or_dir,
@@ -1264,7 +1265,7 @@ def create_python_udf_or_sp(
 
     if comment is not None:
         object_signature_sql = f"{object_name}({','.join(input_sql_types)})"
-        comment = comment.replace("'", "\\'")  # escaping single quotes
+        comment = escape_single_quotes(comment)
         comment_query = f"COMMENT ON {object_type.value.split('_')[-1]} {object_signature_sql} IS '{comment}'"
         session._run_query(
             comment_query,

src/snowflake/snowpark/_internal/utils.py

@@ -266,6 +266,10 @@ def unwrap_single_quote(name: str) -> str:
     return new_name
 
 
+def escape_single_quotes(input_str):
+    return input_str.replace("'", r"\'")
+
+
 def is_sql_select_statement(sql: str) -> bool:
     return (
         SNOWFLAKE_SELECT_SQL_PREFIX_PATTERN.match(sql) is not None