Adds comment on role propagation

src/backend/distributed/commands/database.c

@@ -701,3 +701,31 @@ CreateDatabaseDDLCommand(Oid dbId)
 
 	return outerDbStmt->data;
 }
+
+
+/*
+ * DatabaseCommentObjectAddress resolves the ObjectAddress for the DATABASE
+ * on which the comment is placed. Optionally errors if the database does not
+ * exist based on the missing_ok flag passed in by the caller.
+ */
+List *
+DatabaseCommentObjectAddress(Node *node, bool missing_ok, bool isPostprocess)
+{
+	CommentStmt *stmt = castNode(CommentStmt, node);
+	Assert(stmt->objtype == OBJECT_DATABASE);
+
+	char *databaseName = strVal(stmt->object);
+
+	Oid objid = get_database_oid(databaseName, missing_ok);
+
+	if (!OidIsValid(objid))
+	{
+		ereport(ERROR, (errmsg("database \"%s\" does not exist", databaseName)));
+	}
+	else
+	{
+		ObjectAddress *address = palloc0(sizeof(ObjectAddress));
+		ObjectAddressSet(*address, DatabaseRelationId, objid);
+		return list_make1(address);
+	}
+}

src/backend/distributed/commands/distribute_object_ops.c

@@ -293,6 +293,17 @@ static DistributeObjectOps Any_DropRole = {
 	.address = NULL,
 	.markDistributed = false,
 };
+
+static DistributeObjectOps Role_Comment = {
+	.deparse = DeparseRoleCommentStmt,
+	.qualify = NULL,
+	.preprocess = PreprocessAlterDistributedObjectStmt,
+	.postprocess = NULL,
+	.objectType = OBJECT_DATABASE,
+	.operationType = DIST_OPS_ALTER,
+	.address = RoleCommentObjectAddress,
+	.markDistributed = false,
+};
 static DistributeObjectOps Any_CreateForeignServer = {
 	.deparse = DeparseCreateForeignServerStmt,
 	.qualify = NULL,
@@ -521,8 +532,16 @@ static DistributeObjectOps Database_Set = {
 	.address = NULL,
 	.markDistributed = false,
 };
-
-
+static DistributeObjectOps Database_Comment = {
+	.deparse = DeparseDatabaseCommentStmt,
+	.qualify = NULL,
+	.preprocess = PreprocessAlterDistributedObjectStmt,
+	.postprocess = NULL,
+	.objectType = OBJECT_DATABASE,
+	.operationType = DIST_OPS_ALTER,
+	.address = DatabaseCommentObjectAddress,
+	.markDistributed = false,
+};
 static DistributeObjectOps Domain_Alter = {
 	.deparse = DeparseAlterDomainStmt,
 	.qualify = QualifyAlterDomainStmt,
@@ -1759,6 +1778,16 @@ GetDistributeObjectOps(Node *node)
 					return &TextSearchDict_Comment;
 				}
 
+				case OBJECT_DATABASE:
+				{
+					return &Database_Comment;
+				}
+
+				case OBJECT_ROLE:
+				{
+					return &Role_Comment;
+				}
+
 				default:
 				{
 					return &NoDistributeOps;

src/backend/distributed/commands/role.c

@@ -1412,3 +1412,31 @@ RenameRoleStmtObjectAddress(Node *node, bool missing_ok, bool isPostprocess)
 
 	return list_make1(address);
 }
+
+
+/*
+ * RoleCommentObjectAddress resolves the ObjectAddress for the ROLE
+ * on which the comment is placed. Optionally errors if the role does not
+ * exist based on the missing_ok flag passed in by the caller.
+ */
+List *
+RoleCommentObjectAddress(Node *node, bool missing_ok, bool isPostprocess)
+{
+	CommentStmt *stmt = castNode(CommentStmt, node);
+	Assert(stmt->objtype == OBJECT_ROLE);
+
+	char *roleName = strVal(stmt->object);
+
+	Oid objid = get_role_oid(roleName, missing_ok);
+
+	if (!OidIsValid(objid))
+	{
+		ereport(ERROR, (errmsg("role \"%s\" does not exist", roleName)));
+	}
+	else
+	{
+		ObjectAddress *address = palloc0(sizeof(ObjectAddress));
+		ObjectAddressSet(*address, AuthIdRelationId, objid);
+		return list_make1(address);
+	}
+}

src/backend/distributed/deparser/deparse_database_stmts.c

@@ -312,3 +312,21 @@ DeparseDropDatabaseStmt(Node *node)
 
 	return str.data;
 }
