Handle multi-command queries correctly inside SQL as well as plpgsql functions.

Postgres-XL sends down utility statements to the remote nodes as plain query
strings. When there are multiple commands in a query string, separated by ';',
we were incorrectly sending down the entire query string again and again while
handling each command. This can lead to unpleasant as well as incorrect
behaviour. This was earlier handled for execution via psql, but this patch
fixes it for SPI and other places such as extension creation and SQL function
handling.

2 files changed, 41 insertions, 23 deletions

diff --git a/src/backend/executor/functions.c b/src/backend/executor/functions.c
index 4e482f397f..06b4a57656 100644
--- a/src/backend/executor/functions.c
+++ b/src/backend/executor/functions.c
@@ -72,6 +72,7 @@ typedef struct execution_state
 	bool		lazyEval;		/* true if should fetch one row at a time */
 	Node	   *stmt;			/* PlannedStmt or utility statement */
 	QueryDesc  *qd;				/* null unless status == RUN */
+	char		*src;			/* source query resulting in this state */
 } execution_state;
 
 
@@ -156,6 +157,7 @@ static Node *sql_fn_make_param(SQLFunctionParseInfoPtr pinfo,
 static Node *sql_fn_resolve_param_name(SQLFunctionParseInfoPtr pinfo,
 						  const char *paramname, int location);
 static List *init_execution_state(List *queryTree_list,
+					 List *querySource_list,
 					 SQLFunctionCachePtr fcache,
 					 bool lazyEvalOK);
 static void init_sql_fcache(FmgrInfo *finfo, Oid collation, bool lazyEvalOK);
@@ -474,19 +476,21 @@ sql_fn_resolve_param_name(SQLFunctionParseInfoPtr pinfo,
  */
 static List *
 init_execution_state(List *queryTree_list,
+					 List *querySource_list,
 					 SQLFunctionCachePtr fcache,
 					 bool lazyEvalOK)
 {
 	List	   *eslist = NIL;
 	execution_state *lasttages = NULL;
-	ListCell   *lc1;
+	ListCell   *lc1, *lc3;
 
-	foreach(lc1, queryTree_list)
+	forboth(lc1, queryTree_list, lc3, querySource_list)
 	{
 		List	   *qtlist = (List *) lfirst(lc1);
+		char	   *querysource = (char *) lfirst(lc3);
 		execution_state *firstes = NULL;
 		execution_state *preves = NULL;
-		ListCell   *lc2;
+		ListCell   *lc2, *lc4;
 
 		foreach(lc2, qtlist)
 		{
@@ -553,6 +557,7 @@ init_execution_state(List *queryTree_list,
 			newes->lazyEval = false;	/* might change below */
 			newes->stmt = stmt;
 			newes->qd = NULL;
+			newes->src = pstrdup(querysource);
 
 			if (queryTree->canSetTag)
 				lasttages = newes;
@@ -610,9 +615,10 @@ init_sql_fcache(FmgrInfo *finfo, Oid collation, bool lazyEvalOK)
 	Form_pg_proc procedureStruct;
 	SQLFunctionCachePtr fcache;
 	List	   *raw_parsetree_list;
+	List	   *querysource_list;
 	List	   *queryTree_list;
 	List	   *flat_query_list;
-	ListCell   *lc;
+	ListCell   *lc, *lc2;
 	Datum		tmp;
 	bool		isNull;
 
@@ -713,17 +719,18 @@ init_sql_fcache(FmgrInfo *finfo, Oid collation, bool lazyEvalOK)
 	 * but we'll not worry about it until the module is rewritten to use
 	 * plancache.c.
 	 */
-	raw_parsetree_list = pg_parse_query(fcache->src);
+	raw_parsetree_list = pg_parse_query_get_source(fcache->src, &querysource_list);
 
 	queryTree_list = NIL;
 	flat_query_list = NIL;
-	foreach(lc, raw_parsetree_list)
+	forboth(lc, raw_parsetree_list, lc2, querysource_list)
 	{
 		Node	   *parsetree = (Node *) lfirst(lc);
+		char	   *querysource = (char *) lfirst(lc2);
 		List	   *queryTree_sublist;
 
 		queryTree_sublist = pg_analyze_and_rewrite_params(parsetree,
-														  fcache->src,
+														  querysource,
 									   (ParserSetupHook) sql_fn_parser_setup,
 														  fcache->pinfo);
 		queryTree_list = lappend(queryTree_list, queryTree_sublist);
@@ -774,6 +781,7 @@ init_sql_fcache(FmgrInfo *finfo, Oid collation, bool lazyEvalOK)
 
 	/* Finally, plan the queries */
 	fcache->func_state = init_execution_state(queryTree_list,
+											  querysource_list,
 											  fcache,
 											  lazyEvalOK);
 
@@ -818,14 +826,14 @@ postquel_start(execution_state *es, SQLFunctionCachePtr fcache)
 
 	if (IsA(es->stmt, PlannedStmt))
 		es->qd = CreateQueryDesc((PlannedStmt *) es->stmt,
-								 fcache->src,
+								 es->src,
 								 GetActiveSnapshot(),
 								 InvalidSnapshot,
 								 dest,
 								 fcache->paramLI, 0);
 	else
 		es->qd = CreateUtilityQueryDesc(es->stmt,
-										fcache->src,
+										es->src,
 										GetActiveSnapshot(),
 										dest,
 										fcache->paramLI);
@@ -865,7 +873,7 @@ postquel_getnext(execution_state *es, SQLFunctionCachePtr fcache)
 		ProcessUtility((es->qd->plannedstmt ?
 						(Node *) es->qd->plannedstmt :
 						es->qd->utilitystmt),
-					   fcache->src,
+					   es->src,
 					   PROCESS_UTILITY_QUERY,
 					   es->qd->params,
 					   es->qd->dest,
diff --git a/src/backend/executor/spi.c b/src/backend/executor/spi.c
index d971d936c6..0a3c65e6f7 100644
--- a/src/backend/executor/spi.c
+++ b/src/backend/executor/spi.c
@@ -52,6 +52,7 @@ static Portal SPI_cursor_open_internal(const char *name, SPIPlanPtr plan,
 
 #ifdef PGXC
 static void _SPI_pgxc_prepare_plan(const char *src, List *src_parsetree,
+				  List *query_source,
                   SPIPlanPtr plan);
 #endif
 static void _SPI_prepare_plan(const char *src, SPIPlanPtr plan);
@@ -405,7 +406,8 @@ SPI_execute_direct(const char *remote_sql, char *nodename)
 	plan.cursor_options = 0;
 
 	/* Now pass the ExecDirectStmt parsetree node */
-	_SPI_pgxc_prepare_plan(execdirect.data, list_make1(stmt), &plan);
+	_SPI_pgxc_prepare_plan(execdirect.data, list_make1(stmt),
+			list_make1(execdirect.data), &plan);
 
 	res = _SPI_execute_plan(&plan, NULL,
 							InvalidSnapshot, InvalidSnapshot, false, true, 0);
@@ -1888,7 +1890,7 @@ static void
 _SPI_prepare_plan(const char *src, SPIPlanPtr plan)
 {
 #ifdef PGXC
-	_SPI_pgxc_prepare_plan(src, NULL, plan);
+	_SPI_pgxc_prepare_plan(src, NULL, NULL, plan);
 }
 
 /*
@@ -1898,12 +1900,14 @@ _SPI_prepare_plan(const char *src, SPIPlanPtr plan)
  * transparent to the user.
  */
 static void
-_SPI_pgxc_prepare_plan(const char *src, List *src_parsetree, SPIPlanPtr plan)
+_SPI_pgxc_prepare_plan(const char *src, List *src_parsetree,
+		List *query_source, SPIPlanPtr plan)
 {
 #endif
 	List	   *raw_parsetree_list;
+	List	   *querysource_list;
 	List	   *plancache_list;
-	ListCell   *list_item;
+	ListCell   *list_item, *list_item2;
 	ErrorContextCallback spierrcontext;
 
 	/*
@@ -1920,19 +1924,23 @@ _SPI_pgxc_prepare_plan(const char *src, List *src_parsetree, SPIPlanPtr plan)
 #ifdef PGXC
 	/* Parse it only if there isn't an already parsed tree passed */
 	if (src_parsetree)
+	{
 		raw_parsetree_list = src_parsetree;
+		querysource_list = query_source;
+	}
 	else
 #endif
-		raw_parsetree_list = pg_parse_query(src);
+		raw_parsetree_list = pg_parse_query_get_source(src, &querysource_list);
 	/*
 	 * Do parse analysis and rule rewrite for each raw parsetree, storing the
 	 * results into unsaved plancache entries.
 	 */
 	plancache_list = NIL;
 
-	foreach(list_item, raw_parsetree_list)
+	forboth(list_item, raw_parsetree_list, list_item2, querysource_list)
 	{
 		Node	   *parsetree = (Node *) lfirst(list_item);
+		char	   *querysource = (char *) lfirst (list_item2);
 		List	   *stmt_list;
 		CachedPlanSource *plansource;
 
@@ -1941,7 +1949,7 @@ _SPI_pgxc_prepare_plan(const char *src, List *src_parsetree, SPIPlanPtr plan)
 		 * needs to see the unmodified raw parse tree.
 		 */
 		plansource = CreateCachedPlan(parsetree,
-									  src,
+									  querysource,
 #ifdef PGXC
 									  NULL,
 #endif
@@ -1955,14 +1963,14 @@ _SPI_pgxc_prepare_plan(const char *src, List *src_parsetree, SPIPlanPtr plan)
 		{
 			Assert(plan->nargs == 0);
 			stmt_list = pg_analyze_and_rewrite_params(parsetree,
-													  src,
+													  querysource,
 													  plan->parserSetup,
 													  plan->parserSetupArg);
 		}
 		else
 		{
 			stmt_list = pg_analyze_and_rewrite(parsetree,
-											   src,
+											   querysource,
 											   plan->argtypes,
 											   plan->nargs);
 		}
@@ -2013,8 +2021,9 @@ static void
 _SPI_prepare_oneshot_plan(const char *src, SPIPlanPtr plan)
 {
 	List	   *raw_parsetree_list;
+	List	   *querysource_list;
 	List	   *plancache_list;
-	ListCell   *list_item;
+	ListCell   *list_item, *list_item2;
 	ErrorContextCallback spierrcontext;
 
 	/*
@@ -2028,20 +2037,21 @@ _SPI_prepare_oneshot_plan(const char *src, SPIPlanPtr plan)
 	/*
 	 * Parse the request string into a list of raw parse trees.
 	 */
-	raw_parsetree_list = pg_parse_query(src);
+	raw_parsetree_list = pg_parse_query_get_source(src, &querysource_list);
 
 	/*
 	 * Construct plancache entries, but don't do parse analysis yet.
 	 */
 	plancache_list = NIL;
 
-	foreach(list_item, raw_parsetree_list)
+	forboth(list_item, raw_parsetree_list, list_item2, querysource_list)
 	{
 		Node	   *parsetree = (Node *) lfirst(list_item);
+		char	   *querysource = (char *) lfirst (list_item2);
 		CachedPlanSource *plansource;
 
 		plansource = CreateOneShotCachedPlan(parsetree,
-											 src,
+											 querysource,
 											 CreateCommandTag(parsetree));
 
 		plancache_list = lappend(plancache_list, plansource);