Adds Grant Role support from non-main db

src/backend/distributed/commands/utility_hook.c

@@ -94,6 +94,37 @@
 #define MARK_OBJECT_DISTRIBUTED \
 	"SELECT citus_internal.mark_object_distributed(%d, %s, %d, %s)"
 
+/*
+ * NonMainDbDistributedStatementInfo is used to determine whether a statement is
+ * supported from non-main databases and whether it should be marked as
+ * distributed explicitly (*).
+ *
+ * (*) We always have to mark such objects as "distributed" but while for some
+ * object types we can delegate this to main database, for some others we have
+ * to explicitly send a command to all nodes in this code-path to achieve this.
+ */
+typedef struct NonMainDbDistributedStatementInfo
+{
+	int statementType;
+	bool explicitlyMarkAsDistributed;
+} NonMainDbDistributedStatementInfo;
+
+typedef struct MarkObjectDistributedParams
+{
+	char *name;
+	Oid id;
+	uint16 catalogRelId;
+} MarkObjectDistributedParams;
+
+/*
+ * NonMainDbSupportedStatements is an array of statements that are supported
+ * from non-main databases.
+ */
+static const NonMainDbDistributedStatementInfo NonMainDbSupportedStatements[] = {
+	{ T_GrantRoleStmt, false },
+	{ T_CreateRoleStmt, true }
+};
+
 
 bool EnableDDLPropagation = true; /* ddl propagation is enabled */
 int CreateObjectPropagationMode = CREATE_OBJECT_PROPAGATION_IMMEDIATE;
@@ -122,8 +153,12 @@
 static void DecrementUtilityHookCountersIfNecessary(Node *parsetree);
 static bool IsDropSchemaOrDB(Node *parsetree);
 static bool ShouldCheckUndistributeCitusLocalTables(void);
-static void RunPreprocessMainDBCommand(Node *parsetree, const char *queryString);
+static void RunPreprocessMainDBCommand(Node *parsetree);
 static void RunPostprocessMainDBCommand(Node *parsetree);
+static bool IsStatementSupportedFromNonMainDb(Node *parsetree);
+static bool StatementRequiresMarkDistributedFromNonMainDb(Node *parsetree);
+static void MarkObjectDistributedFromNonMainDb(Node *parsetree);
+static MarkObjectDistributedParams GetMarkObjectDistributedParams(Node *parsetree);
 
 /*
  * ProcessUtilityParseTree is a convenience method to create a PlannedStmt out of
@@ -257,7 +292,7 @@
 	{
 		if (!IsMainDB)
 		{
-			RunPreprocessMainDBCommand(parsetree, queryString);
+			RunPreprocessMainDBCommand(parsetree);
 		}
 
 		/*
@@ -1608,22 +1643,25 @@
  * database before query is run on the local node with PrevProcessUtility
  */
 static void
