Merge branch 'grant_role_2pc' into grant_database_2pc

ci/check_gucs_are_alphabetically_sorted.sh

@@ -4,7 +4,22 @@ set -euo pipefail
 # shellcheck disable=SC1091
 source ci/ci_helpers.sh
 
+# Find the line that exactly matches "RegisterCitusConfigVariables(void)" in
+# shared_library_init.c. grep command returns something like
+# "934:RegisterCitusConfigVariables(void)" and we extract the line number
+# with cut.
+RegisterCitusConfigVariables_begin_linenumber=$(grep -n "^RegisterCitusConfigVariables(void)$" src/backend/distributed/shared_library_init.c | cut -d: -f1)
+
+# Consider the lines starting from $RegisterCitusConfigVariables_begin_linenumber,
+# grep the first line that starts with "}" and extract the line number with cut
+# as in the previous step.
+RegisterCitusConfigVariables_length=$(tail -n +$RegisterCitusConfigVariables_begin_linenumber src/backend/distributed/shared_library_init.c | grep -n -m 1 "^}$" | cut -d: -f1)
+
+# extract the function definition of RegisterCitusConfigVariables into a temp file
+tail -n +$RegisterCitusConfigVariables_begin_linenumber src/backend/distributed/shared_library_init.c | head -n $(($RegisterCitusConfigVariables_length)) > RegisterCitusConfigVariables_func_def.out
+
 # extract citus gucs in the form of <tab><tab>"citus.X"
-grep -P "^[\t][\t]\"citus\.[a-zA-Z_0-9]+\"" src/backend/distributed/shared_library_init.c > gucs.out
+grep -P "^[\t][\t]\"citus\.[a-zA-Z_0-9]+\"" RegisterCitusConfigVariables_func_def.out > gucs.out
 sort -c gucs.out
 rm gucs.out
+rm RegisterCitusConfigVariables_func_def.out

src/backend/distributed/commands/utility_hook.c

@@ -96,28 +96,34 @@
 	"SELECT citus_internal.mark_object_distributed(%d, %s, %d, %s)"
 
 /*
- * TwoPcStatementInfo is used to determine whether a statement is supported in 2PC
+ * NonMainDbDistributedStatementInfo is used to determine whether a statement is
- * and whether it should be marked as distributed in 2PC.
+ * supported from non-main databases and whether it should be marked as
+ * distributed explicitly (*).
+ *
+ * We always have to mark such the objects created "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 TwoPcStatementInfo
+typedef struct NonMainDbDistributedStatementInfo
 {
 	int statementType;
-	ObjectType *supportedObjectTypes;
+	bool explicitlyMarkAsDistributed;
-	int supportedObjectTypesSize;
+} NonMainDbDistributedStatementInfo;
-	bool markAsDistributed;
+
-} TwoPcStatementInfo;
+typedef struct ObjectInfo
+{
+	char *name;
+	Oid id;
+} ObjectInfo;
 
 /*
- * twoPcSupportedStatements is a list of statements that are supported in 2PC.
+ * NonMainDbSupportedStatements is an array of statements that are supported
- * The list is used to determine whether a statement is supported in 2PC and
+ * from non-main databases.
- * whether it should be marked as distributed in 2PC.
- * We use this array to avoid hardcoding the list of supported statements in
- * multiple places.
  */
 ObjectType supportedObjectTypesForGrantStmt[] = { OBJECT_DATABASE };
