diff --git a/src/backend/distributed/commands/utility_hook.c b/src/backend/distributed/commands/utility_hook.c index 9e6b66e3e..154b9ca6c 100644 --- a/src/backend/distributed/commands/utility_hook.c +++ b/src/backend/distributed/commands/utility_hook.c @@ -43,6 +43,7 @@ #include "foreign/foreign.h" #include "lib/stringinfo.h" #include "nodes/makefuncs.h" +#include "nodes/nodes.h" #include "nodes/parsenodes.h" #include "nodes/pg_list.h" #include "postmaster/postmaster.h" @@ -95,6 +96,18 @@ "SELECT citus_internal.mark_object_distributed(%d, %s, %d)" +typedef struct TwoPcStatementInfo +{ + int statementType; + bool markAsDistributed; +} TwoPcStatementInfo; + +const TwoPcStatementInfo twoPcSupportedStatements[] = { + { T_GrantRoleStmt, false }, + { T_CreateRoleStmt, true } +}; + + bool EnableDDLPropagation = true; /* ddl propagation is enabled */ int CreateObjectPropagationMode = CREATE_OBJECT_PROPAGATION_IMMEDIATE; PropSetCmdBehavior PropagateSetCommands = PROPSETCMD_NONE; /* SET prop off */ @@ -124,6 +137,8 @@ static bool IsDropSchemaOrDB(Node *parsetree); static bool ShouldCheckUndistributeCitusLocalTables(void); static void RunPreprocessMainDBCommand(Node *parsetree, const char *queryString); static void RunPostprocessMainDBCommand(Node *parsetree); +static bool IsStatementSupportedIn2Pc(Node *parsetree); +static bool IsStatementMarkDistributedFor2PC(Node *parsetree); /* * ProcessUtilityParseTree is a convenience method to create a PlannedStmt out of @@ -1603,20 +1618,22 @@ DropSchemaOrDBInProgress(void) static void RunPreprocessMainDBCommand(Node *parsetree, const char *queryString) { - if (IsA(parsetree, CreateRoleStmt)) + if (!IsStatementSupportedIn2Pc(parsetree)) { - StringInfo mainDBQuery = makeStringInfo(); - appendStringInfo(mainDBQuery, - START_MANAGEMENT_TRANSACTION, - GetCurrentFullTransactionId().value); - RunCitusMainDBQuery(mainDBQuery->data); - mainDBQuery = makeStringInfo(); - appendStringInfo(mainDBQuery, - EXECUTE_COMMAND_ON_REMOTE_NODES_AS_USER, - quote_literal_cstr(queryString), - quote_literal_cstr(CurrentUserName())); - RunCitusMainDBQuery(mainDBQuery->data); + return; } + + StringInfo mainDBQuery = makeStringInfo(); + appendStringInfo(mainDBQuery, + START_MANAGEMENT_TRANSACTION, + GetCurrentFullTransactionId().value); + RunCitusMainDBQuery(mainDBQuery->data); + mainDBQuery = makeStringInfo(); + appendStringInfo(mainDBQuery, + EXECUTE_COMMAND_ON_REMOTE_NODES_AS_USER, + quote_literal_cstr(queryString), + quote_literal_cstr(CurrentUserName())); + RunCitusMainDBQuery(mainDBQuery->data); } @@ -1627,6 +1644,12 @@ RunPreprocessMainDBCommand(Node *parsetree, const char *queryString) static void RunPostprocessMainDBCommand(Node *parsetree) { + if (!IsStatementSupportedIn2Pc(parsetree) || + !IsStatementMarkDistributedFor2PC(parsetree)) + { + return; + } + if (IsA(parsetree, CreateRoleStmt)) { StringInfo mainDBQuery = makeStringInfo(); @@ -1640,3 +1663,45 @@ RunPostprocessMainDBCommand(Node *parsetree) RunCitusMainDBQuery(mainDBQuery->data); } } + +/* + * IsStatementSupportedIn2Pc returns true if the statement is supported in 2pc + */ + +static bool +IsStatementSupportedIn2Pc(Node *parsetree) +{ + NodeTag type = nodeTag(parsetree); + + for (int i = 0; i < sizeof(twoPcSupportedStatements) / + sizeof(twoPcSupportedStatements[0]); i++) + { + if (type == twoPcSupportedStatements[i].statementType) + { + return true; + } + } + + return false; +} + +/* + * IsStatementMarkDistributedFor2PC returns true if the statement should be marked + * as distributed in 2pc + */ +static bool +IsStatementMarkDistributedFor2PC(Node *parsetree) +{ + NodeTag type = nodeTag(parsetree); + + for (int i = 0; i < sizeof(twoPcSupportedStatements) / + sizeof(twoPcSupportedStatements[0]); i++) + { + if (type == twoPcSupportedStatements[i].statementType) + { + return twoPcSupportedStatements[i].markAsDistributed; + } + } + + return false; +} diff --git a/src/test/regress/expected/grant_role_2pc.out b/src/test/regress/expected/grant_role_2pc.out new file mode 100644 index 000000000..e69de29bb diff --git a/src/test/regress/multi_schedule b/src/test/regress/multi_schedule index 5c9d8a45c..aac6464ae 100644 --- a/src/test/regress/multi_schedule +++ b/src/test/regress/multi_schedule @@ -109,6 +109,7 @@ test: undistribute_table test: run_command_on_all_nodes test: background_task_queue_monitor test: other_databases +test: grant_role_2pc # Causal clock test test: clock diff --git a/src/test/regress/sql/grant_role_2pc.sql b/src/test/regress/sql/grant_role_2pc.sql new file mode 100644 index 000000000..55e4b2e8b --- /dev/null +++ b/src/test/regress/sql/grant_role_2pc.sql @@ -0,0 +1,131 @@ + + +CREATE SCHEMA grant_role2pc; + +SET search_path TO grant_role2pc; + +set citus.enable_create_database_propagation to on; + + +set citus.log_remote_commands to on; + + +SET citus.next_shard_id TO 10231023; + +CREATE DATABASE grant_role2pc_db; + +revoke connect,temp,temporary,create on database grant_role2pc_db from public; + +\c grant_role2pc_db +SHOW citus.main_db; + +-- check that empty citus.superuser gives error +SET citus.superuser TO ''; +CREATE USER empty_superuser; +SET citus.superuser TO 'postgres'; + +CREATE USER grant_role2pc_user1; +CREATE USER grant_role2pc_user2; +CREATE USER grant_role2pc_user3; +CREATE USER grant_role2pc_user4; +CREATE USER grant_role2pc_user5; +CREATE USER grant_role2pc_user6; +CREATE USER grant_role2pc_user7; + + +\c regression + +SELECT * FROM public.check_database_privileges('grant_role2pc_user2', 'grant_role2pc_db', ARRAY['CREATE', 'CONNECT', 'TEMP', 'TEMPORARY']); + +grant create,connect,temporary,temp on database grant_role2pc_db to grant_role2pc_user1; + +\c grant_role2pc_db + +grant grant_role2pc_user1 to grant_role2pc_user2; + +\c regression + + +SELECT * FROM public.check_database_privileges('grant_role2pc_user2', 'grant_role2pc_db', ARRAY['CREATE', 'CONNECT', 'TEMP', 'TEMPORARY']); + + +\c grant_role2pc_db +--test grant under transactional context with multiple operations +BEGIN; +grant grant_role2pc_user1 to grant_role2pc_user3; +grant grant_role2pc_user1 to grant_role2pc_user4; +COMMIT; + +BEGIN; +grant grant_role2pc_user1 to grant_role2pc_user5; +grant grant_role2pc_user1 to grant_role2pc_user6; +ROLLBACK; + + + +BEGIN; +grant grant_role2pc_user1 to grant_role2pc_user7; +SELECT 1/0; +commit; + + +\c regression + +SELECT * FROM public.check_database_privileges('grant_role2pc_user3', 'grant_role2pc_db', ARRAY['CREATE', 'CONNECT', 'TEMP', 'TEMPORARY']); +SELECT * FROM public.check_database_privileges('grant_role2pc_user4', 'grant_role2pc_db', ARRAY['CREATE', 'CONNECT', 'TEMP', 'TEMPORARY']); +SELECT * FROM public.check_database_privileges('grant_role2pc_user5', 'grant_role2pc_db', ARRAY['CREATE', 'CONNECT', 'TEMP', 'TEMPORARY']); +SELECT * FROM public.check_database_privileges('grant_role2pc_user6', 'grant_role2pc_db', ARRAY['CREATE', 'CONNECT', 'TEMP', 'TEMPORARY']); +SELECT * FROM public.check_database_privileges('grant_role2pc_user7', 'grant_role2pc_db', ARRAY['CREATE', 'CONNECT', 'TEMP', 'TEMPORARY']); + +\c grant_role2pc_db + +grant grant_role2pc_user1 to grant_role2pc_user5,grant_role2pc_user6,grant_role2pc_user7; + +\c regression +SELECT * FROM public.check_database_privileges('grant_role2pc_user5', 'grant_role2pc_db', ARRAY['CREATE', 'CONNECT', 'TEMP', 'TEMPORARY']); +SELECT * FROM public.check_database_privileges('grant_role2pc_user6', 'grant_role2pc_db', ARRAY['CREATE', 'CONNECT', 'TEMP', 'TEMPORARY']); +SELECT * FROM public.check_database_privileges('grant_role2pc_user7', 'grant_role2pc_db', ARRAY['CREATE', 'CONNECT', 'TEMP', 'TEMPORARY']); + +\c grant_role2pc_db +revoke grant_role2pc_user1 from grant_role2pc_user2; + +--test revoke under transactional context with multiple operations +BEGIN; +revoke grant_role2pc_user1 from grant_role2pc_user3; +revoke grant_role2pc_user1 from grant_role2pc_user4; +COMMIT; + +BEGIN; +revoke grant_role2pc_user1 from grant_role2pc_user5,grant_role2pc_user6; +revoke grant_role2pc_user1 from grant_role2pc_user7; +COMMIT; + +\c regression + +SELECT * FROM public.check_database_privileges('grant_role2pc_user2', 'grant_role2pc_db', ARRAY['CREATE', 'CONNECT', 'TEMP', 'TEMPORARY']); +SELECT * FROM public.check_database_privileges('grant_role2pc_user3', 'grant_role2pc_db', ARRAY['CREATE', 'CONNECT', 'TEMP', 'TEMPORARY']); +SELECT * FROM public.check_database_privileges('grant_role2pc_user4', 'grant_role2pc_db', ARRAY['CREATE', 'CONNECT', 'TEMP', 'TEMPORARY']); +SELECT * FROM public.check_database_privileges('grant_role2pc_user5', 'grant_role2pc_db', ARRAY['CREATE', 'CONNECT', 'TEMP', 'TEMPORARY']); +SELECT * FROM public.check_database_privileges('grant_role2pc_user6', 'grant_role2pc_db', ARRAY['CREATE', 'CONNECT', 'TEMP', 'TEMPORARY']); +SELECT * FROM public.check_database_privileges('grant_role2pc_user7', 'grant_role2pc_db', ARRAY['CREATE', 'CONNECT', 'TEMP', 'TEMPORARY']); + + + + + + + + +DROP SCHEMA grant_role2pc; + +REVOKE ALL PRIVILEGES ON DATABASE grant_role2pc_db FROM grant_role2pc_user1; +DROP DATABASE grant_role2pc_db; + +drop user grant_role2pc_user2,grant_role2pc_user3,grant_role2pc_user4,grant_role2pc_user5,grant_role2pc_user6,grant_role2pc_user7; + + +drop user grant_role2pc_user1; + +grant connect,temp,temporary on database regression to public; + +reset citus.enable_create_database_propagation; diff --git a/src/test/regress/sql/multi_test_helpers.sql b/src/test/regress/sql/multi_test_helpers.sql index e67b782a5..40bbaaf07 100644 --- a/src/test/regress/sql/multi_test_helpers.sql +++ b/src/test/regress/sql/multi_test_helpers.sql @@ -652,3 +652,16 @@ BEGIN JOIN pg_dist_node USING (nodeid); END; $func$ LANGUAGE plpgsql; + +CREATE OR REPLACE FUNCTION check_database_privileges(role_name text, db_name text, permissions text[]) +RETURNS TABLE(permission text, result text) +AS $func$ +DECLARE + permission text; +BEGIN + FOREACH permission IN ARRAY permissions + LOOP + RETURN QUERY EXECUTE format($inner$SELECT '%s', result FROM run_command_on_all_nodes($$select has_database_privilege('%s','%s', '%s'); $$)$inner$, permission, role_name, db_name, permission); + END LOOP; +END; +$func$ LANGUAGE plpgsql;