Avoid using a cursor in plpgsql's RETURN QUERY statement.

plpgsql has always executed the query given in a RETURN QUERY command
by opening it as a cursor and then fetching a few rows at a time,
which it turns around and dumps into the function's result tuplestore.
The point of this was to keep from blowing out memory with an oversized
SPITupleTable result (note that while a tuplestore can spill tuples
to disk, SPITupleTable cannot).  However, it's rather inefficient, both
because of extra data copying and because of executor entry/exit
overhead.  In recent versions, a new performance problem has emerged:
use of a cursor prevents use of a parallel plan for the executed query.

We can improve matters by skipping use of a cursor and having the
executor push result tuples directly into the function's result
tuplestore.  However, a moderate amount of new infrastructure is needed
to make that idea work:

* We can use the existing tstoreReceiver.c DestReceiver code to funnel
executor output to the tuplestore, but it has to be extended to support
plpgsql's requirement for possibly applying a tuple conversion map.

* SPI needs to be extended to allow use of a caller-supplied
DestReceiver instead of its usual receiver that puts tuples into
a SPITupleTable.  Two new API calls are needed to handle both the
RETURN QUERY and RETURN QUERY EXECUTE cases.

I also felt that I didn't want these new API calls to use the legacy
method of specifying query parameter values with "char" null flags
(the old ' '/'n' convention); rather they should accept ParamListInfo
objects containing the parameter type and value info.  This required
a bit of additional new infrastructure since we didn't yet have any
parse analysis callback that would interpret $N parameter symbols
according to type data supplied in a ParamListInfo.  There seems to be
no harm in letting makeParamList install that callback by default,
rather than leaving a new ParamListInfo's parserSetup hook as NULL.
(Indeed, as of HEAD, I couldn't find anyplace that was using the
parserSetup field at all; plpgsql was using parserSetupArg for its
own purposes, but parserSetup seemed to be write-only.)

We can actually get plpgsql out of the business of using legacy null
flags altogether, and using ParamListInfo instead of its ad-hoc
PreparedParamsData structure; but this requires inventing one more
SPI API call that can replace SPI_cursor_open_with_args.  That seems
worth doing, though.

SPI_execute_with_args and SPI_cursor_open_with_args are now unused
anywhere in the core PG distribution.  Perhaps someday we could
deprecate/remove them.  But cleaning up the crufty bits of the SPI
API is a task for a different patch.

Per bug #16040 from Jeremy Smith.  This is unfortunately too invasive to
consider back-patching.  Patch by me; thanks to Hamid Akhtar for review.

Discussion: https://postgr.es/m/16040-eaacad11fecfb198@postgresql.org

1 files changed, 151 insertions, 129 deletions

diff --git a/src/pl/plpgsql/src/pl_exec.c b/src/pl/plpgsql/src/pl_exec.c
index 9a0fa14ec88..f41d675d656 100644
--- a/src/pl/plpgsql/src/pl_exec.c
+++ b/src/pl/plpgsql/src/pl_exec.c
@@ -27,6 +27,7 @@
 #include "executor/execExpr.h"
 #include "executor/spi.h"
 #include "executor/spi_priv.h"
+#include "executor/tstoreReceiver.h"
 #include "funcapi.h"
 #include "mb/stringinfo_mb.h"
 #include "miscadmin.h"
@@ -51,14 +52,6 @@
 #include "utils/syscache.h"
 #include "utils/typcache.h"
 
