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.

5 files changed, 115 insertions, 8 deletions

diff --git a/src/pl/plpgsql/src/gram.y b/src/pl/plpgsql/src/gram.y
index 795d3c4c4bd..b0bc0ea304c 100644
--- a/src/pl/plpgsql/src/gram.y
+++ b/src/pl/plpgsql/src/gram.y
@@ -9,7 +9,7 @@
  *
  *
  * IDENTIFICATION
- *	  $PostgreSQL: pgsql/src/pl/plpgsql/src/gram.y,v 1.104 2007/07/16 17:01:10 tgl Exp $
+ *	  $PostgreSQL: pgsql/src/pl/plpgsql/src/gram.y,v 1.105 2007/07/25 04:19:08 neilc Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -32,6 +32,7 @@ static	PLpgSQL_stmt	*make_execsql_stmt(const char *sqlstart, int lineno);
 static	PLpgSQL_stmt_fetch *read_fetch_direction(void);
 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);
@@ -187,6 +188,7 @@ static	void			 check_labels(const char *start_label,
 %token	K_NULL
 %token	K_OPEN
 %token	K_OR
+%token	K_QUERY
 %token	K_PERFORM
 %token	K_ROW_COUNT
 %token	K_RAISE
@@ -1171,6 +1173,10 @@ stmt_return		: K_RETURN lno
 						{
 							$$ = make_return_next_stmt($2);
 						}
+						else if (tok == K_QUERY)
+						{
+							$$ = make_return_query_stmt($2);
+						}
 						else
 						{
 							plpgsql_push_back_token(tok);
@@ -2104,7 +2110,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)
 	{
@@ -2200,6 +2207,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)
 {
diff --git a/src/pl/plpgsql/src/pl_exec.c b/src/pl/plpgsql/src/pl_exec.c
index 9527fdc61d5..dd415047f3e 100644
--- a/src/pl/plpgsql/src/pl_exec.c
+++ b/src/pl/plpgsql/src/pl_exec.c
@@ -8,7 +8,7 @@
  *
  *
  * IDENTIFICATION
- *	  $PostgreSQL: pgsql/src/pl/plpgsql/src/pl_exec.c,v 1.198 2007/07/15 02:15:04 tgl Exp $
+ *	  $PostgreSQL: pgsql/src/pl/plpgsql/src/pl_exec.c,v 1.199 2007/07/25 04:19:08 neilc Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -105,6 +105,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,
@@ -1244,6 +1246,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;
@@ -2137,6 +2143,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)
 {
diff --git a/src/pl/plpgsql/src/pl_funcs.c b/src/pl/plpgsql/src/pl_funcs.c
index 894284fc197..55c8d2eeac7 100644
--- a/src/pl/plpgsql/src/pl_funcs.c
+++ b/src/pl/plpgsql/src/pl_funcs.c
@@ -8,7 +8,7 @@
  *
  *
  * IDENTIFICATION
- *	  $PostgreSQL: pgsql/src/pl/plpgsql/src/pl_funcs.c,v 1.62 2007/07/20 16:23:34 petere Exp $
+ *	  $PostgreSQL: pgsql/src/pl/plpgsql/src/pl_funcs.c,v 1.63 2007/07/25 04:19:08 neilc Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -443,6 +443,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:
@@ -484,6 +486,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);
@@ -542,6 +545,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;
@@ -879,6 +885,15 @@ dump_return_next(PLpgSQL_stmt_return_next *stmt)
 }
 
 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)
 {
 	ListCell   *lc;
diff --git a/src/pl/plpgsql/src/plpgsql.h b/src/pl/plpgsql/src/plpgsql.h
index 6ffb8b41958..4a61379fb55 100644
--- a/src/pl/plpgsql/src/plpgsql.h
+++ b/src/pl/plpgsql/src/plpgsql.h
@@ -8,7 +8,7 @@
  *
  *
  * IDENTIFICATION
- *	  $PostgreSQL: pgsql/src/pl/plpgsql/src/plpgsql.h,v 1.90 2007/07/16 17:01:11 tgl Exp $
+ *	  $PostgreSQL: pgsql/src/pl/plpgsql/src/plpgsql.h,v 1.91 2007/07/25 04:19:09 neilc Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -83,6 +83,7 @@ enum
 	PLPGSQL_STMT_EXIT,
 	PLPGSQL_STMT_RETURN,
 	PLPGSQL_STMT_RETURN_NEXT,
+	PLPGSQL_STMT_RETURN_QUERY,
 	PLPGSQL_STMT_RAISE,
 	PLPGSQL_STMT_EXECSQL,
 	PLPGSQL_STMT_DYNEXECUTE,
@@ -494,6 +495,13 @@ typedef struct
 } 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;
 	int			lineno;
diff --git a/src/pl/plpgsql/src/scan.l b/src/pl/plpgsql/src/scan.l
index 0d71dd6436a..b322a4045e6 100644
--- a/src/pl/plpgsql/src/scan.l
+++ b/src/pl/plpgsql/src/scan.l
@@ -9,7 +9,7 @@
  *
  *
  * IDENTIFICATION
- *	  $PostgreSQL: pgsql/src/pl/plpgsql/src/scan.l,v 1.57 2007/04/29 01:21:09 neilc Exp $
+ *	  $PostgreSQL: pgsql/src/pl/plpgsql/src/scan.l,v 1.58 2007/07/25 04:19:09 neilc Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -124,8 +124,8 @@ declare			{ return K_DECLARE;			}
 default			{ return K_DEFAULT;			}
 diagnostics		{ return K_DIAGNOSTICS;		}
 else			{ return K_ELSE;			}
-elseif          { return K_ELSIF;           }
-elsif           { return K_ELSIF;           }
+elseif			{ return K_ELSIF;			}
+elsif			{ return K_ELSIF;			}
 end				{ return K_END;				}
 exception		{ return K_EXCEPTION;		}
 execute			{ return K_EXECUTE;			}
@@ -151,6 +151,7 @@ null			{ return K_NULL;			}
 open			{ return K_OPEN;			}
 or				{ return K_OR;				}
 perform			{ return K_PERFORM;			}
+query			{ return K_QUERY;			}
 raise			{ return K_RAISE;			}
 rename			{ return K_RENAME;			}
 result_oid		{ return K_RESULT_OID;		}