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.

5 files changed, 146 insertions, 38 deletions

diff --git a/src/backend/commands/extension.c b/src/backend/commands/extension.c
index 680c22dee2..be8641b573 100644
--- a/src/backend/commands/extension.c
+++ b/src/backend/commands/extension.c
@@ -683,13 +683,14 @@ static void
 execute_sql_string(const char *sql, const char *filename)
 {
 	List	   *raw_parsetree_list;
+	List	   *querysource_list;
 	DestReceiver *dest;
-	ListCell   *lc1;
+	ListCell   *lc1, *lc3;
 
 	/*
 	 * Parse the SQL string into a list of raw parse trees.
 	 */
-	raw_parsetree_list = pg_parse_query(sql);
+	raw_parsetree_list = pg_parse_query_get_source(sql, &querysource_list);
 
 	/* All output from SELECTs goes to the bit bucket */
 	dest = CreateDestReceiver(DestNone);
@@ -699,14 +700,15 @@ execute_sql_string(const char *sql, const char *filename)
 	 * parsetree.  We must fully execute each query before beginning parse
 	 * analysis on the next one, since there may be interdependencies.
 	 */
-	foreach(lc1, raw_parsetree_list)
+	forboth(lc1, raw_parsetree_list, lc3, querysource_list)
 	{
 		Node	   *parsetree = (Node *) lfirst(lc1);
+		char	   *querysource = (char *) lfirst(lc3);
 		List	   *stmt_list;
 		ListCell   *lc2;
 
 		stmt_list = pg_analyze_and_rewrite(parsetree,
-										   sql,
+										   querysource,
 										   NULL,
 										   0);
 		stmt_list = pg_plan_queries(stmt_list, CURSOR_OPT_PARALLEL_OK, NULL);
@@ -730,7 +732,7 @@ execute_sql_string(const char *sql, const char *filename)
 				QueryDesc  *qdesc;
 
 				qdesc = CreateQueryDesc((PlannedStmt *) stmt,
-										sql,
+										querysource,
 										GetActiveSnapshot(), NULL,
 										dest, NULL, 0);
 
@@ -744,7 +746,7 @@ execute_sql_string(const char *sql, const char *filename)
 			else
 			{
 				ProcessUtility(stmt,
-							   sql,
+							   querysource,
 							   PROCESS_UTILITY_QUERY,
 							   NULL,
 							   dest,
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);
diff --git a/src/test/regress/expected/xc_misc.out b/src/test/regress/expected/xc_misc.out
index 75d207cccb..b378aa495e 100644
--- a/src/test/regress/expected/xc_misc.out
+++ b/src/test/regress/expected/xc_misc.out
@@ -51,14 +51,79 @@ drop table t1_misc;
 create table my_tab1 (a int);
 insert into my_tab1 values(1);
 create function f1 () returns setof my_tab1 as $$ create table my_tab2 (a int); select * from my_tab1; $$ language sql;
+select f1();
+ f1  
+-----
+ (1)
+(1 row)
+
+drop function f1();
+-- fail since my_tab4 does not exist
+create function f1 () returns setof my_tab2 as $$ create table my_tab3 (a int); select * from my_tab4; $$ language sql;
+ERROR:  relation "my_tab4" does not exist
+LINE 1: ...as $$ create table my_tab3 (a int); select * from my_tab4; $...
+                                                             ^
 SET check_function_bodies = false;
-create function f1 () returns setof my_tab1 as $$ create table my_tab2 (a int); select * from my_tab1; $$ language sql;
-ERROR:  function "f1" already exists with same argument types
+-- should be created since check_function_bodies is false
+create function f1 () returns setof my_tab2 as $$ create table my_tab3 (a int); select * from my_tab4; $$ language sql;
+-- execution would fail though
 select f1();
-ERROR:  Unexpected response from Datanode
-CONTEXT:  SQL function "f1" statement 1
-SET check_function_bodies = true;
+ERROR:  relation "my_tab4" does not exist
+LINE 1:  create table my_tab3 (a int); select * from my_tab4; 
+                               ^
+QUERY:   create table my_tab3 (a int); select * from my_tab4; 
+CONTEXT:  SQL function "f1" during startup
 drop function f1();
+SET check_function_bodies = true;
+-- check handling of multi-command statements in plpgsql block
+do
+$$
+declare
+begin
+	execute 'create table my_tab6(a int); create table my_tab7(a int)';
+end
+$$ language plpgsql;
+-- check handling of multi-command statements in sql function
+create or replace function f2() returns void as $$ create table my_tab8(a int); create table my_tab9(a int); $$ language sql;
+select f2();
+ f2 
+----
+ 
+(1 row)
+
+\d+ my_tab6
+                       Table "public.my_tab6"
+ Column |  Type   | Modifiers | Storage | Stats target | Description 
+--------+---------+-----------+---------+--------------+-------------
+ a      | integer |           | plain   |              | 
+Distribute By: HASH(a)
+Location Nodes: ALL DATANODES
+
+\d+ my_tab7
+                       Table "public.my_tab7"
+ Column |  Type   | Modifiers | Storage | Stats target | Description 
+--------+---------+-----------+---------+--------------+-------------
+ a      | integer |           | plain   |              | 
+Distribute By: HASH(a)
+Location Nodes: ALL DATANODES
+
+\d+ my_tab8
+                       Table "public.my_tab8"
+ Column |  Type   | Modifiers | Storage | Stats target | Description 
+--------+---------+-----------+---------+--------------+-------------
+ a      | integer |           | plain   |              | 
+Distribute By: HASH(a)
+Location Nodes: ALL DATANODES
+
+\d+ my_tab9
+                       Table "public.my_tab9"
+ Column |  Type   | Modifiers | Storage | Stats target | Description 
+--------+---------+-----------+---------+--------------+-------------
+ a      | integer |           | plain   |              | 
+Distribute By: HASH(a)
+Location Nodes: ALL DATANODES
+
+drop table my_tab6, my_tab7, my_tab8, my_tab9;
 -- Test pl-pgsql functions containing utility statements
 CREATE OR REPLACE FUNCTION test_fun_2() RETURNS SETOF my_tab1 AS '
 DECLARE
diff --git a/src/test/regress/sql/xc_misc.sql b/src/test/regress/sql/xc_misc.sql
index 30db55a1a3..5b5057a440 100644
--- a/src/test/regress/sql/xc_misc.sql
+++ b/src/test/regress/sql/xc_misc.sql
@@ -44,16 +44,39 @@ create table my_tab1 (a int);
 insert into my_tab1 values(1);
 
 create function f1 () returns setof my_tab1 as $$ create table my_tab2 (a int); select * from my_tab1; $$ language sql;
+select f1();
+drop function f1();
 
+-- fail since my_tab4 does not exist
+create function f1 () returns setof my_tab2 as $$ create table my_tab3 (a int); select * from my_tab4; $$ language sql;
 SET check_function_bodies = false;
-
-create function f1 () returns setof my_tab1 as $$ create table my_tab2 (a int); select * from my_tab1; $$ language sql;
-
+-- should be created since check_function_bodies is false
+create function f1 () returns setof my_tab2 as $$ create table my_tab3 (a int); select * from my_tab4; $$ language sql;
+-- execution would fail though
 select f1();
+drop function f1();
 
 SET check_function_bodies = true;
 
-drop function f1();
+-- check handling of multi-command statements in plpgsql block
+do
+$$
+declare
+begin
+	execute 'create table my_tab6(a int); create table my_tab7(a int)';
+end
+$$ language plpgsql;
+
+-- check handling of multi-command statements in sql function
+create or replace function f2() returns void as $$ create table my_tab8(a int); create table my_tab9(a int); $$ language sql;
+select f2();
+
+\d+ my_tab6
+\d+ my_tab7
+\d+ my_tab8
+\d+ my_tab9
+
+drop table my_tab6, my_tab7, my_tab8, my_tab9;
 
 -- Test pl-pgsql functions containing utility statements