-RunPreprocessMainDBCommand(Node *parsetree, const char *queryString)
+RunPreprocessMainDBCommand(Node *parsetree)
 {
-	if (IsA(parsetree, CreateRoleStmt))
+	if (!IsStatementSupportedFromNonMainDb(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;
 	}
+
+	char *queryString = DeparseTreeNode(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);
 }
 
 
@@ -1633,18 +1671,99 @@
  */
 static void
 RunPostprocessMainDBCommand(Node *parsetree)
+{
+	if (IsStatementSupportedFromNonMainDb(parsetree) &&
+		StatementRequiresMarkDistributedFromNonMainDb(parsetree))
+	{
+		MarkObjectDistributedFromNonMainDb(parsetree);
+	}
+}
+
+
+/*
+ * IsStatementSupportedFromNonMainDb returns true if the statement is supported from a
+ * non-main database.
+ */
+static bool
+IsStatementSupportedFromNonMainDb(Node *parsetree)
+{
+	NodeTag type = nodeTag(parsetree);
+
+	for (int i = 0; i < sizeof(NonMainDbSupportedStatements) /
+		 sizeof(NonMainDbSupportedStatements[0]); i++)
+	{
+		if (type == NonMainDbSupportedStatements[i].statementType)
+		{
+			return true;
+		}
+	}
+
+	return false;
+}
+
+
+/*
+ * StatementRequiresMarkDistributedFromNonMainDb returns true if the statement should be marked
+ * as distributed when executed from a non-main database.
+ */
+static bool
+StatementRequiresMarkDistributedFromNonMainDb(Node *parsetree)
+{
+	NodeTag type = nodeTag(parsetree);
+
+	for (int i = 0; i < sizeof(NonMainDbSupportedStatements) /
+		 sizeof(NonMainDbSupportedStatements[0]); i++)
+	{
+		if (type == NonMainDbSupportedStatements[i].statementType)
+		{
+			return NonMainDbSupportedStatements[i].explicitlyMarkAsDistributed;
+		}
+	}
+
+	return false;
+}
+
+
+/*
+ * MarkObjectDistributedFromNonMainDb marks the given object as distributed on the
+ * non-main database.
+ */
+static void
+MarkObjectDistributedFromNonMainDb(Node *parsetree)
+{
+	MarkObjectDistributedParams markObjectDistributedParams =
+		GetMarkObjectDistributedParams(parsetree);
+	StringInfo mainDBQuery = makeStringInfo();
+	appendStringInfo(mainDBQuery,
+					 MARK_OBJECT_DISTRIBUTED,
+					 markObjectDistributedParams.catalogRelId,
+					 quote_literal_cstr(markObjectDistributedParams.name),
+					 markObjectDistributedParams.id,
+					 quote_literal_cstr(CurrentUserName()));
+	RunCitusMainDBQuery(mainDBQuery->data);
+}
+
+
+/*
+ * GetMarkObjectDistributedParams returns MarkObjectDistributedParams for the target
+ * object of given parsetree.
+ */
+static MarkObjectDistributedParams
+GetMarkObjectDistributedParams(Node *parsetree)
 {
 	if (IsA(parsetree, CreateRoleStmt))
 	{
-		StringInfo mainDBQuery = makeStringInfo();
-		CreateRoleStmt *createRoleStmt = castNode(CreateRoleStmt, parsetree);
-		Oid roleOid = get_role_oid(createRoleStmt->role, false);
-		appendStringInfo(mainDBQuery,
-						 MARK_OBJECT_DISTRIBUTED,
-						 AuthIdRelationId,
-						 quote_literal_cstr(createRoleStmt->role),
-						 roleOid,
-						 quote_literal_cstr(CurrentUserName()));
-		RunCitusMainDBQuery(mainDBQuery->data);
+		CreateRoleStmt *stmt = castNode(CreateRoleStmt, parsetree);
+		MarkObjectDistributedParams info = {
+			.name = stmt->role,
+			.catalogRelId = AuthIdRelationId,
+			.id = get_role_oid(stmt->role, false)
+		};
+
+		return info;
 	}
+
+	/* Add else if branches for other statement types */
+
+	elog(ERROR, "unsupported statement type");
 }

src/test/regress/expected/grant_role_from_non_maindb.out

@@ -0,0 +1,160 @@
+CREATE SCHEMA grant_role2pc;
+SET search_path TO grant_role2pc;
+set citus.enable_create_database_propagation to on;
+CREATE DATABASE grant_role2pc_db;
+\c grant_role2pc_db
+SHOW citus.main_db;
+ citus.main_db
+---------------------------------------------------------------------
+ regression
+(1 row)
+
+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 grant_role2pc_db
+--test with empty superuser
+SET citus.superuser TO '';
+grant grant_role2pc_user1 to grant_role2pc_user2;
+ERROR:  No superuser role is given for Citus main database connection
+HINT:  Set citus.superuser to a superuser role name
+SET citus.superuser TO 'postgres';
+grant grant_role2pc_user1 to grant_role2pc_user2 with admin option granted by CURRENT_USER;
+\c regression
+select result FROM run_command_on_all_nodes(
+    $$
+    SELECT array_to_json(array_agg(row_to_json(t)))
+    FROM (
+    SELECT member::regrole, roleid::regrole as role, grantor::regrole, admin_option
+    FROM pg_auth_members
+    WHERE member::regrole::text = 'grant_role2pc_user2'
+    order by member::regrole::text, roleid::regrole::text
+    ) t
+    $$
+);
+                                                  result
+---------------------------------------------------------------------
+ [{"member":"grant_role2pc_user2","role":"grant_role2pc_user1","grantor":"postgres","admin_option":true}]
+ [{"member":"grant_role2pc_user2","role":"grant_role2pc_user1","grantor":"postgres","admin_option":true}]
+ [{"member":"grant_role2pc_user2","role":"grant_role2pc_user1","grantor":"postgres","admin_option":true}]
+(3 rows)
+
+\c grant_role2pc_db
+--test grant under transactional context with multiple operations
+BEGIN;
+grant grant_role2pc_user1,grant_role2pc_user2 to grant_role2pc_user3 WITH ADMIN OPTION;
+grant grant_role2pc_user1 to grant_role2pc_user4 granted by grant_role2pc_user3 ;
+COMMIT;
+BEGIN;
+grant grant_role2pc_user1 to grant_role2pc_user5 WITH ADMIN OPTION granted by grant_role2pc_user3;
+grant grant_role2pc_user1 to grant_role2pc_user6;
+ROLLBACK;
+BEGIN;
+grant grant_role2pc_user1 to grant_role2pc_user7;
+SELECT 1/0;
+ERROR:  division by zero
+commit;
+\c regression
+select result FROM run_command_on_all_nodes($$
+SELECT array_to_json(array_agg(row_to_json(t)))
+FROM (
+    SELECT member::regrole, roleid::regrole as role, grantor::regrole, admin_option
+    FROM pg_auth_members
+    WHERE member::regrole::text in
+        ('grant_role2pc_user3','grant_role2pc_user4','grant_role2pc_user5','grant_role2pc_user6','grant_role2pc_user7')
+    order by member::regrole::text, roleid::regrole::text
+) t
+$$);
+                                                                                                                                                               result
+---------------------------------------------------------------------
+ [{"member":"grant_role2pc_user3","role":"grant_role2pc_user1","grantor":"postgres","admin_option":true},{"member":"grant_role2pc_user3","role":"grant_role2pc_user2","grantor":"postgres","admin_option":true},{"member":"grant_role2pc_user4","role":"grant_role2pc_user1","grantor":"grant_role2pc_user3","admin_option":false}]
+ [{"member":"grant_role2pc_user3","role":"grant_role2pc_user1","grantor":"postgres","admin_option":true},{"member":"grant_role2pc_user3","role":"grant_role2pc_user2","grantor":"postgres","admin_option":true},{"member":"grant_role2pc_user4","role":"grant_role2pc_user1","grantor":"grant_role2pc_user3","admin_option":false}]
+ [{"member":"grant_role2pc_user3","role":"grant_role2pc_user1","grantor":"postgres","admin_option":true},{"member":"grant_role2pc_user3","role":"grant_role2pc_user2","grantor":"postgres","admin_option":true},{"member":"grant_role2pc_user4","role":"grant_role2pc_user1","grantor":"grant_role2pc_user3","admin_option":false}]
+(3 rows)
+
+\c grant_role2pc_db
+grant grant_role2pc_user1,grant_role2pc_user2 to grant_role2pc_user5,grant_role2pc_user6,grant_role2pc_user7 granted by grant_role2pc_user3;
+\c regression
+select result FROM run_command_on_all_nodes($$
+SELECT array_to_json(array_agg(row_to_json(t)))
+FROM (
+    SELECT member::regrole, roleid::regrole as role, grantor::regrole, admin_option
+    FROM pg_auth_members
+    WHERE member::regrole::text in
+        ('grant_role2pc_user5','grant_role2pc_user6','grant_role2pc_user7')
+    order by member::regrole::text, roleid::regrole::text
+) t
+$$);
+                                                                                                                                                                                                                                                                                                                                                       result
+---------------------------------------------------------------------
+ [{"member":"grant_role2pc_user5","role":"grant_role2pc_user1","grantor":"grant_role2pc_user3","admin_option":false},{"member":"grant_role2pc_user5","role":"grant_role2pc_user2","grantor":"grant_role2pc_user3","admin_option":false},{"member":"grant_role2pc_user6","role":"grant_role2pc_user1","grantor":"grant_role2pc_user3","admin_option":false},{"member":"grant_role2pc_user6","role":"grant_role2pc_user2","grantor":"grant_role2pc_user3","admin_option":false},{"member":"grant_role2pc_user7","role":"grant_role2pc_user1","grantor":"grant_role2pc_user3","admin_option":false},{"member":"grant_role2pc_user7","role":"grant_role2pc_user2","grantor":"grant_role2pc_user3","admin_option":false}]
+ [{"member":"grant_role2pc_user5","role":"grant_role2pc_user1","grantor":"grant_role2pc_user3","admin_option":false},{"member":"grant_role2pc_user5","role":"grant_role2pc_user2","grantor":"grant_role2pc_user3","admin_option":false},{"member":"grant_role2pc_user6","role":"grant_role2pc_user1","grantor":"grant_role2pc_user3","admin_option":false},{"member":"grant_role2pc_user6","role":"grant_role2pc_user2","grantor":"grant_role2pc_user3","admin_option":false},{"member":"grant_role2pc_user7","role":"grant_role2pc_user1","grantor":"grant_role2pc_user3","admin_option":false},{"member":"grant_role2pc_user7","role":"grant_role2pc_user2","grantor":"grant_role2pc_user3","admin_option":false}]
+ [{"member":"grant_role2pc_user5","role":"grant_role2pc_user1","grantor":"grant_role2pc_user3","admin_option":false},{"member":"grant_role2pc_user5","role":"grant_role2pc_user2","grantor":"grant_role2pc_user3","admin_option":false},{"member":"grant_role2pc_user6","role":"grant_role2pc_user1","grantor":"grant_role2pc_user3","admin_option":false},{"member":"grant_role2pc_user6","role":"grant_role2pc_user2","grantor":"grant_role2pc_user3","admin_option":false},{"member":"grant_role2pc_user7","role":"grant_role2pc_user1","grantor":"grant_role2pc_user3","admin_option":false},{"member":"grant_role2pc_user7","role":"grant_role2pc_user2","grantor":"grant_role2pc_user3","admin_option":false}]
+(3 rows)
+
+\c grant_role2pc_db
+revoke admin option for grant_role2pc_user1 from grant_role2pc_user5 granted by grant_role2pc_user3;
+--test revoke under transactional context with multiple operations
+BEGIN;
+revoke grant_role2pc_user1 from grant_role2pc_user5 granted by grant_role2pc_user3 ;
+revoke grant_role2pc_user1 from grant_role2pc_user4 granted by grant_role2pc_user3;
+COMMIT;
+\c grant_role2pc_db - - :worker_1_port
+BEGIN;
+revoke grant_role2pc_user1 from grant_role2pc_user6,grant_role2pc_user7 granted by grant_role2pc_user3;
+revoke grant_role2pc_user1 from grant_role2pc_user3 cascade;
+COMMIT;
+\c regression
+select result FROM run_command_on_all_nodes($$
+SELECT array_to_json(array_agg(row_to_json(t)))
+FROM (
+    SELECT member::regrole, roleid::regrole as role, grantor::regrole, admin_option
+    FROM pg_auth_members
+    WHERE member::regrole::text in
+        ('grant_role2pc_user2','grant_role2pc_user3','grant_role2pc_user4','grant_role2pc_user5','grant_role2pc_user6','grant_role2pc_user7')
+    order by member::regrole::text, roleid::regrole::text
+) t
+$$);
+                                                                                                                                                                                                                                                                                  result
+---------------------------------------------------------------------
+ [{"member":"grant_role2pc_user2","role":"grant_role2pc_user1","grantor":"postgres","admin_option":true},{"member":"grant_role2pc_user3","role":"grant_role2pc_user2","grantor":"postgres","admin_option":true},{"member":"grant_role2pc_user5","role":"grant_role2pc_user2","grantor":"grant_role2pc_user3","admin_option":false},{"member":"grant_role2pc_user6","role":"grant_role2pc_user2","grantor":"grant_role2pc_user3","admin_option":false},{"member":"grant_role2pc_user7","role":"grant_role2pc_user2","grantor":"grant_role2pc_user3","admin_option":false}]
+ [{"member":"grant_role2pc_user2","role":"grant_role2pc_user1","grantor":"postgres","admin_option":true},{"member":"grant_role2pc_user3","role":"grant_role2pc_user2","grantor":"postgres","admin_option":true},{"member":"grant_role2pc_user5","role":"grant_role2pc_user2","grantor":"grant_role2pc_user3","admin_option":false},{"member":"grant_role2pc_user6","role":"grant_role2pc_user2","grantor":"grant_role2pc_user3","admin_option":false},{"member":"grant_role2pc_user7","role":"grant_role2pc_user2","grantor":"grant_role2pc_user3","admin_option":false}]
+ [{"member":"grant_role2pc_user2","role":"grant_role2pc_user1","grantor":"postgres","admin_option":true},{"member":"grant_role2pc_user3","role":"grant_role2pc_user2","grantor":"postgres","admin_option":true},{"member":"grant_role2pc_user5","role":"grant_role2pc_user2","grantor":"grant_role2pc_user3","admin_option":false},{"member":"grant_role2pc_user6","role":"grant_role2pc_user2","grantor":"grant_role2pc_user3","admin_option":false},{"member":"grant_role2pc_user7","role":"grant_role2pc_user2","grantor":"grant_role2pc_user3","admin_option":false}]
+(3 rows)
+
+\c grant_role2pc_db - - :worker_1_port
+BEGIN;
+grant grant_role2pc_user1 to grant_role2pc_user5 WITH ADMIN OPTION;
+grant grant_role2pc_user1 to grant_role2pc_user6;
+COMMIT;
+\c regression - - :master_port
+select result FROM run_command_on_all_nodes($$
+SELECT array_to_json(array_agg(row_to_json(t)))
+FROM (
+    SELECT member::regrole, roleid::regrole as role, grantor::regrole, admin_option
+    FROM pg_auth_members
+    WHERE member::regrole::text in
+        ('grant_role2pc_user5','grant_role2pc_user6')
+    order by member::regrole::text, roleid::regrole::text
+) t
+$$);
+                                                                                                                                                                                                                         result
+---------------------------------------------------------------------
+ [{"member":"grant_role2pc_user5","role":"grant_role2pc_user1","grantor":"postgres","admin_option":true},{"member":"grant_role2pc_user5","role":"grant_role2pc_user2","grantor":"grant_role2pc_user3","admin_option":false},{"member":"grant_role2pc_user6","role":"grant_role2pc_user1","grantor":"postgres","admin_option":false},{"member":"grant_role2pc_user6","role":"grant_role2pc_user2","grantor":"grant_role2pc_user3","admin_option":false}]
+ [{"member":"grant_role2pc_user5","role":"grant_role2pc_user1","grantor":"postgres","admin_option":true},{"member":"grant_role2pc_user5","role":"grant_role2pc_user2","grantor":"grant_role2pc_user3","admin_option":false},{"member":"grant_role2pc_user6","role":"grant_role2pc_user1","grantor":"postgres","admin_option":false},{"member":"grant_role2pc_user6","role":"grant_role2pc_user2","grantor":"grant_role2pc_user3","admin_option":false}]
+ [{"member":"grant_role2pc_user5","role":"grant_role2pc_user1","grantor":"postgres","admin_option":true},{"member":"grant_role2pc_user5","role":"grant_role2pc_user2","grantor":"grant_role2pc_user3","admin_option":false},{"member":"grant_role2pc_user6","role":"grant_role2pc_user1","grantor":"postgres","admin_option":false},{"member":"grant_role2pc_user6","role":"grant_role2pc_user2","grantor":"grant_role2pc_user3","admin_option":false}]
+(3 rows)
+
+revoke grant_role2pc_user1 from grant_role2pc_user5,grant_role2pc_user6;
+--clean resources
+DROP SCHEMA grant_role2pc;
+set citus.enable_create_database_propagation to on;
+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;
+reset citus.enable_create_database_propagation;