Adds Grant Role support into 2PC

grant_database_2pc_onur_1
gurkanindibay 2024-01-05 13:12:21 +03:00
parent d940cfa992
commit 20dc90289b
5 changed files with 222 additions and 12 deletions

View File

@ -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,8 +1618,11 @@ DropSchemaOrDBInProgress(void)
static void
RunPreprocessMainDBCommand(Node *parsetree, const char *queryString)
{
if (IsA(parsetree, CreateRoleStmt))
if (!IsStatementSupportedIn2Pc(parsetree))
{
return;
}
StringInfo mainDBQuery = makeStringInfo();
appendStringInfo(mainDBQuery,
START_MANAGEMENT_TRANSACTION,
@ -1617,7 +1635,6 @@ RunPreprocessMainDBCommand(Node *parsetree, const char *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;
}

View File

@ -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

View File

@ -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;

View File

@ -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;