Improve pl/pgsql to support composite-type expressions in RETURN.

For some reason lost in the mists of prehistory, RETURN was only coded to
allow a simple reference to a composite variable when the function's return
type is composite.  Allow an expression instead, while preserving the
efficiency of the original code path in the case where the expression is
indeed just a composite variable's name.  Likewise for RETURN NEXT.

As is true in various other places, the supplied expression must yield
exactly the number and data types of the required columns.  There was some
discussion of relaxing that for pl/pgsql, but no consensus yet, so this
patch doesn't address that.

Asif Rehman, reviewed by Pavel Stehule

doc/src/sgml/plpgsql.sgml

@@ -1571,11 +1571,11 @@ RETURN <replaceable>expression</replaceable>;
      </para>
 
      <para>
-      When returning a scalar type, any expression can be used. The
-      expression's result will be automatically cast into the
-      function's return type as described for assignments. To return a
-      composite (row) value, you must write a record or row variable
-      as the <replaceable>expression</replaceable>.
+      In a function that returns a scalar type, the expression's result will
+      automatically be cast into the function's return type as described for
+      assignments.  But to return a composite (row) value, you must write an
+      expression delivering exactly the requested column set.  This may
+      require use of explicit casting.
      </para>
 
      <para>
@@ -1600,6 +1600,20 @@ RETURN <replaceable>expression</replaceable>;
       however.  In those cases a <command>RETURN</command> statement is
       automatically executed if the top-level block finishes.
      </para>
+
+     <para>
+      Some examples:
+
+<programlisting>
+-- functions returning a scalar type
+RETURN 1 + 2;
+RETURN scalar_var;
+
+-- functions returning a composite type
+RETURN composite_type_var;
+RETURN (1, 2, 'three'::text);  -- must cast columns to correct types
+</programlisting>
+     </para>
     </sect3>
 
     <sect3>

src/pl/plpgsql/src/pl_gram.y

@@ -2926,32 +2926,27 @@ make_return_stmt(int location)
 	}
 	else if (plpgsql_curr_compile->fn_retistuple)
 	{
-		switch (yylex())
-		{
-			case K_NULL:
-				/* we allow this to support RETURN NULL in triggers */
-				break;
+		/*
+		 * We want to special-case simple row or record references for
+		 * efficiency.  So peek ahead to see if that's what we have.
+		 */
+		int		tok = yylex();
 
-			case T_DATUM:
-				if (yylval.wdatum.datum->dtype == PLPGSQL_DTYPE_ROW ||
-					yylval.wdatum.datum->dtype == PLPGSQL_DTYPE_REC)
+		if (tok == T_DATUM && plpgsql_peek() == ';' &&
+			(yylval.wdatum.datum->dtype == PLPGSQL_DTYPE_ROW ||
+			 yylval.wdatum.datum->dtype == PLPGSQL_DTYPE_REC))
+		{
 			new->retvarno = yylval.wdatum.datum->dno;
+			/* eat the semicolon token that we only peeked at above */
+			tok = yylex();
+			Assert(tok == ';');
+		}
 		else
-					ereport(ERROR,
-							(errcode(ERRCODE_DATATYPE_MISMATCH),
-							 errmsg("RETURN must specify a record or row variable in function returning row"),
-							 parser_errposition(yylloc)));
-				break;
-
-			default:
-				ereport(ERROR,
-						(errcode(ERRCODE_DATATYPE_MISMATCH),
-						 errmsg("RETURN must specify a record or row variable in function returning row"),
-						 parser_errposition(yylloc)));
-				break;
+		{
+			/* Not (just) a row/record name, so treat as expression */
+			plpgsql_push_back_token(tok);
+			new->expr = read_sql_expression(';', ";");
 		}
-		if (yylex() != ';')
-			yyerror("syntax error");
 	}
 	else
 	{
@@ -2994,28 +2989,27 @@ make_return_next_stmt(int location)
 	}
 	else if (plpgsql_curr_compile->fn_retistuple)
 	{
-		switch (yylex())
+		/*
+		 * We want to special-case simple row or record references for
+		 * efficiency.  So peek ahead to see if that's what we have.
+		 */
+		int		tok = yylex();
+
+		if (tok == T_DATUM && plpgsql_peek() == ';' &&
+			(yylval.wdatum.datum->dtype == PLPGSQL_DTYPE_ROW ||
+			 yylval.wdatum.datum->dtype == PLPGSQL_DTYPE_REC))
 		{
-			case T_DATUM:
-				if (yylval.wdatum.datum->dtype == PLPGSQL_DTYPE_ROW ||
-					yylval.wdatum.datum->dtype == PLPGSQL_DTYPE_REC)
 			new->retvarno = yylval.wdatum.datum->dno;
+			/* eat the semicolon token that we only peeked at above */
+			tok = yylex();
+			Assert(tok == ';');
+		}
 		else
-					ereport(ERROR,
-							(errcode(ERRCODE_DATATYPE_MISMATCH),
-							 errmsg("RETURN NEXT must specify a record or row variable in function returning row"),
-							 parser_errposition(yylloc)));
-				break;
-
-			default:
-				ereport(ERROR,
-						(errcode(ERRCODE_DATATYPE_MISMATCH),
-						 errmsg("RETURN NEXT must specify a record or row variable in function returning row"),
-						 parser_errposition(yylloc)));
-				break;
+		{
+			/* Not (just) a row/record name, so treat as expression */
+			plpgsql_push_back_token(tok);
+			new->expr = read_sql_expression(';', ";");
 		}
-		if (yylex() != ';')
-			yyerror("syntax error");
 	}
 	else
 		new->expr = read_sql_expression(';', ";");

src/pl/plpgsql/src/pl_scanner.c

@@ -442,9 +442,27 @@ plpgsql_append_source_text(StringInfo buf,
 						   endlocation - startlocation);
 }
 