-typedef struct
-{
-	int			nargs;			/* number of arguments */
-	Oid		   *types;			/* types of arguments */
-	Datum	   *values;			/* evaluated argument values */
-	char	   *nulls;			/* null markers (' '/'n' style) */
-} PreparedParamsData;
-
 /*
  * All plpgsql function executions within a single transaction share the same
  * executor EState for evaluating "simple" expressions.  Each function call
@@ -441,15 +434,15 @@ static void assign_text_var(PLpgSQL_execstate *estate, PLpgSQL_var *var,
 							const char *str);
 static void assign_record_var(PLpgSQL_execstate *estate, PLpgSQL_rec *rec,
 							  ExpandedRecordHeader *erh);
-static PreparedParamsData *exec_eval_using_params(PLpgSQL_execstate *estate,
-												  List *params);
+static ParamListInfo exec_eval_using_params(PLpgSQL_execstate *estate,
+											List *params);
 static Portal exec_dynquery_with_params(PLpgSQL_execstate *estate,
 										PLpgSQL_expr *dynquery, List *params,
 										const char *portalname, int cursorOptions);
 static char *format_expr_params(PLpgSQL_execstate *estate,
 								const PLpgSQL_expr *expr);
 static char *format_preparedparamsdata(PLpgSQL_execstate *estate,
-									   const PreparedParamsData *ppd);
+									   ParamListInfo paramLI);
 
 
 /* ----------
@@ -3513,9 +3506,11 @@ static int
 exec_stmt_return_query(PLpgSQL_execstate *estate,
 					   PLpgSQL_stmt_return_query *stmt)
 {
-	Portal		portal;
-	uint64		processed = 0;
-	TupleConversionMap *tupmap;
+	int64		tcount;
+	DestReceiver *treceiver;
+	int			rc;
+	uint64		processed;
+	MemoryContext stmt_mcontext = get_stmt_mcontext(estate);
 	MemoryContext oldcontext;
 
 	if (!estate->retisset)
@@ -3525,60 +3520,99 @@ exec_stmt_return_query(PLpgSQL_execstate *estate,
 
 	if (estate->tuple_store == NULL)
 		exec_init_tuple_store(estate);
+	/* There might be some tuples in the tuplestore already */
+	tcount = tuplestore_tuple_count(estate->tuple_store);
+
+	/*
+	 * Set up DestReceiver to transfer results directly to tuplestore,
+	 * converting rowtype if necessary.  DestReceiver lives in mcontext.
+	 */
+	oldcontext = MemoryContextSwitchTo(stmt_mcontext);
+	treceiver = CreateDestReceiver(DestTuplestore);
+	SetTuplestoreDestReceiverParams(treceiver,
+									estate->tuple_store,
+									estate->tuple_store_cxt,
+									false,
+									estate->tuple_store_desc,
+									gettext_noop("structure of query does not match function result type"));
+	MemoryContextSwitchTo(oldcontext);
 
 	if (stmt->query != NULL)
 	{
 		/* static query */
-		exec_run_select(estate, stmt->query, 0, &portal);
-	}
-	else
-	{
-		/* RETURN QUERY EXECUTE */
-		Assert(stmt->dynquery != NULL);
-		portal = exec_dynquery_with_params(estate, stmt->dynquery,
-										   stmt->params, NULL,
-										   0);
-	}
+		PLpgSQL_expr *expr = stmt->query;
+		ParamListInfo paramLI;
 
-	/* Use eval_mcontext for tuple conversion work */
-	oldcontext = MemoryContextSwitchTo(get_eval_mcontext(estate));
+		/*
+		 * On the first call for this expression generate the plan.
+		 */
+		if (expr->plan == NULL)
+			exec_prepare_plan(estate, expr, CURSOR_OPT_PARALLEL_OK, true);
 
-	tupmap = convert_tuples_by_position(portal->tupDesc,
-										estate->tuple_store_desc,
-										gettext_noop("structure of query does not match function result type"));
+		/*
+		 * Set up ParamListInfo to pass to executor
+		 */
+		paramLI = setup_param_list(estate, expr);
 
