Adds Grant Role support into 2PC

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();
+		return;
-		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);
 	}
+
+	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;
+}

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

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;

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;