Recreate table and column comments on undistribute_table (#6768)

src/backend/distributed/commands/alter_table.c

@@ -47,6 +47,7 @@
 #include "distributed/colocation_utils.h"
 #include "distributed/commands.h"
 #include "distributed/commands/utility_hook.h"
+#include "distributed/comment.h"
 #include "distributed/coordinator_protocol.h"
 #include "distributed/deparser.h"
 #include "distributed/distribution_column.h"
@@ -753,6 +754,28 @@ ConvertTableInternal(TableConversionState *con)
 	postLoadCommands = list_concat(postLoadCommands,
 								   WrapTableDDLCommands(alterPublicationCommands));
 
+	if (con->conversionType == UNDISTRIBUTE_TABLE)
+	{
+		List *commentDDLCommandsTable = GetCommentPropagationCommandsX(
+			con->relationId, RelationRelationId, con->relationName, OBJECT_TABLE,
+			NULL, 0);
+		postLoadCommands = list_concat(postLoadCommands,
+									   WrapTableDDLCommands(commentDDLCommandsTable));
+		List *nonStoredColumnNameList = GetNonGeneratedStoredColumnNameList(
+			con->relationId);
+		char *columnName = NULL;
+		int columnCount = 0;
+		foreach_ptr(columnName, nonStoredColumnNameList)
+		{
+			List *commentDDLCommandsColumn = GetCommentPropagationCommandsX(
+				con->relationId, RelationRelationId, columnName, OBJECT_COLUMN,
+				con->relationName, ++columnCount);
+			postLoadCommands = list_concat(postLoadCommands,
+										   WrapTableDDLCommands(
+											   commentDDLCommandsColumn));
+		}
+	}
+
 	List *foreignKeyCommands = NIL;
 	if (con->conversionType == ALTER_DISTRIBUTED_TABLE)
 	{

src/backend/distributed/commands/comment.c

@@ -15,6 +15,7 @@
 #include "access/htup_details.h"
 #include "access/table.h"
 #include "catalog/pg_shdescription.h"
+#include "commands/comment.h"
 #include "nodes/parsenodes.h"
 #include "utils/builtins.h"
 #include "utils/fmgroids.h"
@@ -25,23 +26,43 @@
 static char * GetCommentForObject(Oid classOid, Oid objectOid);
 
 
-List *
+inline List *
 GetCommentPropagationCommands(Oid classOid, Oid objOoid, char *objectName, ObjectType
 							  objectType)
+{
+	return GetCommentPropagationCommandsX(classOid, objOoid, objectName, objectType, NULL,
+										  0);
+}
+
+
+List *
+GetCommentPropagationCommandsX(Oid classOid, Oid objOoid, char *objectName, ObjectType
+							   objectType, char *qualifier, int32 subid)
 {
 	List *commands = NIL;
 
 	StringInfo commentStmt = makeStringInfo();
 
-	/* Get the comment for the database */
-	char *comment = GetCommentForObject(classOid, objOoid);
+	char *comment = NULL;
+
+	if ((objectType == OBJECT_DATABASE) || (objectType == OBJECT_ROLE) || (objectType ==
+																		   OBJECT_TABLESPACE))
+	{
+		/* Get the comment for the shared object */
+		comment = GetCommentForObject(classOid, objOoid);
+	}
+	else
+	{
+		comment = GetComment(classOid, objOoid, subid);
+	}
+
 	char const *commentObjectType = ObjectTypeNames[objectType];
 
 	/* Create the SQL command to propagate the comment to other nodes */
 	if (comment != NULL)
 	{
 		appendStringInfo(commentStmt, "COMMENT ON %s %s IS %s;", commentObjectType,
-						 quote_identifier(objectName),
+						 quote_qualified_identifier(qualifier, objectName),
 						 quote_literal_cstr(comment));
 	}
 

src/backend/distributed/deparser/deparse_comment_stmts.c

@@ -35,6 +35,8 @@ const char *ObjectTypeNames[] =
 	[OBJECT_ROLE] = "ROLE",
 	[OBJECT_TSCONFIGURATION] = "TEXT SEARCH CONFIGURATION",
 	[OBJECT_TSDICTIONARY] = "TEXT SEARCH DICTIONARY",
+	[OBJECT_TABLE] = "TABLE",
+	[OBJECT_COLUMN] = "COLUMN",
 
 	/* When support for propagating comments to new objects is introduced, an entry for each
 	 * statement type should be added to this list. The first element in each entry is the 'object_type' keyword

src/include/distributed/comment.h

@@ -8,8 +8,8 @@
  *-------------------------------------------------------------------------
  */
 
-#ifndef COMMENT_H
-#define COMMENT_H
+#ifndef CITUS_COMMENT_H
+#define CITUS_COMMENT_H
 
 #include "postgres.h"
 
@@ -21,6 +21,9 @@ extern const char *ObjectTypeNames[];
 
 extern List * GetCommentPropagationCommands(Oid classOid, Oid oid, char *objectName,
 											ObjectType objectType);
+extern List * GetCommentPropagationCommandsX(Oid classOid, Oid oid, char *objectName,
+											 ObjectType objectType, char *qualifier,
+											 int32 subid);
 extern List * CommentObjectAddress(Node *node, bool missing_ok, bool isPostprocess);
 
-# endif /* COMMENT_H */
+# endif /* CITUS_COMMENT_H */

src/test/regress/expected/comment_on_table_column.out

@@ -0,0 +1,41 @@
+CREATE SCHEMA comment_on_table_and_column;
+SET search_path TO comment_on_table_and_column;
+create table tbl (a int, b text);
+comment on table tbl is 'table comment';
+comment on column tbl.b is 'column b comment';
+select col_description('tbl'::regclass,0) as table_comment, col_description('tbl'::regclass,1) as column1_comment, col_description('tbl'::regclass,2) as column2_comment;
+ table_comment | column1_comment | column2_comment
+---------------------------------------------------------------------
+ table comment |                 | column b comment
+(1 row)
+
+select create_distributed_table('tbl','a');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+select col_description('tbl'::regclass,0) as table_comment, col_description('tbl'::regclass,1) as column1_comment, col_description('tbl'::regclass,2) as column2_comment;
+ table_comment | column1_comment | column2_comment
+---------------------------------------------------------------------
+ table comment |                 | column b comment
+(1 row)
+
+select undistribute_table('tbl');
+NOTICE:  creating a new table for comment_on_table_and_column.tbl
+NOTICE:  moving the data of comment_on_table_and_column.tbl
+NOTICE:  dropping the old comment_on_table_and_column.tbl
+NOTICE:  renaming the new table to comment_on_table_and_column.tbl
+ undistribute_table
+---------------------------------------------------------------------
+
+(1 row)
+
+select col_description('tbl'::regclass,0) as table_comment, col_description('tbl'::regclass,1) as column1_comment, col_description('tbl'::regclass,2) as column2_comment;
+ table_comment | column1_comment | column2_comment
+---------------------------------------------------------------------
+ table comment |                 | column b comment
+(1 row)
+
+DROP SCHEMA comment_on_table_and_column CASCADE;
+NOTICE:  drop cascades to table tbl

src/test/regress/multi_schedule

@@ -111,6 +111,7 @@ test: background_task_queue_monitor
 test: other_databases grant_role_from_non_maindb role_operations_from_non_maindb seclabel_non_maindb
 test: citus_internal_access
 test: function_with_case_when
+test: comment_on_table_column
 
 # Causal clock test
 test: clock

src/test/regress/sql/comment_on_table_column.sql

@@ -0,0 +1,15 @@
+CREATE SCHEMA comment_on_table_and_column;
+SET search_path TO comment_on_table_and_column;
+
+create table tbl (a int, b text);
+comment on table tbl is 'table comment';
+comment on column tbl.b is 'column b comment';
+select col_description('tbl'::regclass,0) as table_comment, col_description('tbl'::regclass,1) as column1_comment, col_description('tbl'::regclass,2) as column2_comment;
+
+select create_distributed_table('tbl','a');
+select col_description('tbl'::regclass,0) as table_comment, col_description('tbl'::regclass,1) as column1_comment, col_description('tbl'::regclass,2) as column2_comment;
+
+select undistribute_table('tbl');
+select col_description('tbl'::regclass,0) as table_comment, col_description('tbl'::regclass,1) as column1_comment, col_description('tbl'::regclass,2) as column2_comment;
+
+DROP SCHEMA comment_on_table_and_column CASCADE;