Adds REASSIGN OWNED BY propagation

src/backend/distributed/commands/distribute_object_ops.c

@@ -274,6 +274,17 @@ static DistributeObjectOps Any_CreateRole = {
 	.address = CreateRoleStmtObjectAddress,
 	.markDistributed = true,
 };
+
+static DistributeObjectOps Any_ReassignOwned = {
+	.deparse = DeparseReassignOwnedStmt,
+	.qualify = NULL,
+	.preprocess = PreprocessReassignOwnedStmt,
+	.postprocess = NULL,
+	.operationType = DIST_OPS_DROP,
+	.address = NULL,
+	.markDistributed = false,
+};
+
 static DistributeObjectOps Any_DropOwned = {
 	.deparse = DeparseDropOwnedStmt,
 	.qualify = NULL,
@@ -1826,6 +1837,11 @@ GetDistributeObjectOps(Node *node)
 			return &Any_DropOwned;
 		}
 
+		case T_ReassignOwnedStmt:
+		{
+			return &Any_ReassignOwned;
+		}
+
 		case T_DropStmt:
 		{
 			DropStmt *stmt = castNode(DropStmt, node);

src/backend/distributed/commands/owned.c

@@ -88,3 +88,36 @@ PreprocessDropOwnedStmt(Node *node, const char *queryString,
 
 	return NodeDDLTaskList(NON_COORDINATOR_NODES, commands);
 }
+
+
+List *
+PreprocessReassignOwnedStmt(Node *node, const char *queryString,
+							ProcessUtilityContext processUtilityContext)
+{
+	ReassignOwnedStmt *stmt = castNode(ReassignOwnedStmt, node);
+	List *allReassignRoles = stmt->roles;
+
+	List *distributedReassignRoles = FilterDistributedRoles(allReassignRoles);
+
+	if (list_length(distributedReassignRoles) <= 0)
+	{
+		return NIL;
+	}
+
+	if (!ShouldPropagate())
+	{
+		return NIL;
+	}
+
+	EnsureCoordinator();
+
+	stmt->roles = distributedReassignRoles;
+	char *sql = DeparseTreeNode((Node *) stmt);
+	stmt->roles = allReassignRoles;
+
+	List *commands = list_make3(DISABLE_DDL_PROPAGATION,
+								sql,
+								ENABLE_DDL_PROPAGATION);
+
+	return NodeDDLTaskList(NON_COORDINATOR_NODES, commands);
+}

src/backend/distributed/deparser/deparse_owned_stmts.c

@@ -82,3 +82,27 @@ AppendRoleList(StringInfo buf, List *roleList)
 		}
 	}
 }
+
+
+static void
+AppendReassignOwnedStmt(StringInfo buf, ReassignOwnedStmt *stmt)
+{
+	appendStringInfo(buf, "REASSIGN OWNED BY ");
+
+	AppendRoleList(buf, stmt->roles);
+	char const *newRoleName = RoleSpecString(stmt->newrole, true);
+	appendStringInfo(buf, " TO %s", quote_identifier(newRoleName));
+}
+
+
+char *
+DeparseReassignOwnedStmt(Node *node)
+{
+	ReassignOwnedStmt *stmt = castNode(ReassignOwnedStmt, node);
+	StringInfoData buf = { 0 };
+	initStringInfo(&buf);
+
+	AppendReassignOwnedStmt(&buf, stmt);
+
+	return buf.data;
+}

src/include/distributed/commands.h