-	while (true)
+		/*
+		 * Execute the query
+		 */
+		rc = SPI_execute_plan_with_receiver(expr->plan, paramLI,
+											estate->readonly_func, 0,
+											treceiver);
+		if (rc != SPI_OK_SELECT)
+			ereport(ERROR,
+					(errcode(ERRCODE_SYNTAX_ERROR),
+					 errmsg("query \"%s\" is not a SELECT", expr->query)));
+	}
+	else
 	{
-		uint64		i;
-
-		SPI_cursor_fetch(portal, true, 50);
+		/* RETURN QUERY EXECUTE */
+		Datum		query;
+		bool		isnull;
+		Oid			restype;
+		int32		restypmod;
+		char	   *querystr;
 
-		/* SPI will have changed CurrentMemoryContext */
-		MemoryContextSwitchTo(get_eval_mcontext(estate));
+		/*
+		 * Evaluate the string expression after the EXECUTE keyword. Its
+		 * result is the querystring we have to execute.
+		 */
+		Assert(stmt->dynquery != NULL);
+		query = exec_eval_expr(estate, stmt->dynquery,
+							   &isnull, &restype, &restypmod);
+		if (isnull)
+			ereport(ERROR,
+					(errcode(ERRCODE_NULL_VALUE_NOT_ALLOWED),
+					 errmsg("query string argument of EXECUTE is null")));
 
-		if (SPI_processed == 0)
-			break;
+		/* Get the C-String representation */
+		querystr = convert_value_to_string(estate, query, restype);
 
-		for (i = 0; i < SPI_processed; i++)
-		{
-			HeapTuple	tuple = SPI_tuptable->vals[i];
+		/* copy it into the stmt_mcontext before we clean up */
+		querystr = MemoryContextStrdup(stmt_mcontext, querystr);
 
-			if (tupmap)
-				tuple = execute_attr_map_tuple(tuple, tupmap);
-			tuplestore_puttuple(estate->tuple_store, tuple);
-			if (tupmap)
-				heap_freetuple(tuple);
-			processed++;
-		}
+		exec_eval_cleanup(estate);
 
-		SPI_freetuptable(SPI_tuptable);
+		/* Execute query, passing params if necessary */
+		rc = SPI_execute_with_receiver(querystr,
+									   exec_eval_using_params(estate,
+															  stmt->params),
+									   estate->readonly_func,
+									   0,
+									   treceiver);
+		if (rc < 0)
+			elog(ERROR, "SPI_execute_with_receiver failed executing query \"%s\": %s",
+				 querystr, SPI_result_code_string(rc));
 	}
 
-	SPI_freetuptable(SPI_tuptable);
-	SPI_cursor_close(portal);
-
-	MemoryContextSwitchTo(oldcontext);
+	/* Clean up */
+	treceiver->rDestroy(treceiver);
 	exec_eval_cleanup(estate);
+	MemoryContextReset(stmt_mcontext);
+
+	/* Count how many tuples we got */
+	processed = tuplestore_tuple_count(estate->tuple_store) - tcount;
 
 	estate->eval_processed = processed;
 	exec_set_found(estate, processed != 0);
