Support RETURN QUERY for PL/PgSQL (#644)

Backport from upstream and related commits are

    commit e5fe2e84
    Author: Tom Lane <tgl@sss.pgh.pa.us>
    Date:   Fri Nov 9 23:58:32 2007 +0000

        Recognize RETURN QUERY via a textual test, so that QUERY doesn't need to be
        a plpgsql keyword.  This avoids springing a new reserved word on plpgsql
        programmers.
        For consistency, handle RETURN NEXT the same way.

    commit b2b9b4d5
    Author: Neil Conway <neilc@samurai.com>
    Date:   Wed Jul 25 04:19:09 2007 +0000

        Implement RETURN QUERY for PL/PgSQL. This provides some convenient syntax
        sugar for PL/PgSQL set-returning functions that want to return the result
        of evaluating a query; it should also be more efficient than repeated
        RETURN NEXT statements. Based on an earlier patch from Pavel Stehule.

Committers: Kuien Liu and Haozhou Wang

src/pl/plpgsql/src/gram.y

@@ -45,6 +45,7 @@ static	PLpgSQL_stmt	*make_execsql_stmt(const char *sqlstart, int lineno);
 static	PLpgSQL_stmt	*make_fetch_stmt(int lineno, int curvar);
 static	PLpgSQL_stmt	*make_return_stmt(int lineno);
 static	PLpgSQL_stmt	*make_return_next_stmt(int lineno);
+static	PLpgSQL_stmt	*make_return_query_stmt(int lineno);
 static	void			 check_assignable(PLpgSQL_datum *datum);
 static	void			 read_into_target(PLpgSQL_rec **rec, PLpgSQL_row **row,
 										  bool *strict);
@@ -189,7 +190,6 @@ static	void			 check_labels(const char *start_label,
 %token	K_IS
 %token	K_LOG
 %token	K_LOOP
-%token	K_NEXT
 %token	K_NOT
 %token	K_NOTICE
 %token	K_NULL
@@ -1167,10 +1167,21 @@ stmt_return		: K_RETURN lno
 						int	tok;
 
 						tok = yylex();
-						if (tok == K_NEXT)
+						if (tok == 0)
+							yyerror("unexpected end of function definition");
+						/*
+						 * To avoid making NEXT and QUERY effectively be
+						 * reserved words within plpgsql, recognize them
+						 * via yytext.
+						 */
+						if (pg_strcasecmp(yytext, "next") == 0)
 						{
 							$$ = make_return_next_stmt($2);
 						}
+						else if (pg_strcasecmp(yytext, "query") == 0)
+						{
+							$$ = make_return_query_stmt($2);
+						}
 						else
 						{
 							plpgsql_push_back_token(tok);
@@ -2017,7 +2028,8 @@ make_return_stmt(int lineno)
 	if (plpgsql_curr_compile->fn_retset)
 	{
 		if (yylex() != ';')
-			yyerror("RETURN cannot have a parameter in function returning set; use RETURN NEXT");
+			yyerror("RETURN cannot have a parameter in function "
+					"returning set; use RETURN NEXT or RETURN QUERY");
 	}
 	else if (plpgsql_curr_compile->out_param_varno >= 0)
 	{
@@ -2113,6 +2125,23 @@ make_return_next_stmt(int lineno)
 }
 
 
+static PLpgSQL_stmt *
+make_return_query_stmt(int lineno)
+{
+	PLpgSQL_stmt_return_query *new;
+
+	if (!plpgsql_curr_compile->fn_retset)
+		yyerror("cannot use RETURN QUERY in a non-SETOF function");
+
+	new = palloc0(sizeof(PLpgSQL_stmt_return_query));
+	new->cmd_type	= PLPGSQL_STMT_RETURN_QUERY;
+	new->lineno		= lineno;
+	new->query		= read_sql_construct(';', 0, ")", "", false, true, NULL);
+
+	return (PLpgSQL_stmt *) new;
+}
+
+
 static void
 check_assignable(PLpgSQL_datum *datum)
 {

src/pl/plpgsql/src/pl_exec.c

@@ -106,6 +106,8 @@ static int exec_stmt_return(PLpgSQL_execstate *estate,
 				 PLpgSQL_stmt_return *stmt);
 static int exec_stmt_return_next(PLpgSQL_execstate *estate,
 					  PLpgSQL_stmt_return_next *stmt);
+static int exec_stmt_return_query(PLpgSQL_execstate *estate,
+					  PLpgSQL_stmt_return_query *stmt);
 static int exec_stmt_raise(PLpgSQL_execstate *estate,
 				PLpgSQL_stmt_raise *stmt);
 static int exec_stmt_execsql(PLpgSQL_execstate *estate,
@@ -1249,6 +1251,10 @@ exec_stmt(PLpgSQL_execstate *estate, PLpgSQL_stmt *stmt)
 			rc = exec_stmt_return_next(estate, (PLpgSQL_stmt_return_next *) stmt);
 			break;
 
+		case PLPGSQL_STMT_RETURN_QUERY:
+			rc = exec_stmt_return_query(estate, (PLpgSQL_stmt_return_query *) stmt);
+			break;
+
 		case PLPGSQL_STMT_RAISE:
 			rc = exec_stmt_raise(estate, (PLpgSQL_stmt_raise *) stmt);
 			break;
@@ -2120,6 +2126,59 @@ exec_stmt_return_next(PLpgSQL_execstate *estate,
 	return PLPGSQL_RC_OK;
 }
 
+/* ----------
+ * exec_stmt_return_query		Evaluate a query and add it to the
+ *								list of tuples returned by the current
+ *								SRF.
+ * ----------
+ */
+static int
+exec_stmt_return_query(PLpgSQL_execstate *estate,
+					   PLpgSQL_stmt_return_query *stmt)
+{
+	Portal	portal;
+
+	if (!estate->retisset)
+		ereport(ERROR,
+				(errcode(ERRCODE_SYNTAX_ERROR),
+				 errmsg("cannot use RETURN QUERY in a non-SETOF function")));
+
+	if (estate->tuple_store == NULL)
+		exec_init_tuple_store(estate);
+
+	exec_run_select(estate, stmt->query, 0, &portal);
+
+	if (!compatible_tupdesc(estate->rettupdesc, portal->tupDesc))
+		ereport(ERROR,
+				(errcode(ERRCODE_DATATYPE_MISMATCH),
+				 errmsg("structure of query does not match function result type")));
+
+	while (true)
+	{
+		MemoryContext	old_cxt;
+		int				i;
+
+		SPI_cursor_fetch(portal, true, 50);
+		if (SPI_processed == 0)
+			break;
+
+		old_cxt = MemoryContextSwitchTo(estate->tuple_store_cxt);
+		for (i = 0; i < SPI_processed; i++)
+		{
+			HeapTuple tuple = SPI_tuptable->vals[i];
+			tuplestore_puttuple(estate->tuple_store, tuple);
+		}
+		MemoryContextSwitchTo(old_cxt);
+
+		SPI_freetuptable(SPI_tuptable);
+	}
+
+	SPI_freetuptable(SPI_tuptable);
+	SPI_cursor_close(portal);
+
+	return PLPGSQL_RC_OK;
+}
+
 static void
 exec_init_tuple_store(PLpgSQL_execstate *estate)
 {

src/pl/plpgsql/src/pl_funcs.c

@@ -452,6 +452,8 @@ plpgsql_stmt_typename(PLpgSQL_stmt *stmt)
 			return _("return");
 		case PLPGSQL_STMT_RETURN_NEXT:
 			return _("return next");
+		case PLPGSQL_STMT_RETURN_QUERY:
+			return "RETURN QUERY";
 		case PLPGSQL_STMT_RAISE:
 			return _("raise");
 		case PLPGSQL_STMT_EXECSQL:
@@ -495,6 +497,7 @@ static void free_fors(PLpgSQL_stmt_fors *stmt);
 static void free_exit(PLpgSQL_stmt_exit *stmt);
 static void free_return(PLpgSQL_stmt_return *stmt);
 static void free_return_next(PLpgSQL_stmt_return_next *stmt);
+static void free_return_query(PLpgSQL_stmt_return_query *stmt);
 static void free_raise(PLpgSQL_stmt_raise *stmt);
 static void free_execsql(PLpgSQL_stmt_execsql *stmt);
 static void free_dynexecute(PLpgSQL_stmt_dynexecute *stmt);
@@ -542,6 +545,9 @@ free_stmt(PLpgSQL_stmt *stmt)
 		case PLPGSQL_STMT_RETURN_NEXT:
 			free_return_next((PLpgSQL_stmt_return_next *) stmt);
 			break;
+		case PLPGSQL_STMT_RETURN_QUERY:
+			free_return_query((PLpgSQL_stmt_return_query *) stmt);
+			break;
 		case PLPGSQL_STMT_RAISE:
 			free_raise((PLpgSQL_stmt_raise *) stmt);
 			break;
@@ -688,6 +694,14 @@ free_return_next(PLpgSQL_stmt_return_next *stmt)
 	free_expr(stmt->expr);
 }
 
+static void
+free_return_query(PLpgSQL_stmt_return_query *stmt)
+{
+	ListCell   *lc;
+
+	free_expr(stmt->query);
+}
+
 static void
 free_raise(PLpgSQL_stmt_raise *stmt)
 {
@@ -804,6 +818,7 @@ static void dump_fors(PLpgSQL_stmt_fors *stmt);
 static void dump_exit(PLpgSQL_stmt_exit *stmt);
 static void dump_return(PLpgSQL_stmt_return *stmt);
 static void dump_return_next(PLpgSQL_stmt_return_next *stmt);
+static void dump_return_query(PLpgSQL_stmt_return_query *stmt);
 static void dump_raise(PLpgSQL_stmt_raise *stmt);
 static void dump_execsql(PLpgSQL_stmt_execsql *stmt);
 static void dump_dynexecute(PLpgSQL_stmt_dynexecute *stmt);
@@ -861,6 +876,9 @@ dump_stmt(PLpgSQL_stmt *stmt)
 		case PLPGSQL_STMT_RETURN_NEXT:
 			dump_return_next((PLpgSQL_stmt_return_next *) stmt);
 			break;
+		case PLPGSQL_STMT_RETURN_QUERY:
+			dump_return_query((PLpgSQL_stmt_return_query *) stmt);
+			break;
 		case PLPGSQL_STMT_RAISE:
 			dump_raise((PLpgSQL_stmt_raise *) stmt);
 			break;
@@ -1154,6 +1172,15 @@ dump_return_next(PLpgSQL_stmt_return_next *stmt)
 	printf("\n");
 }
 
+static void
+dump_return_query(PLpgSQL_stmt_return_query *stmt)
+{
+	dump_ind();
+	printf("RETURN QUERY ");
+	dump_expr(stmt->query);
+	printf("\n");
+}
+
 static void
 dump_raise(PLpgSQL_stmt_raise *stmt)
 {

src/pl/plpgsql/src/plpgsql.h

@@ -83,6 +83,7 @@ enum PLpgSQL_stmt_types
 	PLPGSQL_STMT_EXIT,
 	PLPGSQL_STMT_RETURN,
 	PLPGSQL_STMT_RETURN_NEXT,
+	PLPGSQL_STMT_RETURN_QUERY,
 	PLPGSQL_STMT_RAISE,
 	PLPGSQL_STMT_EXECSQL,
 	PLPGSQL_STMT_DYNEXECUTE,
@@ -488,6 +489,13 @@ typedef struct
 	int			retvarno;
 } PLpgSQL_stmt_return_next;
 
+typedef struct
+{								/* RETURN QUERY statement */
+	int			cmd_type;
+	int			lineno;
+	PLpgSQL_expr *query;
+} PLpgSQL_stmt_return_query;
+
 typedef struct
 {								/* RAISE statement			*/
 	int			cmd_type;

src/pl/plpgsql/src/scan.l

@@ -159,7 +159,6 @@ into			{ return K_INTO;			}
 is				{ return K_IS;				}
 log				{ return K_LOG;				}
 loop			{ return K_LOOP;			}
-next			{ return K_NEXT;			}
 not				{ return K_NOT;				}
 notice			{ return K_NOTICE;			}
 null			{ return K_NULL;			}

src/test/regress/expected/plpgsql.out

@@ -2977,3 +2977,52 @@ NOTICE:  caught division by zero
 NOTICE:  caught division by zero
 NOTICE:  caught division by zero
 NOTICE:  caught division by zero
+-- tests for RETURN QUERY
+create function ret_query1(out int, out int) returns setof record as $$
+begin
+    $1 := -1;
+    $2 := -2;
+    return next;
+    return query select x + 1, x * 10 from generate_series(0, 10) s (x);
+    return next;
+end;
+$$ language plpgsql;
+select * from ret_query1();
+ column1 | column2
+---------+---------
+      -1 |      -2
+       1 |       0
+       2 |      10
+       3 |      20
+       4 |      30
+       5 |      40
+       6 |      50
+       7 |      60
+       8 |      70
+       9 |      80
+      10 |      90
+      11 |     100
+      -1 |      -2
+(13 rows)
+
+create type record_type as (x text, y int, z boolean);
+create or replace function ret_query2(lim int) returns setof record_type as $$
+begin
+    return query select md5(s.x::text), s.x, s.x > 0
+                 from generate_series(-8, lim) s (x) where s.x % 2 = 0;
+end;
+$$ language plpgsql;
+select * from ret_query2(8);
+                x                 | y  | z
+----------------------------------+----+---
+ a8d2ec85eaf98407310b72eb73dda247 | -8 | f
+ 596a3d04481816330f07e4f97510c28f | -6 | f
+ 0267aaf632e87a63288a08331f22c7c3 | -4 | f
+ 5d7b9adcbe1c629ec722529dd12e5129 | -2 | f
+ cfcd208495d565ef66e7dff9f98764da |  0 | f
+ c81e728d9d4c2f636f067f89cc14862c |  2 | t
+ a87ff679a2f3e71d9181a67b7542122c |  4 | t
+ 1679091c5a880faf6fb5e6087eb1b2dc |  6 | t
+ c9f0f895fb98ab9159f51fd0297e236d |  8 | t
+(9 rows)
+

src/test/regress/sql/plpgsql.sql

@@ -2483,3 +2483,27 @@ begin
   end loop;
 end;
 $outer$;
+
+-- tests for RETURN QUERY
+create function ret_query1(out int, out int) returns setof record as $$
+begin
+    $1 := -1;
+    $2 := -2;
+    return next;
+    return query select x + 1, x * 10 from generate_series(0, 10) s (x);
+    return next;
+end;
+$$ language plpgsql;
+
+select * from ret_query1();
+
+create type record_type as (x text, y int, z boolean);
+
+create or replace function ret_query2(lim int) returns setof record_type as $$
+begin
+    return query select md5(s.x::text), s.x, s.x > 0
+                 from generate_series(-8, lim) s (x) where s.x % 2 = 0;
+end;
+$$ language plpgsql;
+
+select * from ret_query2(8);