+
+
+char *
+DeparseDatabaseCommentStmt(Node *node)
+{
+	CommentStmt *stmt = castNode(CommentStmt, node);
+	StringInfoData str = { 0 };
+	initStringInfo(&str);
+
+	char *databaseName = strVal(stmt->object);
+
+
+	appendStringInfo(&str, "COMMENT ON DATABASE %s IS %s;",
+					 databaseName,
+					 quote_literal_cstr(stmt->comment));
+
+	return str.data;
+}

src/backend/distributed/deparser/deparse_role_stmts.c

@@ -533,3 +533,21 @@ AppendAlterRoleSetStmt(StringInfo buf, AlterRoleSetStmt *stmt)
 	VariableSetStmt *setStmt = castNode(VariableSetStmt, stmt->setstmt);
 	AppendVariableSet(buf, setStmt);
 }
+
+
+char *
+DeparseRoleCommentStmt(Node *node)
+{
+	CommentStmt *stmt = castNode(CommentStmt, node);
+	StringInfoData str = { 0 };
+	initStringInfo(&str);
+
+	char *roleName = strVal(stmt->object);
+
+
+	appendStringInfo(&str, "COMMENT ON ROLE %s IS %s;",
+					 roleName,
+					 quote_literal_cstr(stmt->comment));
+
+	return str.data;
+}

src/include/distributed/commands.h

@@ -246,6 +246,8 @@ extern List * CreateDatabaseStmtObjectAddress(Node *node, bool missingOk,
 											  bool isPostprocess);
 extern void EnsureSupportedCreateDatabaseCommand(CreatedbStmt *stmt);
 extern char * CreateDatabaseDDLCommand(Oid dbId);
+extern List * DatabaseCommentObjectAddress(Node *node, bool missing_ok, bool
+										   isPostprocess);
 
 
 /* domain.c - forward declarations */
