Merge pull request #3889 from citusdata/fix/stage_generates_utility_commands

Execute shard creation as utility tasks

src/backend/distributed/executor/adaptive_executor.c

@@ -855,6 +855,29 @@ ExecuteUtilityTaskList(List *utilityTaskList, bool localExecutionSupported)
 }
 
 
+/*
+ * ExecuteUtilityTaskListExtended is a wrapper around executing task
+ * list for utility commands.
+ */
+uint64
+ExecuteUtilityTaskListExtended(List *utilityTaskList, int poolSize,
+							   bool localExecutionSupported)
+{
+	RowModifyLevel modLevel = ROW_MODIFY_NONE;
+	ExecutionParams *executionParams = CreateBasicExecutionParams(
+		modLevel, utilityTaskList, poolSize, localExecutionSupported
+		);
+
+	bool excludeFromXact = false;
+	executionParams->xactProperties =
+		DecideTransactionPropertiesForTaskList(modLevel, utilityTaskList,
+											   excludeFromXact);
+	executionParams->isUtilityCommand = true;
+
+	return ExecuteTaskListExtended(executionParams);
+}
+
+
 /*
  * ExecuteTaskListOutsideTransaction is a proxy to ExecuteTaskListExtended
  * with defaults for some of the arguments.

src/backend/distributed/executor/local_executor.c

@@ -366,8 +366,11 @@ ExtractParametersForLocalExecution(ParamListInfo paramListInfo, Oid **parameterT
 static void
 LocallyExecuteUtilityTask(const char *localTaskQueryCommand)
 {
-	RawStmt *localTaskRawStmt = (RawStmt *) ParseTreeRawStmt(localTaskQueryCommand);
+	List *parseTreeList = pg_parse_query(localTaskQueryCommand);
+	RawStmt *localTaskRawStmt = NULL;
 
+	foreach_ptr(localTaskRawStmt, parseTreeList)
+	{
 		Node *localTaskRawParseTree = localTaskRawStmt->stmt;
 
 		/*
@@ -396,6 +399,7 @@ LocallyExecuteUtilityTask(const char *localTaskQueryCommand)
 			CitusProcessUtility(localTaskRawParseTree, localTaskQueryCommand,
 								PROCESS_UTILITY_TOPLEVEL, NULL, None_Receiver, NULL);
 		}
+	}
 }
 
 

src/backend/distributed/master/master_stage_protocol.c

@@ -550,7 +550,7 @@ CreateShardsOnWorkers(Oid distributedRelationId, List *shardPlacements,
 		poolSize = MaxAdaptiveExecutorPoolSize;
 	}
 	bool localExecutionSupported = true;
-	ExecuteTaskList(ROW_MODIFY_NONE, taskList, poolSize, localExecutionSupported);
+	ExecuteUtilityTaskListExtended(taskList, poolSize, localExecutionSupported);
 }
 
 

src/include/distributed/adaptive_executor.h

@@ -13,6 +13,8 @@ extern int ExecutorSlowStartInterval;
 extern uint64 ExecuteTaskList(RowModifyLevel modLevel, List *taskList,
 							  int targetPoolSize, bool localExecutionSupported);
 extern uint64 ExecuteUtilityTaskList(List *utilityTaskList, bool localExecutionSupported);
+extern uint64 ExecuteUtilityTaskListExtended(List *utilityTaskList, int poolSize,
+											 bool localExecutionSupported);
 extern uint64 ExecuteTaskListOutsideTransaction(RowModifyLevel modLevel, List *taskList,
 												int targetPoolSize, List *jobIdList);