@@ -4344,7 +4378,7 @@ exec_stmt_dynexecute(PLpgSQL_execstate *estate,
 	int32		restypmod;
 	char	   *querystr;
 	int			exec_res;
-	PreparedParamsData *ppd = NULL;
+	ParamListInfo paramLI;
 	MemoryContext stmt_mcontext = get_stmt_mcontext(estate);
 
 	/*
@@ -4368,16 +4402,9 @@ exec_stmt_dynexecute(PLpgSQL_execstate *estate,
 	/*
 	 * Execute the query without preparing a saved plan.
 	 */
-	if (stmt->params)
-	{
-		ppd = exec_eval_using_params(estate, stmt->params);
-		exec_res = SPI_execute_with_args(querystr,
-										 ppd->nargs, ppd->types,
-										 ppd->values, ppd->nulls,
-										 estate->readonly_func, 0);
-	}
-	else
-		exec_res = SPI_execute(querystr, estate->readonly_func, 0);
+	paramLI = exec_eval_using_params(estate, stmt->params);
+	exec_res = SPI_execute_with_receiver(querystr, paramLI,
+										 estate->readonly_func, 0, NULL);
 
 	switch (exec_res)
 	{
@@ -4429,7 +4456,7 @@ exec_stmt_dynexecute(PLpgSQL_execstate *estate,
 			break;
 
 		default:
-			elog(ERROR, "SPI_execute failed executing query \"%s\": %s",
+			elog(ERROR, "SPI_execute_with_receiver failed executing query \"%s\": %s",
 				 querystr, SPI_result_code_string(exec_res));
 			break;
 	}
@@ -4465,7 +4492,7 @@ exec_stmt_dynexecute(PLpgSQL_execstate *estate,
 				char	   *errdetail;
 
 				if (estate->func->print_strict_params)
-					errdetail = format_preparedparamsdata(estate, ppd);
+					errdetail = format_preparedparamsdata(estate, paramLI);
 				else
 					errdetail = NULL;
 
@@ -4484,7 +4511,7 @@ exec_stmt_dynexecute(PLpgSQL_execstate *estate,
 				char	   *errdetail;
 
 				if (estate->func->print_strict_params)
-					errdetail = format_preparedparamsdata(estate, ppd);
+					errdetail = format_preparedparamsdata(estate, paramLI);
 				else
 					errdetail = NULL;
 
@@ -6308,9 +6335,9 @@ exec_eval_simple_expr(PLpgSQL_execstate *estate,
 /*
  * Create a ParamListInfo to pass to SPI
  *
- * We use a single ParamListInfo struct for all SPI calls made from this
- * estate; it contains no per-param data, just hook functions, so it's
- * effectively read-only for SPI.
+ * We use a single ParamListInfo struct for all SPI calls made to evaluate
+ * PLpgSQL_exprs in this estate.  It contains no per-param data, just hook
+ * functions, so it's effectively read-only for SPI.
  *
  * An exception from pure read-only-ness is that the parserSetupArg points
  * to the specific PLpgSQL_expr being evaluated.  This is not an issue for
@@ -8575,65 +8602,68 @@ assign_record_var(PLpgSQL_execstate *estate, PLpgSQL_rec *rec,
  * The result data structure is created in the stmt_mcontext, and should
  * be freed by resetting that context.
  */
-static PreparedParamsData *
+static ParamListInfo
 exec_eval_using_params(PLpgSQL_execstate *estate, List *params)
 {
-	PreparedParamsData *ppd;
-	MemoryContext stmt_mcontext = get_stmt_mcontext(estate);
+	ParamListInfo paramLI;
 	int			nargs;
+	MemoryContext stmt_mcontext;
+	MemoryContext oldcontext;
 	int			i;
 	ListCell   *lc;
 
-	ppd = (PreparedParamsData *)
-		MemoryContextAlloc(stmt_mcontext, sizeof(PreparedParamsData));
-	nargs = list_length(params);
+	/* Fast path for no parameters: we can just return NULL */
+	if (params == NIL)
+		return NULL;
 
-	ppd->nargs = nargs;
-	ppd->types = (Oid *)
-		MemoryContextAlloc(stmt_mcontext, nargs * sizeof(Oid));
-	ppd->values = (Datum *)
-		MemoryContextAlloc(stmt_mcontext, nargs * sizeof(Datum));
-	ppd->nulls = (char *)
-		MemoryContextAlloc(stmt_mcontext, nargs * sizeof(char));
+	nargs = list_length(params);
+	stmt_mcontext = get_stmt_mcontext(estate);
+	oldcontext = MemoryContextSwitchTo(stmt_mcontext);
+	paramLI = makeParamList(nargs);
+	MemoryContextSwitchTo(oldcontext);
 
 	i = 0;
 	foreach(lc, params)
 	{
 		PLpgSQL_expr *param = (PLpgSQL_expr *) lfirst(lc);
-		bool		isnull;
+		ParamExternData *prm = &paramLI->params[i];
 		int32		ppdtypmod;
-		MemoryContext oldcontext;
 
-		ppd->values[i] = exec_eval_expr(estate, param,
-										&isnull,
-										&ppd->types[i],
-										&ppdtypmod);
-		ppd->nulls[i] = isnull ? 'n' : ' ';
+		/*
+		 * Always mark params as const, since we only use the result with
+		 * one-shot plans.
+		 */
+		prm->pflags = PARAM_FLAG_CONST;
+
+		prm->value = exec_eval_expr(estate, param,
+									&prm->isnull,
+									&prm->ptype,
+									&ppdtypmod);
 
 		oldcontext = MemoryContextSwitchTo(stmt_mcontext);
 
-		if (ppd->types[i] == UNKNOWNOID)
+		if (prm->ptype == UNKNOWNOID)
 		{
 			/*
 			 * Treat 'unknown' parameters as text, since that's what most
-			 * people would expect. SPI_execute_with_args can coerce unknown
+			 * people would expect.  The SPI functions can coerce unknown
 			 * constants in a more intelligent way, but not unknown Params.
 			 * This code also takes care of copying into the right context.
 			 * Note we assume 'unknown' has the representation of C-string.
 			 */
-			ppd->types[i] = TEXTOID;
-			if (!isnull)
-				ppd->values[i] = CStringGetTextDatum(DatumGetCString(ppd->values[i]));
+			prm->ptype = TEXTOID;
+			if (!prm->isnull)
+				prm->value = CStringGetTextDatum(DatumGetCString(prm->value));
 		}
 		/* pass-by-ref non null values must be copied into stmt_mcontext */
-		else if (!isnull)
+		else if (!prm->isnull)
 		{
 			int16		typLen;
 			bool		typByVal;
 
-			get_typlenbyval(ppd->types[i], &typLen, &typByVal);
+			get_typlenbyval(prm->ptype, &typLen, &typByVal);
 			if (!typByVal)
-				ppd->values[i] = datumCopy(ppd->values[i], typByVal, typLen);
+				prm->value = datumCopy(prm->value, typByVal, typLen);
 		}
 
 		MemoryContextSwitchTo(oldcontext);
@@ -8643,7 +8673,7 @@ exec_eval_using_params(PLpgSQL_execstate *estate, List *params)
 		i++;
 	}
 
-	return ppd;
+	return paramLI;
 }
 
 /*
@@ -8689,30 +8719,15 @@ exec_dynquery_with_params(PLpgSQL_execstate *estate,
 
 	/*
 	 * Open an implicit cursor for the query.  We use
-	 * SPI_cursor_open_with_args even when there are no params, because this
-	 * avoids making and freeing one copy of the plan.
+	 * SPI_cursor_parse_open_with_paramlist even when there are no params,
+	 * because this avoids making and freeing one copy of the plan.
 	 */
-	if (params)
-	{
-		PreparedParamsData *ppd;
-
-		ppd = exec_eval_using_params(estate, params);
-		portal = SPI_cursor_open_with_args(portalname,
-										   querystr,
-										   ppd->nargs, ppd->types,
-										   ppd->values, ppd->nulls,
-										   estate->readonly_func,
-										   cursorOptions);
-	}
-	else
-	{
-		portal = SPI_cursor_open_with_args(portalname,
-										   querystr,
-										   0, NULL,
-										   NULL, NULL,
-										   estate->readonly_func,
-										   cursorOptions);
-	}
+	portal = SPI_cursor_parse_open_with_paramlist(portalname,
+												  querystr,
+												  exec_eval_using_params(estate,
+																		 params),
+												  estate->readonly_func,
+												  cursorOptions);
 
 	if (portal == NULL)
 		elog(ERROR, "could not open implicit cursor for query \"%s\": %s",
@@ -8782,37 +8797,44 @@ format_expr_params(PLpgSQL_execstate *estate,
 }
 
 /*
- * Return a formatted string with information about PreparedParamsData, or NULL
- * if there are no parameters.
+ * Return a formatted string with information about the parameter values,
+ * or NULL if there are no parameters.
  * The result is in the eval_mcontext.
  */
 static char *
 format_preparedparamsdata(PLpgSQL_execstate *estate,
-						  const PreparedParamsData *ppd)
+						  ParamListInfo paramLI)
 {
 	int			paramno;
 	StringInfoData paramstr;
 	MemoryContext oldcontext;
 
-	if (!ppd)
+	if (!paramLI)
 		return NULL;
 
 	oldcontext = MemoryContextSwitchTo(get_eval_mcontext(estate));
 
 	initStringInfo(&paramstr);
-	for (paramno = 0; paramno < ppd->nargs; paramno++)
+	for (paramno = 0; paramno < paramLI->numParams; paramno++)
 	{
+		ParamExternData *prm = &paramLI->params[paramno];
+
+		/*
+		 * Note: for now, this is only used on ParamListInfos produced by
+		 * exec_eval_using_params(), so we don't worry about invoking the
+		 * paramFetch hook or skipping unused parameters.
+		 */
 		appendStringInfo(&paramstr, "%s$%d = ",
 						 paramno > 0 ? ", " : "",
 						 paramno + 1);
 
-		if (ppd->nulls[paramno] == 'n')
+		if (prm->isnull)
 			appendStringInfoString(&paramstr, "NULL");
 		else
 			appendStringInfoStringQuoted(&paramstr,
 										 convert_value_to_string(estate,
-																 ppd->values[paramno],
-																 ppd->types[paramno]),
+																 prm->value,
+																 prm->ptype),
 										 -1);
 	}