-TwoPcStatementInfo twoPcSupportedStatements[] = {
+static const NonMainDbDistributedStatementInfo NonMainDbSupportedStatements[] = {
-	{ T_GrantRoleStmt, NULL, 0, false },
+	{ T_GrantRoleStmt, false },
-	{ T_CreateRoleStmt, NULL, 0, true },
+	{ T_CreateRoleStmt, true }
 	{ T_GrantStmt, supportedObjectTypesForGrantStmt,
 	  sizeof(supportedObjectTypesForGrantStmt) / sizeof(ObjectType), true }
 };
@@ -152,10 +158,10 @@ 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 IsStatementSupportedInNonMainDb(Node *parsetree);
-static bool DoesStatementRequireMarkDistributedFor2PC(Node *parsetree);
+static bool StatementRequiresMarkDistributedFromNonMainDb(Node *parsetree);
-static bool IsObjectTypeSupported(Node *parsetree, TwoPcStatementInfo
+static ObjectInfo GetObjectInfo(Node *parsetree);
-								  twoPcSupportedStatement);
+static void MarkObjectDistributedInNonMainDb(Node *parsetree, ObjectInfo objectInfo);
 
 /*
  * ProcessUtilityParseTree is a convenience method to create a PlannedStmt out of
@@ -1635,7 +1641,7 @@ DropSchemaOrDBInProgress(void)
 static void
 RunPreprocessMainDBCommand(Node *parsetree, const char *queryString)
 {
-	if (!IsStatementSupportedIn2PC(parsetree))
+	if (!IsStatementSupportedInNonMainDb(parsetree))
 	{
 		return;
 	}
@@ -1661,40 +1667,69 @@ RunPreprocessMainDBCommand(Node *parsetree, const char *queryString)
 static void
 RunPostprocessMainDBCommand(Node *parsetree)
 {
-	if (!IsStatementSupportedIn2PC(parsetree) ||
+	if (!IsStatementSupportedInNonMainDb(parsetree) ||
-		!DoesStatementRequireMarkDistributedFor2PC(parsetree))
+		!StatementRequiresMarkDistributedFromNonMainDb(parsetree))
 	{
 		return;
 	}
 
-	if (IsA(parsetree, CreateRoleStmt))
+	ObjectInfo objectInfo = GetObjectInfo(parsetree);
-	{
+	MarkObjectDistributedInNonMainDb(parsetree, objectInfo);
-		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);
-	}
 }
 
 
 /*
- * IsStatementSupportedIn2Pc returns true if the statement is supported in 2pc
+ * GetObjectInfo returns the name and oid of the object in the given parsetree.
+ */
+static ObjectInfo
+GetObjectInfo(Node *parsetree)
+{
+	ObjectInfo info;
+
+	if (IsA(parsetree, CreateRoleStmt))
+	{
+		CreateRoleStmt *stmt = castNode(CreateRoleStmt, parsetree);
+		info.name = stmt->role;
+		info.id = get_role_oid(stmt->role, false);
+	}
+
+	/* Add else if branches for other statement types */
+
+	return info;
+}
+
+
+/*
+ * MarkObjectDistributedInNonMainDb marks the given object as distributed on the
+ * non-main database.
+ */
+static void
+MarkObjectDistributedInNonMainDb(Node *parsetree, ObjectInfo objectInfo)
+{
+	StringInfo mainDBQuery = makeStringInfo();
+	appendStringInfo(mainDBQuery,
+					 MARK_OBJECT_DISTRIBUTED,
+					 AuthIdRelationId,
+					 quote_literal_cstr(objectInfo.name),
+					 objectInfo.id,
+					 quote_literal_cstr(CurrentUserName()));
+	RunCitusMainDBQuery(mainDBQuery->data);
+}
+
+
+/*
+ * IsStatementSupportedIn2Pc returns true if the statement is supported from a
+ * non-main database.
  */
 static bool
-IsStatementSupportedIn2PC(Node *parsetree)
+IsStatementSupportedInNonMainDb(Node *parsetree)
 {
 	NodeTag type = nodeTag(parsetree);
 
-	for (int i = 0; i < sizeof(twoPcSupportedStatements) /
+	for (int i = 0; i < sizeof(NonMainDbSupportedStatements) /
-		 sizeof(twoPcSupportedStatements[0]); i++)
+		 sizeof(NonMainDbSupportedStatements[0]); i++)
 	{
-		if (type == twoPcSupportedStatements[i].statementType)
+		if (type == NonMainDbSupportedStatements[i].statementType)
 		{
 			if (twoPcSupportedStatements[i].supportedObjectTypes == NULL)
 			{
@@ -1741,19 +1776,19 @@ IsObjectTypeSupported(Node *parsetree, TwoPcStatementInfo twoPcSupportedStatemen
 
 /*
  * DoesStatementRequireMarkDistributedFor2PC returns true if the statement should be marked
- * as distributed in 2pc
+ * as distributed when executed from a non-main database.
  */
 static bool
-DoesStatementRequireMarkDistributedFor2PC(Node *parsetree)
+StatementRequiresMarkDistributedFromNonMainDb(Node *parsetree)
 {
 	NodeTag type = nodeTag(parsetree);
 
-	for (int i = 0; i < sizeof(twoPcSupportedStatements) /
+	for (int i = 0; i < sizeof(NonMainDbSupportedStatements) /
-		 sizeof(twoPcSupportedStatements[0]); i++)
+		 sizeof(NonMainDbSupportedStatements[0]); i++)
 	{
-		if (type == twoPcSupportedStatements[i].statementType)
+		if (type == NonMainDbSupportedStatements[i].statementType)
 		{
-			return twoPcSupportedStatements[i].markAsDistributed;
+			return NonMainDbSupportedStatements[i].explicitlyMarkAsDistributed;
 		}
 	}