@@ -517,6 +519,7 @@ extern List * RenameRoleStmtObjectAddress(Node *stmt, bool missing_ok, bool
 
 extern void UnmarkRolesDistributed(List *roles);
 extern List * FilterDistributedRoles(List *roles);
+List * RoleCommentObjectAddress(Node *node, bool missing_ok, bool isPostprocess);
 
 /* schema.c - forward declarations */
 extern List * PostprocessCreateSchemaStmt(Node *node, const char *queryString);

src/include/distributed/deparser.h

@@ -231,6 +231,7 @@ extern void QualifyAlterRoleSetStmt(Node *stmt);
 extern char * DeparseCreateRoleStmt(Node *stmt);
 extern char * DeparseDropRoleStmt(Node *stmt);
 extern char * DeparseGrantRoleStmt(Node *stmt);
+extern char * DeparseRoleCommentStmt(Node *node);
 
 /* forward declarations for deparse_owned_stmts.c */
 extern char * DeparseDropOwnedStmt(Node *node);
@@ -251,6 +252,7 @@ extern char * DeparseAlterDatabaseRefreshCollStmt(Node *node);
 extern char * DeparseAlterDatabaseSetStmt(Node *node);
 extern char * DeparseCreateDatabaseStmt(Node *node);
 extern char * DeparseDropDatabaseStmt(Node *node);
+extern char * DeparseDatabaseCommentStmt(Node *node);
 
 
 /* forward declaration for deparse_publication_stmts.c */

src/test/regress/expected/comment_on_database.out

@@ -0,0 +1,40 @@
+set citus.log_remote_commands to on;
+set citus.enable_create_database_propagation to on;
+set citus.grep_remote_commands to 'COMMENT ON DATABASE';
+create database test1;
+comment on DATABASE test1 is 'test-comment';
+SELECT result FROM run_command_on_all_nodes(
+  $$
+  SELECT ds.description AS database_comment
+FROM pg_database d
+LEFT JOIN pg_shdescription ds ON d.oid = ds.objoid
+WHERE d.datname = 'test1';
+  $$
+);
+    result
+---------------------------------------------------------------------
+ test-comment
+ test-comment
+ test-comment
+(3 rows)
+
+comment on DATABASE test1 is 'comment-needs\!escape';
+SELECT result FROM run_command_on_all_nodes(
+  $$
+  SELECT ds.description AS database_comment
+FROM pg_database d
+LEFT JOIN pg_shdescription ds ON d.oid = ds.objoid
+WHERE d.datname = 'test1';
+  $$
+);
+        result
+---------------------------------------------------------------------
+ comment-needs\!escape
+ comment-needs\!escape
+ comment-needs\!escape
+(3 rows)
+
+drop DATABASE test1;
+reset citus.enable_create_database_propagation;
+reset citus.grep_remote_commands;
+reset citus.log_remote_commands;

src/test/regress/multi_1_schedule

@@ -59,6 +59,8 @@ test: grant_on_database_propagation
 test: alter_database_propagation
 
 test: citus_shards
+test: comment_on_database
+test: comment_on_role
 
 # ----------
 # multi_citus_tools tests utility functions written for citus tools

src/test/regress/sql/comment_on_database.sql

@@ -0,0 +1,33 @@
+set citus.log_remote_commands to on;
+
+set citus.enable_create_database_propagation to on;
+set citus.grep_remote_commands to 'COMMENT ON DATABASE';
+
+create database test1;
+
+comment on DATABASE test1 is 'test-comment';
+
+SELECT result FROM run_command_on_all_nodes(
+  $$
+    SELECT ds.description AS database_comment
+    FROM pg_database d
+    LEFT JOIN pg_shdescription ds ON d.oid = ds.objoid
+    WHERE d.datname = 'test1';
+  $$
+);
+
+comment on DATABASE test1 is 'comment-needs\!escape';
+
+SELECT result FROM run_command_on_all_nodes(
+  $$
+    SELECT ds.description AS database_comment
+    FROM pg_database d
+    LEFT JOIN pg_shdescription ds ON d.oid = ds.objoid
+    WHERE d.datname = 'test1';
+  $$
+);
+
+drop DATABASE test1;
+reset citus.enable_create_database_propagation;
+reset citus.grep_remote_commands;
+reset citus.log_remote_commands;

src/test/regress/sql/comment_on_role.sql

@@ -0,0 +1,31 @@
+set citus.log_remote_commands to on;
+
+set citus.grep_remote_commands to 'COMMENT ON ROLE';
+
+create role role1;
+
+comment on ROLE role1 is 'test-comment';
+
+SELECT result FROM run_command_on_all_nodes(
+  $$
+    SELECT ds.description AS role_comment
+    FROM pg_roles r
+    LEFT JOIN pg_shdescription ds ON r.oid = ds.objoid
+    WHERE r.rolname = 'role1';
+  $$
+);
+
+comment on role role1 is 'comment-needs\!escape';
+
+SELECT result FROM run_command_on_all_nodes(
+  $$
+    SELECT ds.description AS role_comment
+    FROM pg_roles r
+    LEFT JOIN pg_shdescription ds ON r.oid = ds.objoid
+    WHERE r.rolname = 'role1';
+  $$
+);
+
+drop role role1;
+reset citus.grep_remote_commands;
+reset citus.log_remote_commands;