diff --git a/src/backend/distributed/commands/distribute_object_ops.c b/src/backend/distributed/commands/distribute_object_ops.c index a0a306e8d..f3ef31667 100644 --- a/src/backend/distributed/commands/distribute_object_ops.c +++ b/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); diff --git a/src/backend/distributed/commands/owned.c b/src/backend/distributed/commands/owned.c index c8f6a4bbe..ba90cf8fa 100644 --- a/src/backend/distributed/commands/owned.c +++ b/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); +} diff --git a/src/backend/distributed/deparser/deparse_owned_stmts.c b/src/backend/distributed/deparser/deparse_owned_stmts.c index 888071165..61dcc7cf3 100644 --- a/src/backend/distributed/deparser/deparse_owned_stmts.c +++ b/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; +} diff --git a/src/include/distributed/commands.h b/src/include/distributed/commands.h index 43429278f..0c681c118 100644 --- a/src/include/distributed/commands.h +++ b/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); diff --git a/src/include/distributed/deparser.h b/src/include/distributed/deparser.h index 95d948bc9..baccec78f 100644 --- a/src/include/distributed/deparser.h +++ b/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); diff --git a/src/test/regress/expected/reassure_owned.out b/src/test/regress/expected/reassure_owned.out new file mode 100644 index 000000000..580bfdd30 --- /dev/null +++ b/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 diff --git a/src/test/regress/multi_1_schedule b/src/test/regress/multi_1_schedule index 4e2b19795..c7093b2e5 100644 --- a/src/test/regress/multi_1_schedule +++ b/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 diff --git a/src/test/regress/sql/reassure_owned.sql b/src/test/regress/sql/reassure_owned.sql new file mode 100644 index 000000000..a45f174da --- /dev/null +++ b/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 ;