From f0b52fe1844bf1be0a3fcc718069c7c1a41393e1 Mon Sep 17 00:00:00 2001
From: Tom Lane <tgl@sss.pgh.pa.us>
Date: Sat, 29 Mar 2025 16:56:57 -0400
Subject: [PATCH v10 4/6] Restructure check_sql_fn_retval().

To support using the plan cache for SQL functions, we'll need to be
able to redo the work of check_sql_fn_retval() on just one query's
list-of-rewritten-queries at a time, since the plan cache will treat
each query independently.  This would be simple enough, except for a
bizarre historical behavior: the existing code will take the last
canSetTag query in the function as determining the result, even if
it came from not-the-last original query.  (The case is only possible
when the last original query(s) are deleted by a DO INSTEAD NOTHING
rule.)  This behavior is undocumented except in source code comments,
and it seems hard to believe that anyone's relying on it.  It would
be a mess to support with the plan cache, because a change in the
rules applicable to some table could change which CachedPlanSource
is supposed to produce the function result, even if the function
itself has not changed.  Let's just get rid of that silliness and
insist that the last source query in the function is the one that
must produce the result.

Having mandated that, we can refactor check_sql_fn_retval() into
an outer and an inner function where the inner one considers only
a single list-of-rewritten-queries; the inner one will be usable
in a post-rewrite callback hook as contemplated by the previous
commit.

Likewise refactor check_sql_fn_statements() so that we have a
version that can be applied to just one list of Queries.
(As things stand, it's not really necessary to recheck that
during a replan, but maybe future changes in the rule system
would create cases where it matters.)

Also remove check_sql_fn_retval()'s targetlist output argument,
putting the equivalent functionality into a separate function.
This is needed because the plan cache would be in the way of
passing that data directly.  No outside caller needed that
anyway.

Author: Tom Lane <tgl@sss.pgh.pa.us>
Discussion: https://postgr.es/m/8216639.NyiUUSuA9g@aivenlaptop
---
 src/backend/catalog/pg_proc.c        |   2 +-
 src/backend/executor/functions.c     | 176 ++++++++++++++++++---------
 src/backend/optimizer/util/clauses.c |   4 +-
 src/include/executor/functions.h     |   3 +-
 4 files changed, 121 insertions(+), 64 deletions(-)

diff --git a/src/backend/catalog/pg_proc.c b/src/backend/catalog/pg_proc.c
index fe0490259e9..880b597fb3a 100644
--- a/src/backend/catalog/pg_proc.c
+++ b/src/backend/catalog/pg_proc.c
@@ -960,7 +960,7 @@ fmgr_sql_validator(PG_FUNCTION_ARGS)
 			(void) check_sql_fn_retval(querytree_list,
 									   rettype, rettupdesc,
 									   proc->prokind,
-									   false, NULL);
+									   false);
 		}
 
 		error_context_stack = sqlerrcontext.previous;
diff --git a/src/backend/executor/functions.c b/src/backend/executor/functions.c
index 6aa8e9c4d8a..5b06df84335 100644
--- a/src/backend/executor/functions.c
+++ b/src/backend/executor/functions.c
@@ -153,11 +153,16 @@ static Datum postquel_get_single_result(TupleTableSlot *slot,
 										MemoryContext resultcontext);
 static void sql_exec_error_callback(void *arg);
 static void ShutdownSQLFunction(Datum arg);
+static void check_sql_fn_statement(List *queryTreeList);
+static bool check_sql_stmt_retval(List *queryTreeList,
+								  Oid rettype, TupleDesc rettupdesc,
+								  char prokind, bool insertDroppedCols);
 static bool coerce_fn_result_column(TargetEntry *src_tle,
 									Oid res_type, int32 res_typmod,
 									bool tlist_is_modifiable,
 									List **upper_tlist,
 									bool *upper_tlist_nontrivial);
+static List *get_sql_fn_result_tlist(List *queryTreeList);
 static void sqlfunction_startup(DestReceiver *self, int operation, TupleDesc typeinfo);
 static bool sqlfunction_receive(TupleTableSlot *slot, DestReceiver *self);
 static void sqlfunction_shutdown(DestReceiver *self);
@@ -592,7 +597,6 @@ init_sql_fcache(FunctionCallInfo fcinfo, Oid collation, bool lazyEvalOK)
 	Form_pg_proc procedureStruct;
 	SQLFunctionCachePtr fcache;
 	List	   *queryTree_list;
-	List	   *resulttlist;
 	ListCell   *lc;
 	Datum		tmp;
 	bool		isNull;
@@ -748,8 +752,7 @@ init_sql_fcache(FunctionCallInfo fcinfo, Oid collation, bool lazyEvalOK)
 											   rettype,
 											   rettupdesc,
 											   procedureStruct->prokind,
