suggestions

src/backend/distributed/commands/utility_hook.c

@@ -107,8 +107,14 @@ typedef struct NonMainDbDistributedStatementInfo
 {
 	int statementType;
 	bool explicitlyMarkAsDistributed;
-	ObjectType *supportedObjectTypes;
-	int supportedObjectTypesSize;
+
+	/*
+	 * checkSupportedObjectTypes is a callback function that checks whether
+	 * type of the object referred to by given statement is supported.
+	 *
+	 * Can be NULL if not applicable for the statement type.
+	 */
+	bool (*checkSupportedObjectTypes)(Node *node);
 } NonMainDbDistributedStatementInfo;
 
 /*
@@ -122,18 +128,6 @@ typedef struct MarkObjectDistributedParams
 	uint16 catalogRelId;
 } MarkObjectDistributedParams;
 
-/*
- * NonMainDbSupportedStatements is an array of statements that are supported
- * from non-main databases.
- */
-ObjectType supportedObjectTypesForGrantStmt[] = { OBJECT_DATABASE };
-static const NonMainDbDistributedStatementInfo NonMainDbSupportedStatements[] = {
-	{ T_GrantRoleStmt, false, NULL, 0 },
-	{ T_CreateRoleStmt, true, NULL, 0 },
-	{ T_GrantStmt, false, supportedObjectTypesForGrantStmt,
-	  sizeof(supportedObjectTypesForGrantStmt) / sizeof(ObjectType) }
-};
-
 
 bool EnableDDLPropagation = true; /* ddl propagation is enabled */
 int CreateObjectPropagationMode = CREATE_OBJECT_PROPAGATION_IMMEDIATE;
@@ -162,6 +156,12 @@ static void PostStandardProcessUtility(Node *parsetree);
 static void DecrementUtilityHookCountersIfNecessary(Node *parsetree);
 static bool IsDropSchemaOrDB(Node *parsetree);
 static bool ShouldCheckUndistributeCitusLocalTables(void);
+
+
+/*
+ * Functions to support commands used to manage node-wide objects from non-main
+ * databases.
+ */
 static void RunPreprocessMainDBCommand(Node *parsetree);
 static void RunPostprocessMainDBCommand(Node *parsetree);
 static bool IsStatementSupportedFromNonMainDb(Node *parsetree);
@@ -169,6 +169,25 @@ static bool StatementRequiresMarkDistributedFromNonMainDb(Node *parsetree);
 static void MarkObjectDistributedFromNonMainDb(Node *parsetree);
 static MarkObjectDistributedParams GetMarkObjectDistributedParams(Node *parsetree);
 
+/*
+ * checkSupportedObjectTypes callbacks for
+ * NonMainDbDistributedStatementInfo objects.
+ */
+static bool NonMainDbCheckSupportedObjectTypeForGrant(Node *node);
+
+
+/*
+ * NonMainDbSupportedStatements is an array of statements that are supported
+ * from non-main databases.
+ */
+ObjectType supportedObjectTypesForGrantStmt[] = { OBJECT_DATABASE };
+static const NonMainDbDistributedStatementInfo NonMainDbSupportedStatements[] = {
+	{ T_GrantRoleStmt, false, NULL },
+	{ T_CreateRoleStmt, true, NULL },
+	{ T_GrantStmt, false, NonMainDbCheckSupportedObjectTypeForGrant}
+};
+
+
 /*
  * ProcessUtilityParseTree is a convenience method to create a PlannedStmt out of
  * pieces of a utility statement before invoking ProcessUtility.
@@ -1701,31 +1720,11 @@ IsStatementSupportedFromNonMainDb(Node *parsetree)
 	for (int i = 0; i < sizeof(NonMainDbSupportedStatements) /
 		 sizeof(NonMainDbSupportedStatements[0]); i++)
 	{
-		if (type == NonMainDbSupportedStatements[i].statementType)
+		if (type == NonMainDbSupportedStatements[i].statementType &&
+			(!NonMainDbSupportedStatements[i].checkSupportedObjectTypes ||
+			 NonMainDbSupportedStatements[i].checkSupportedObjectTypes(parsetree)))
 		{
-			if (NonMainDbSupportedStatements[i].supportedObjectTypes == NULL)
-			{
-				return true;
-			}
-			else
-			{
-				if (type == T_GrantStmt)
-				{
-					GrantStmt *stmt = castNode(GrantStmt, parsetree);
-
-					/* check if stmt->objtype is in supportedObjectTypes */
-					for (int j = 0; j <
-						 NonMainDbSupportedStatements[i].supportedObjectTypesSize; j++)
-					{
-						if (stmt->objtype ==
-							NonMainDbSupportedStatements[i].supportedObjectTypes[j])
-						{
-							return true;
-						}
-					}
-					return false;
-				}
-			}
+			return true;
 		}
 	}
 
@@ -1798,3 +1797,15 @@ GetMarkObjectDistributedParams(Node *parsetree)
 
 	elog(ERROR, "unsupported statement type");
 }
+
+
+/*
+ * NonMainDbCheckSupportedObjectTypeForGrant implements checkSupportedObjectTypes
+ * callback for GrantStmt.
+ */
+static bool
+NonMainDbCheckSupportedObjectTypeForGrant(Node *node)
+{
+	GrantStmt *stmt = castNode(GrantStmt, node);
+	return stmt->objtype == OBJECT_DATABASE;
+}