When replanning a plpgsql "simple expression", check it's still simple.

The previous coding here assumed that we didn't need to recheck any
of the querytree tests made in exec_simple_check_plan().  I think
we supposed that those properties were fully determined by the
syntax of the source text and hence couldn't change.  That is true
for most of them, but at least hasTargetSRFs and hasAggs can change
by dint of forcibly dropping an originally-referenced function and
recreating it with new properties.  That leads to "unexpected plan
node type" or similar failures.

These tests are pretty cheap compared to the cost of replanning, so
rather than sweat over exactly which properties need to be rechecked,
let's just recheck them all.  Hence, factor out those tests into a new
function exec_is_simple_query(), and rearrange callers as needed.

A second problem in the same area was that if we failed during
replanning or during exec_save_simple_expr(), we'd potentially
leave behind now-dangling pointers to the old simple expression,
potentially resulting in crashes later.  To fix, clear those pointers
before replanning.

The v12 code looks quite different in this area but still has the
bug about needing to recheck query simplicity.  I chose to back-patch
all of the plpgsql_simple.sql test script, which formerly didn't exist
in this branch.

Per bug #18497 from Nikita Kalinin.  Back-patch to all supported
branches.

Discussion: https://postgr.es/m/18497-fe93b6da82ce31d4@postgresql.org

Author: tglsfdc

src/pl/plpgsql/src/expected/plpgsql_simple.out

@@ -89,3 +89,32 @@ select simplecaller();
             4
 (1 row)
 