-											   false,
-											   &resulttlist);
+											   false);
 
 	/*
 	 * Construct a JunkFilter we can use to coerce the returned rowtype to the
@@ -762,6 +765,14 @@ init_sql_fcache(FunctionCallInfo fcinfo, Oid collation, bool lazyEvalOK)
 	{
 		TupleTableSlot *slot = MakeSingleTupleTableSlot(NULL,
 														&TTSOpsMinimalTuple);
+		List	   *resulttlist;
+
+		/*
+		 * Re-fetch the (possibly modified) output tlist of the final
+		 * statement.  By this point, we should have thrown an error if there
+		 * is not one.
+		 */
+		resulttlist = get_sql_fn_result_tlist(llast_node(List, queryTree_list));
 
 		/*
 		 * If the result is composite, *and* we are returning the whole tuple
@@ -1541,29 +1552,39 @@ check_sql_fn_statements(List *queryTreeLists)
 	foreach(lc, queryTreeLists)
 	{
 		List	   *sublist = lfirst_node(List, lc);
-		ListCell   *lc2;
 
-		foreach(lc2, sublist)
-		{
-			Query	   *query = lfirst_node(Query, lc2);
+		check_sql_fn_statement(sublist);
+	}
+}
 
-			/*
-			 * Disallow calling procedures with output arguments.  The current
-			 * implementation would just throw the output values away, unless
-			 * the statement is the last one.  Per SQL standard, we should
-			 * assign the output values by name.  By disallowing this here, we
-			 * preserve an opportunity for future improvement.
-			 */
-			if (query->commandType == CMD_UTILITY &&
-				IsA(query->utilityStmt, CallStmt))
-			{
-				CallStmt   *stmt = (CallStmt *) query->utilityStmt;
+/*
+ * As above, for a single sublist of Queries.
+ */
+static void
+check_sql_fn_statement(List *queryTreeList)
+{
+	ListCell   *lc;
 
-				if (stmt->outargs != NIL)
-					ereport(ERROR,
-							(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
-							 errmsg("calling procedures with output arguments is not supported in SQL functions")));
-			}
+	foreach(lc, queryTreeList)
+	{
+		Query	   *query = lfirst_node(Query, lc);
+
+		/*
+		 * Disallow calling procedures with output arguments.  The current
+		 * implementation would just throw the output values away, unless the
+		 * statement is the last one.  Per SQL standard, we should assign the
+		 * output values by name.  By disallowing this here, we preserve an
+		 * opportunity for future improvement.
+		 */
+		if (query->commandType == CMD_UTILITY &&
+			IsA(query->utilityStmt, CallStmt))
+		{
+			CallStmt   *stmt = (CallStmt *) query->utilityStmt;
+
+			if (stmt->outargs != NIL)
+				ereport(ERROR,
+						(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+						 errmsg("calling procedures with output arguments is not supported in SQL functions")));
 		}
 	}
 }
@@ -1602,17 +1623,45 @@ check_sql_fn_statements(List *queryTreeLists)
  * In addition to coercing individual output columns, we can modify the
  * output to include dummy NULL columns for any dropped columns appearing
  * in rettupdesc.  This is done only if the caller asks for it.