+/*
+ * Peek one token ahead in the input stream.  Only the token code is
+ * made available, not any of the auxiliary info such as location.
+ *
+ * NB: no variable or unreserved keyword lookup is performed here, they will
+ * be returned as IDENT. Reserved keywords are resolved as usual.
+ */
+int
+plpgsql_peek(void)
+{
+	int			tok1;
+	TokenAuxData aux1;
+
+	tok1 = internal_yylex(&aux1);
+	push_back_token(tok1, &aux1);
+	return tok1;
+}
+
 /*
  * Peek two tokens ahead in the input stream. The first token and its
- * location the query are returned in *tok1_p and *tok1_loc, second token
+ * location in the query are returned in *tok1_p and *tok1_loc, second token
  * and its location in *tok2_p and *tok2_loc.
  *
  * NB: no variable or unreserved keyword lookup is performed here, they will

src/pl/plpgsql/src/plpgsql.h

@@ -976,6 +976,7 @@ extern void plpgsql_push_back_token(int token);
 extern bool plpgsql_token_is_unreserved_keyword(int token);
 extern void plpgsql_append_source_text(StringInfo buf,
 						   int startlocation, int endlocation);
+extern int	plpgsql_peek(void);
 extern void plpgsql_peek2(int *tok1_p, int *tok2_p, int *tok1_loc,
 			  int *tok2_loc);
 extern int	plpgsql_scanner_errposition(int location);

src/test/regress/expected/plpgsql.out

@@ -3624,7 +3624,139 @@ select * from returnqueryf();
 
 drop function returnqueryf();
 drop table tabwithcols;
+--
+-- Tests for composite-type results
+--
+create type footype as (x int, y varchar);
+-- test: use of variable of composite type in return statement
+create or replace function foo() returns footype as $$
+declare
+  v footype;
+begin
+  v := (1, 'hello');
+  return v;
+end;
+$$ language plpgsql;
+select foo();
+    foo    
+-----------
+ (1,hello)
+(1 row)
+
+-- test: use of variable of record type in return statement
+create or replace function foo() returns footype as $$
+declare
+  v record;
+begin
+  v := (1, 'hello'::varchar);
+  return v;
+end;
+$$ language plpgsql;
+select foo();
+    foo    
+-----------
+ (1,hello)
+(1 row)
+
+-- test: use of row expr in return statement
+create or replace function foo() returns footype as $$
+begin
+  return (1, 'hello'::varchar);
+end;
+$$ language plpgsql;
+select foo();
+    foo    
+-----------
+ (1,hello)
+(1 row)
+
+-- this does not work currently (no implicit casting)
+create or replace function foo() returns footype as $$
+begin
+  return (1, 'hello');
+end;
+$$ language plpgsql;
+select foo();
+ERROR:  returned record type does not match expected record type
+DETAIL:  Returned type unknown does not match expected type character varying in column 2.
+CONTEXT:  PL/pgSQL function foo() while casting return value to function's return type
+-- ... but this does
+create or replace function foo() returns footype as $$
+begin
+  return (1, 'hello')::footype;
+end;
+$$ language plpgsql;
+select foo();
+    foo    
+-----------
+ (1,hello)
+(1 row)
+
+drop function foo();
+-- test: return a row expr as record.
+create or replace function foorec() returns record as $$
+declare
+  v record;
+begin
+  v := (1, 'hello');
+  return v;
+end;
+$$ language plpgsql;
+select foorec();
+  foorec   
+-----------
+ (1,hello)
+(1 row)
+
+-- test: return row expr in return statement.
+create or replace function foorec() returns record as $$
+begin
+  return (1, 'hello');
+end;
+$$ language plpgsql;
+select foorec();
+  foorec   
+-----------
+ (1,hello)
+(1 row)
+
+drop function foorec();
+-- test: row expr in RETURN NEXT statement.
+create or replace function foo() returns setof footype as $$
+begin
+  for i in 1..3
+  loop
+    return next (1, 'hello'::varchar);
+  end loop;
+  return next null::footype;
+  return next (2, 'goodbye')::footype;
+end;
+$$ language plpgsql;
+select * from foo();
+ x |    y    
+---+---------
+ 1 | hello
+ 1 | hello
+ 1 | hello
+   | 
+ 2 | goodbye
+(5 rows)
+
+drop function foo();
+-- test: use invalid expr in return statement.
+create or replace function foo() returns footype as $$
+begin
+  return 1 + 1;
+end;
+$$ language plpgsql;
+select foo();
+ERROR:  cannot return non-composite value from function returning composite type
+CONTEXT:  PL/pgSQL function foo() line 3 at RETURN
+drop function foo();
+drop type footype;
+--
 -- Tests for 8.4's new RAISE features
+--
 create or replace function raise_test() returns void as $$
 begin
   raise notice '% % %', 1, 2, 3

src/test/regress/sql/plpgsql.sql

@@ -2937,7 +2937,119 @@ select * from returnqueryf();
 drop function returnqueryf();
 drop table tabwithcols;
 
+--
+-- Tests for composite-type results
+--
+
+create type footype as (x int, y varchar);
+
+-- test: use of variable of composite type in return statement
+create or replace function foo() returns footype as $$
+declare
+  v footype;
+begin
+  v := (1, 'hello');
+  return v;
+end;
+$$ language plpgsql;
+
+select foo();
+
+-- test: use of variable of record type in return statement
+create or replace function foo() returns footype as $$
+declare
+  v record;
+begin
+  v := (1, 'hello'::varchar);
+  return v;
+end;
+$$ language plpgsql;
+
+select foo();
+
+-- test: use of row expr in return statement
+create or replace function foo() returns footype as $$
+begin
+  return (1, 'hello'::varchar);
+end;
+$$ language plpgsql;
+
+select foo();
+
+-- this does not work currently (no implicit casting)
+create or replace function foo() returns footype as $$
+begin
+  return (1, 'hello');
+end;
+$$ language plpgsql;
+
+select foo();
+
+-- ... but this does
+create or replace function foo() returns footype as $$
+begin
+  return (1, 'hello')::footype;
+end;
+$$ language plpgsql;
+
+select foo();
+
+drop function foo();
+
+-- test: return a row expr as record.
+create or replace function foorec() returns record as $$
+declare
+  v record;
+begin
+  v := (1, 'hello');
+  return v;
+end;
+$$ language plpgsql;
+
+select foorec();
+
+-- test: return row expr in return statement.
+create or replace function foorec() returns record as $$
+begin
+  return (1, 'hello');
+end;
+$$ language plpgsql;
+
+select foorec();
+
+drop function foorec();
+
+-- test: row expr in RETURN NEXT statement.
+create or replace function foo() returns setof footype as $$
+begin
+  for i in 1..3
+  loop
+    return next (1, 'hello'::varchar);
+  end loop;
+  return next null::footype;
+  return next (2, 'goodbye')::footype;
+end;
+$$ language plpgsql;
+
+select * from foo();
+
+drop function foo();
+
+-- test: use invalid expr in return statement.
+create or replace function foo() returns footype as $$
+begin
+  return 1 + 1;
+end;
+$$ language plpgsql;
+
+select foo();
+
+drop function foo();
+drop type footype;
+
+--
 -- Tests for 8.4's new RAISE features
+--
 
 create or replace function raise_test() returns void as $$
 begin