+-- Check case where called function changes from non-SRF to SRF (bug #18497)
+create or replace function simplecaller() returns int language plpgsql
+as $$
+declare x int;
+begin
+  x := simplesql();
+  return x;
+end$$;
+select simplecaller();
+ simplecaller 
+--------------
+            4
+(1 row)
+
+drop function simplesql();
+create function simplesql() returns setof int language sql
+as $$select 22 + 22$$;
+select simplecaller();
+ simplecaller 
+--------------
+           44
+(1 row)
+
+select simplecaller();
+ simplecaller 
+--------------
+           44
+(1 row)
+

src/pl/plpgsql/src/pl_exec.c

@@ -339,6 +339,7 @@ static void exec_eval_cleanup(PLpgSQL_execstate *estate);
 static void exec_prepare_plan(PLpgSQL_execstate *estate,
 							  PLpgSQL_expr *expr, int cursorOptions);
 static void exec_simple_check_plan(PLpgSQL_execstate *estate, PLpgSQL_expr *expr);
+static bool exec_is_simple_query(PLpgSQL_expr *expr);
 static void exec_save_simple_expr(PLpgSQL_expr *expr, CachedPlan *cplan);
 static void exec_check_rw_parameter(PLpgSQL_expr *expr);
 static void exec_check_assignable(PLpgSQL_execstate *estate, int dno);
@@ -6094,12 +6095,18 @@ exec_eval_simple_expr(PLpgSQL_execstate *estate,
 		 * release it, so we don't leak plans intra-transaction.
 		 */
 		if (expr->expr_simple_plan_lxid == curlxid)
-		{
 			ReleaseCachedPlan(expr->expr_simple_plan,
 							  estate->simple_eval_resowner);
-			expr->expr_simple_plan = NULL;
-			expr->expr_simple_plan_lxid = InvalidLocalTransactionId;
-		}
+
+		/*
+		 * Reset to "not simple" to leave sane state (with no dangling
+		 * pointers) in case we fail while replanning.  expr_simple_plansource
+		 * can be left alone however, as that cannot move.
+		 */
+		expr->expr_simple_expr = NULL;
+		expr->expr_rw_param = NULL;
+		expr->expr_simple_plan = NULL;
+		expr->expr_simple_plan_lxid = InvalidLocalTransactionId;
 
 		/* Do the replanning work in the eval_mcontext */
 		oldcontext = MemoryContextSwitchTo(get_eval_mcontext(estate));
@@ -6115,11 +6122,15 @@ exec_eval_simple_expr(PLpgSQL_execstate *estate,
 		Assert(cplan != NULL);
 
 		/*
-		 * This test probably can't fail either, but if it does, cope by
-		 * declaring the plan to be non-simple.  On success, we'll acquire a
-		 * refcount on the new plan, stored in simple_eval_resowner.
+		 * Recheck exec_is_simple_query, which could now report false in
+		 * edge-case scenarios such as a non-SRF having been replaced with a
+		 * SRF.  Also recheck CachedPlanAllowsSimpleValidityCheck, just to be
+		 * sure.  If either test fails, cope by declaring the plan to be
+		 * non-simple.  On success, we'll acquire a refcount on the new plan,
+		 * stored in simple_eval_resowner.
 		 */
-		if (CachedPlanAllowsSimpleValidityCheck(expr->expr_simple_plansource,
+		if (exec_is_simple_query(expr) &&
+			CachedPlanAllowsSimpleValidityCheck(expr->expr_simple_plansource,
 												cplan,
 												estate->simple_eval_resowner))
 		{
@@ -6131,9 +6142,6 @@ exec_eval_simple_expr(PLpgSQL_execstate *estate,
 		{
 			/* Release SPI_plan_get_cached_plan's refcount */
 			ReleaseCachedPlan(cplan, CurrentResourceOwner);
-			/* Mark expression as non-simple, and fail */
-			expr->expr_simple_expr = NULL;
-			expr->expr_rw_param = NULL;
 			return false;
 		}
 
@@ -7974,7 +7982,6 @@ exec_simple_check_plan(PLpgSQL_execstate *estate, PLpgSQL_expr *expr)
 {
 	List	   *plansources;
 	CachedPlanSource *plansource;
-	Query	   *query;
 	CachedPlan *cplan;
 	MemoryContext oldcontext;
 
@@ -7990,31 +7997,88 @@ exec_simple_check_plan(PLpgSQL_execstate *estate, PLpgSQL_expr *expr)
 	 * called immediately after creating the CachedPlanSource, we need not
 	 * worry about the query being stale.
 	 */
+	if (!exec_is_simple_query(expr))
+		return;
+
+	/* exec_is_simple_query verified that there's just one CachedPlanSource */
+	plansources = SPI_plan_get_plan_sources(expr->plan);
+	plansource = (CachedPlanSource *) linitial(plansources);
 
 	/*
-	 * We can only test queries that resulted in exactly one CachedPlanSource
+	 * Get the generic plan for the query.  If replanning is needed, do that
+	 * work in the eval_mcontext.  (Note that replanning could throw an error,
+	 * in which case the expr is left marked "not simple", which is fine.)
+	 */
+	oldcontext = MemoryContextSwitchTo(get_eval_mcontext(estate));
+	cplan = SPI_plan_get_cached_plan(expr->plan);
+	MemoryContextSwitchTo(oldcontext);
+
+	/* Can't fail, because we checked for a single CachedPlanSource above */
+	Assert(cplan != NULL);
+
+	/*
+	 * Verify that plancache.c thinks the plan is simple enough to use
+	 * CachedPlanIsSimplyValid.  Given the restrictions above, it's unlikely
+	 * that this could fail, but if it does, just treat plan as not simple. On
+	 * success, save a refcount on the plan in the simple-expression resowner.
+	 */
+	if (CachedPlanAllowsSimpleValidityCheck(plansource, cplan,
+											estate->simple_eval_resowner))
+	{
+		/* Remember that we have the refcount */
+		expr->expr_simple_plansource = plansource;
+		expr->expr_simple_plan = cplan;
+		expr->expr_simple_plan_lxid = MyProc->vxid.lxid;
+
+		/* Share the remaining work with the replan code path */
+		exec_save_simple_expr(expr, cplan);
+	}
+
+	/*
+	 * Release the plan refcount obtained by SPI_plan_get_cached_plan.  (This
+	 * refcount is held by the wrong resowner, so we can't just repurpose it.)
+	 */
+	ReleaseCachedPlan(cplan, CurrentResourceOwner);
+}
+
+/*
+ * exec_is_simple_query - precheck a query tree to see if it might be simple
+ *
+ * Check the analyzed-and-rewritten form of a query to see if we will be
+ * able to treat it as a simple expression.  It is caller's responsibility
+ * that the CachedPlanSource be up-to-date.
+ */
+static bool
+exec_is_simple_query(PLpgSQL_expr *expr)
+{
+	List	   *plansources;
+	CachedPlanSource *plansource;
+	Query	   *query;
+
+	/*
+	 * We can only test queries that resulted in exactly one CachedPlanSource.
 	 */
 	plansources = SPI_plan_get_plan_sources(expr->plan);
 	if (list_length(plansources) != 1)
-		return;
+		return false;
 	plansource = (CachedPlanSource *) linitial(plansources);
 
 	/*
 	 * 1. There must be one single querytree.
 	 */
 	if (list_length(plansource->query_list) != 1)
-		return;
+		return false;
 	query = (Query *) linitial(plansource->query_list);
 
 	/*
-	 * 2. It must be a plain SELECT query without any input tables
+	 * 2. It must be a plain SELECT query without any input tables.
 	 */
 	if (!IsA(query, Query))
-		return;
+		return false;
 	if (query->commandType != CMD_SELECT)
-		return;
+		return false;
 	if (query->rtable != NIL)
-		return;
+		return false;
 
 	/*
 	 * 3. Can't have any subplans, aggregates, qual clauses either.  (These
@@ -8038,51 +8102,18 @@ exec_simple_check_plan(PLpgSQL_execstate *estate, PLpgSQL_expr *expr)
 		query->limitOffset ||
 		query->limitCount ||
 		query->setOperations)
-		return;
+		return false;
 
 	/*
-	 * 4. The query must have a single attribute as result
+	 * 4. The query must have a single attribute as result.
 	 */
 	if (list_length(query->targetList) != 1)
-		return;
+		return false;
 
 	/*
 	 * OK, we can treat it as a simple plan.
-	 *
-	 * Get the generic plan for the query.  If replanning is needed, do that
-	 * work in the eval_mcontext.  (Note that replanning could throw an error,
-	 * in which case the expr is left marked "not simple", which is fine.)
-	 */
-	oldcontext = MemoryContextSwitchTo(get_eval_mcontext(estate));
-	cplan = SPI_plan_get_cached_plan(expr->plan);
-	MemoryContextSwitchTo(oldcontext);
-
-	/* Can't fail, because we checked for a single CachedPlanSource above */
-	Assert(cplan != NULL);
-
-	/*
-	 * Verify that plancache.c thinks the plan is simple enough to use
-	 * CachedPlanIsSimplyValid.  Given the restrictions above, it's unlikely
-	 * that this could fail, but if it does, just treat plan as not simple. On
-	 * success, save a refcount on the plan in the simple-expression resowner.
-	 */
-	if (CachedPlanAllowsSimpleValidityCheck(plansource, cplan,
-											estate->simple_eval_resowner))
-	{
-		/* Remember that we have the refcount */
-		expr->expr_simple_plansource = plansource;
-		expr->expr_simple_plan = cplan;
-		expr->expr_simple_plan_lxid = MyProc->vxid.lxid;
-
-		/* Share the remaining work with the replan code path */
-		exec_save_simple_expr(expr, cplan);
-	}
-
-	/*
-	 * Release the plan refcount obtained by SPI_plan_get_cached_plan.  (This
-	 * refcount is held by the wrong resowner, so we can't just repurpose it.)
 	 */
-	ReleaseCachedPlan(cplan, CurrentResourceOwner);
+	return true;
 }
 
 /*

src/pl/plpgsql/src/sql/plpgsql_simple.sql

@@ -80,3 +80,25 @@ create or replace function simplesql() returns int language sql
 as $$select 2 + 2$$;
 
 select simplecaller();
+
+
+-- Check case where called function changes from non-SRF to SRF (bug #18497)
+
+create or replace function simplecaller() returns int language plpgsql
+as $$
+declare x int;
+begin
+  x := simplesql();
+  return x;
+end$$;
+
+select simplecaller();
+
+drop function simplesql();
+
+create function simplesql() returns setof int language sql
+as $$select 22 + 22$$;
+
+select simplecaller();
+
+select simplecaller();