- *
- * If resultTargetList isn't NULL, then *resultTargetList is set to the
- * targetlist that defines the final statement's result.  Exception: if the
- * function is defined to return VOID then *resultTargetList is set to NIL.
  */
 bool
 check_sql_fn_retval(List *queryTreeLists,
 					Oid rettype, TupleDesc rettupdesc,
 					char prokind,
-					bool insertDroppedCols,
-					List **resultTargetList)
+					bool insertDroppedCols)
+{
+	List	   *queryTreeList;
+
+	/*
+	 * We consider only the last sublist of Query nodes, so that only the last
+	 * original statement is a candidate to produce the result.  This is a
+	 * change from pre-v18 versions, which would back up to the last statement
+	 * that includes a canSetTag query, thus ignoring any ending statement(s)
+	 * that rewrite to DO INSTEAD NOTHING.  That behavior was undocumented and
+	 * there seems no good reason for it, except that it was an artifact of
+	 * the original coding.
+	 *
+	 * If the function body is completely empty, handle that the same as if
+	 * the last query had rewritten to nothing.
+	 */
+	if (queryTreeLists != NIL)
+		queryTreeList = llast_node(List, queryTreeLists);
+	else
+		queryTreeList = NIL;
+
+	return check_sql_stmt_retval(queryTreeList,
+								 rettype, rettupdesc,
+								 prokind, insertDroppedCols);
+}
+
+/*
+ * As for check_sql_fn_retval, but we are given just the last query's
+ * rewritten-queries list.
+ */
+static bool
+check_sql_stmt_retval(List *queryTreeList,
+					  Oid rettype, TupleDesc rettupdesc,
+					  char prokind, bool insertDroppedCols)
 {
 	bool		is_tuple_result = false;
 	Query	   *parse;
@@ -1625,9 +1674,6 @@ check_sql_fn_retval(List *queryTreeLists,
 	bool		upper_tlist_nontrivial = false;
 	ListCell   *lc;
 
-	if (resultTargetList)
-		*resultTargetList = NIL;	/* initialize in case of VOID result */
-
 	/*
 	 * If it's declared to return VOID, we don't care what's in the function.
 	 * (This takes care of procedures with no output parameters, as well.)
@@ -1636,30 +1682,20 @@ check_sql_fn_retval(List *queryTreeLists,
 		return false;
 
 	/*
-	 * Find the last canSetTag query in the function body (which is presented
-	 * to us as a list of sublists of Query nodes).  This isn't necessarily
-	 * the last parsetree, because rule rewriting can insert queries after
-	 * what the user wrote.  Note that it might not even be in the last
-	 * sublist, for example if the last query rewrites to DO INSTEAD NOTHING.
-	 * (It might not be unreasonable to throw an error in such a case, but
-	 * this is the historical behavior and it doesn't seem worth changing.)
+	 * Find the last canSetTag query in the list of Query nodes.  This isn't
+	 * necessarily the last parsetree, because rule rewriting can insert
+	 * queries after what the user wrote.
 	 */
 	parse = NULL;
 	parse_cell = NULL;
-	foreach(lc, queryTreeLists)
+	foreach(lc, queryTreeList)
 	{
-		List	   *sublist = lfirst_node(List, lc);
-		ListCell   *lc2;
+		Query	   *q = lfirst_node(Query, lc);
 
-		foreach(lc2, sublist)
+		if (q->canSetTag)
 		{
-			Query	   *q = lfirst_node(Query, lc2);
-
-			if (q->canSetTag)
-			{
-				parse = q;
-				parse_cell = lc2;
-			}
+			parse = q;
+			parse_cell = lc;
 		}
 	}
 
@@ -1812,12 +1848,7 @@ check_sql_fn_retval(List *queryTreeLists,
 		 * further checking.  Assume we're returning the whole tuple.
 		 */
 		if (rettupdesc == NULL)
-		{
-			/* Return tlist if requested */
-			if (resultTargetList)
-				*resultTargetList = tlist;
 			return true;
-		}
 
 		/*
 		 * Verify that the targetlist matches the return tuple type.  We scan
@@ -1984,10 +2015,6 @@ tlist_coercion_finished:
 		lfirst(parse_cell) = newquery;
 	}
 
-	/* Return tlist (possibly modified) if requested */
-	if (resultTargetList)
-		*resultTargetList = upper_tlist;
-
 	return is_tuple_result;
 }
 
@@ -2063,6 +2090,37 @@ coerce_fn_result_column(TargetEntry *src_tle,
 	return true;
 }
 
+/*
+ * Extract the targetlist of the last canSetTag query in the given list
+ * of parsed-and-rewritten Queries.  Returns NIL if there is none.
+ */
+static List *
+get_sql_fn_result_tlist(List *queryTreeList)
+{
+	Query	   *parse = NULL;
+	ListCell   *lc;
+
+	foreach(lc, queryTreeList)
+	{
+		Query	   *q = lfirst_node(Query, lc);
+
+		if (q->canSetTag)
+			parse = q;
+	}
+	if (parse &&
+		parse->commandType == CMD_SELECT)
+		return parse->targetList;
+	else if (parse &&
+			 (parse->commandType == CMD_INSERT ||
+			  parse->commandType == CMD_UPDATE ||
+			  parse->commandType == CMD_DELETE ||
+			  parse->commandType == CMD_MERGE) &&
+			 parse->returningList)
+		return parse->returningList;
+	else
+		return NIL;
+}
+
 
 /*
  * CreateSQLFunctionDestReceiver -- create a suitable DestReceiver object
diff --git a/src/backend/optimizer/util/clauses.c b/src/backend/optimizer/util/clauses.c
index 43dfecfb47f..816536ab865 100644
--- a/src/backend/optimizer/util/clauses.c
+++ b/src/backend/optimizer/util/clauses.c
@@ -4742,7 +4742,7 @@ inline_function(Oid funcid, Oid result_type, Oid result_collid,
 	if (check_sql_fn_retval(list_make1(querytree_list),
 							result_type, rettupdesc,
 							funcform->prokind,
-							false, NULL))
+							false))
 		goto fail;				/* reject whole-tuple-result cases */
 
 	/*
@@ -5288,7 +5288,7 @@ inline_set_returning_function(PlannerInfo *root, RangeTblEntry *rte)
 	if (!check_sql_fn_retval(list_make1(querytree_list),
 							 fexpr->funcresulttype, rettupdesc,
 							 funcform->prokind,
-							 true, NULL) &&
+							 true) &&
 		(functypclass == TYPEFUNC_COMPOSITE ||
 		 functypclass == TYPEFUNC_COMPOSITE_DOMAIN ||
 		 functypclass == TYPEFUNC_RECORD))
diff --git a/src/include/executor/functions.h b/src/include/executor/functions.h
index a6ae2e72d79..58bdff9b039 100644
--- a/src/include/executor/functions.h
+++ b/src/include/executor/functions.h
@@ -48,8 +48,7 @@ extern void check_sql_fn_statements(List *queryTreeLists);
 extern bool check_sql_fn_retval(List *queryTreeLists,
 								Oid rettype, TupleDesc rettupdesc,
 								char prokind,
-								bool insertDroppedCols,
-								List **resultTargetList);
+								bool insertDroppedCols);
 
 extern DestReceiver *CreateSQLFunctionDestReceiver(void);
 
-- 
2.43.5