@@ -427,6 +427,8 @@ extern List * CreateExtensionStmtObjectAddress(Node *stmt, bool missing_ok, bool
 /* owned.c -  forward declarations */
 extern List * PreprocessDropOwnedStmt(Node *node, const char *queryString,
 									  ProcessUtilityContext processUtilityContext);
+extern List * PreprocessReassignOwnedStmt(Node *node, const char *queryString,
+										  ProcessUtilityContext processUtilityContext);
 
 /* policy.c -  forward declarations */
 extern List * CreatePolicyCommands(Oid relationId);

src/include/distributed/deparser.h

@@ -209,6 +209,7 @@ extern void QualifyAlterRoleSetStmt(Node *stmt);
 extern char * DeparseCreateRoleStmt(Node *stmt);
 extern char * DeparseDropRoleStmt(Node *stmt);
 extern char * DeparseGrantRoleStmt(Node *stmt);
+extern char * DeparseReassignOwnedStmt(Node *node);
 
 /* forward declarations for deparse_owned_stmts.c */
 extern char * DeparseDropOwnedStmt(Node *node);

src/test/regress/expected/reassure_owned.out

@@ -0,0 +1,93 @@
+CREATE ROLE role1;
+CREATE ROLE role2;
+GRANT CREATE ON SCHEMA public TO role1;
+SET ROLE role1;
+CREATE TABLE public.test_table (col1 int);
+RESET ROLE;
+select create_distributed_table('test_table', 'col1');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+SELECT result from run_command_on_all_nodes(
+  $$
+  SELECT jsonb_agg(to_jsonb(q2.*)) FROM (
+   SELECT
+    schemaname,
+    tablename,
+    tableowner
+FROM
+    pg_tables
+WHERE
+    tablename = 'test_table'
+  ) q2
+  $$
+) ORDER BY result;
+                                    result
+---------------------------------------------------------------------
+ [{"tablename": "test_table", "schemaname": "public", "tableowner": "role1"}]
+ [{"tablename": "test_table", "schemaname": "public", "tableowner": "role1"}]
+ [{"tablename": "test_table", "schemaname": "public", "tableowner": "role1"}]
+(3 rows)
+
+set citus.log_remote_commands to on;
+set citus.grep_remote_commands = '%REASSIGN OWNED BY%';
+REASSIGN OWNED BY role1 TO role2;
+NOTICE:  issuing REASSIGN OWNED BY role1 TO role2
+DETAIL:  on server postgres@localhost:xxxxx connectionId: xxxxxxx
+NOTICE:  issuing REASSIGN OWNED BY role1 TO role2
+DETAIL:  on server postgres@localhost:xxxxx connectionId: xxxxxxx
+SET citus.log_remote_commands = false;
+SELECT result from run_command_on_all_nodes(
+  $$
+  SELECT jsonb_agg(to_jsonb(q2.*)) FROM (
+   SELECT
+    schemaname,
+    tablename,
+    tableowner
+FROM
+    pg_tables
+WHERE
+    tablename = 'test_table'
+  ) q2
+  $$
+) ORDER BY result;
+                                    result
+---------------------------------------------------------------------
+ [{"tablename": "test_table", "schemaname": "public", "tableowner": "role2"}]
+ [{"tablename": "test_table", "schemaname": "public", "tableowner": "role2"}]
+ [{"tablename": "test_table", "schemaname": "public", "tableowner": "role2"}]
+(3 rows)
+
+SET citus.log_remote_commands = true;
+DROP OWNED BY role1, role2;
+SET citus.log_remote_commands = false;
+SELECT result from run_command_on_all_nodes(
+  $$
+  SELECT jsonb_agg(to_jsonb(q2.*)) FROM (
+   SELECT
+    schemaname,
+    tablename,
+    tableowner
+FROM
+    pg_tables
+WHERE
+    tablename = 'test_table'
+  ) q2
+  $$
+) ORDER BY result;
+ result
+---------------------------------------------------------------------
+
+
+
+(3 rows)
+
+SET citus.log_remote_commands = true;
+set citus.grep_remote_commands = '%DROP ROLE%';
+drop role role1, role2 ;
+NOTICE:  issuing DROP ROLE role1, role2
+DETAIL:  on server postgres@localhost:xxxxx connectionId: xxxxxxx
+NOTICE:  issuing DROP ROLE role1, role2
+DETAIL:  on server postgres@localhost:xxxxx connectionId: xxxxxxx

src/test/regress/multi_1_schedule

@@ -54,6 +54,7 @@ test: grant_on_database_propagation
 test: alter_database_propagation
 
 test: citus_shards
+test: reassure_owned
 
 # ----------
 # multi_citus_tools tests utility functions written for citus tools

src/test/regress/sql/reassure_owned.sql

@@ -0,0 +1,70 @@
+CREATE ROLE role1;
+CREATE ROLE role2;
+
+GRANT CREATE ON SCHEMA public TO role1;
+
+SET ROLE role1;
+CREATE TABLE public.test_table (col1 int);
+RESET ROLE;
+select create_distributed_table('test_table', 'col1');
+
+SELECT result from run_command_on_all_nodes(
+  $$
+  SELECT jsonb_agg(to_jsonb(q2.*)) FROM (
+   SELECT
+    schemaname,
+    tablename,
+    tableowner
+FROM
+    pg_tables
+WHERE
+    tablename = 'test_table'
+  ) q2
+  $$
+) ORDER BY result;
+
+set citus.log_remote_commands to on;
+set citus.grep_remote_commands = '%REASSIGN OWNED BY%';
+REASSIGN OWNED BY role1 TO role2;
+
+
+SET citus.log_remote_commands = false;
+SELECT result from run_command_on_all_nodes(
+  $$
+  SELECT jsonb_agg(to_jsonb(q2.*)) FROM (
+   SELECT
+    schemaname,
+    tablename,
+    tableowner
+FROM
+    pg_tables
+WHERE
+    tablename = 'test_table'
+  ) q2
+  $$
+) ORDER BY result;
+
+
+
+SET citus.log_remote_commands = true;
+DROP OWNED BY role1, role2;
+
+SET citus.log_remote_commands = false;
+SELECT result from run_command_on_all_nodes(
+  $$
+  SELECT jsonb_agg(to_jsonb(q2.*)) FROM (
+   SELECT
+    schemaname,
+    tablename,
+    tableowner
+FROM
+    pg_tables
+WHERE
+    tablename = 'test_table'
+  ) q2
+  $$
+) ORDER BY result;
+
+SET citus.log_remote_commands = true;
+set citus.grep_remote_commands = '%DROP ROLE%';
+drop role role1, role2 ;