Skip to content

Commit

Permalink
Adds Grant Role support from non-main db (#7404)
Browse files Browse the repository at this point in the history 
DESCRIPTION: Adds support for distributed role-membership management
commands from the databases where Citus is not installed (`GRANT <role>
TO <role>`)

This PR also refactors the code-path that allows executing some of the
node-wide commands so that we use send deparsed query string to other
nodes instead of the `queryString` passed into utility hook.

---------

Co-authored-by: Onur Tirtir <onurcantirtir@gmail.com>
  • Loading branch information
@gurkanindibay @onurctirtir
gurkanindibay and onurctirtir committed Feb 19, 2024
1 parent 9a0cdbf commit 2cbfdbf
Show file tree
Hide file tree
Showing 9 changed files with 591 additions and 27 deletions.
169 changes: 144 additions & 25 deletions src/backend/distributed/commands/utility_hook.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -122,8 +153,12 @@ static void PostStandardProcessUtility(Node *parsetree);
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
Expand Down Expand Up @@ -257,7 +292,7 @@ citus_ProcessUtility(PlannedStmt *pstmt,
{
if (!IsMainDB)
{
RunPreprocessMainDBCommand(parsetree, queryString);
RunPreprocessMainDBCommand(parsetree);
}

/*
Expand Down Expand Up @@ -1608,22 +1643,25 @@ DropSchemaOrDBInProgress(void)
* 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);
}


Expand All @@ -1633,18 +1671,99 @@ RunPreprocessMainDBCommand(Node *parsetree, const char *queryString)
*/
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");
}
160 changes: 160 additions & 0 deletions src/test/regress/expected/grant_role_from_non_maindb.out
View file
Original file line number Diff line number Diff line change
@@ -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;
Loading

0 comments on commit 2cbfdbf

Please sign in to comment.