plpgsql.extra_warnings='num_into_expressions'

*** a/doc/src/sgml/plpgsql.sgml
--- b/doc/src/sgml/plpgsql.sgml
***************
*** 4719,4725 **** a_output := a_output || $$ if v_$$ || referrer_keys.kind || $$ like '$$
    <varname>plpgsql.extra_errors</> for errors. Both can be set either to
    a comma-separated list of checks, <literal>"none"</> or <literal>"all"</>.
    The default is <literal>"none"</>. Currently the list of available checks
!   includes only one:
    <variablelist>
     <varlistentry>
      <term><varname>shadowed_variables</varname></term>
--- 4719,4725 ----
    <varname>plpgsql.extra_errors</> for errors. Both can be set either to
    a comma-separated list of checks, <literal>"none"</> or <literal>"all"</>.
    The default is <literal>"none"</>. Currently the list of available checks
!   is as follows:
    <variablelist>
     <varlistentry>
      <term><varname>shadowed_variables</varname></term>
***************
*** 4729,4734 **** a_output := a_output || $$ if v_$$ || referrer_keys.kind || $$ like '$$
--- 4729,4744 ----
       </para>
      </listitem>
     </varlistentry>
+    <varlistentry>
+     <term><varname>num_into_expressions</varname></term>
+     <listitem>
+      <para>
+       When possible, checks that the number of expressions specified in the
+       SELECT or RETURNING list of a query matches the number expected by the
+       INTO target.  This check is done on a "best effort" basis.
+      </para>
+     </listitem>
+    </varlistentry>
    </variablelist>
  
    The following example shows the effect of <varname>plpgsql.extra_warnings</>
*** a/src/pl/plpgsql/src/pl_gram.y
--- b/src/pl/plpgsql/src/pl_gram.y
***************
*** 22,27 ****
--- 22,28 ----
  #include "parser/scanner.h"
  #include "parser/scansup.h"
  #include "utils/builtins.h"
+ #include "nodes/nodefuncs.h"
  
  
  /* Location tracking support --- simpler than bison's default */
***************
*** 97,103 **** static	PLpgSQL_row		*make_scalar_list1(char *initial_name,
  										   PLpgSQL_datum *initial_datum,
  										   int lineno, int location);
  static	void			 check_sql_expr(const char *stmt, int location,
! 										int leaderlen);
  static	void			 plpgsql_sql_error_callback(void *arg);
  static	PLpgSQL_type	*parse_datatype(const char *string, int location);
  static	void			 check_labels(const char *start_label,
--- 98,104 ----
  										   PLpgSQL_datum *initial_datum,
  										   int lineno, int location);
  static	void			 check_sql_expr(const char *stmt, int location,
! 										int leaderlen, PLpgSQL_row *check_row);
  static	void			 plpgsql_sql_error_callback(void *arg);
  static	PLpgSQL_type	*parse_datatype(const char *string, int location);
  static	void			 check_labels(const char *start_label,
***************
*** 106,111 **** static	void			 check_labels(const char *start_label,
--- 107,115 ----
  static	PLpgSQL_expr	*read_cursor_args(PLpgSQL_var *cursor,
  										  int until, const char *expected);
  static	List			*read_raise_options(void);
+ static	bool			find_a_star_walker(Node *node, void *context);
+ static	int				tlist_result_column_count(Node *stmt);
+ 
  
  %}
  
***************
*** 1438,1444 **** for_control		: for_variable K_IN
  								PLpgSQL_stmt_fori	*new;
  
  								/* Check first expression is well-formed */
! 								check_sql_expr(expr1->query, expr1loc, 7);
  
  								/* Read and check the second one */
  								expr2 = read_sql_expression2(K_LOOP, K_BY,
--- 1442,1448 ----
  								PLpgSQL_stmt_fori	*new;
  
  								/* Check first expression is well-formed */
! 								check_sql_expr(expr1->query, expr1loc, 7, NULL);
  
  								/* Read and check the second one */
  								expr2 = read_sql_expression2(K_LOOP, K_BY,
***************
*** 1500,1506 **** for_control		: for_variable K_IN
  								pfree(expr1->query);
  								expr1->query = tmp_query;
  
! 								check_sql_expr(expr1->query, expr1loc, 0);
  
  								new = palloc0(sizeof(PLpgSQL_stmt_fors));
  								new->cmd_type = PLPGSQL_STMT_FORS;
--- 1504,1510 ----
  								pfree(expr1->query);
  								expr1->query = tmp_query;
  
! 								check_sql_expr(expr1->query, expr1loc, 0, NULL);
  
  								new = palloc0(sizeof(PLpgSQL_stmt_fors));
  								new->cmd_type = PLPGSQL_STMT_FORS;
***************
*** 2592,2598 **** read_sql_construct(int until,
  	pfree(ds.data);
  
  	if (valid_sql)
! 		check_sql_expr(expr->query, startlocation, strlen(sqlstart));
  
  	return expr;
  }
--- 2596,2602 ----
  	pfree(ds.data);
  
  	if (valid_sql)
! 		check_sql_expr(expr->query, startlocation, strlen(sqlstart), NULL);
  
  	return expr;
  }
***************
*** 2815,2821 **** make_execsql_stmt(int firsttoken, int location)
  	expr->ns			= plpgsql_ns_top();
  	pfree(ds.data);
  
! 	check_sql_expr(expr->query, location, 0);
  
  	execsql = palloc(sizeof(PLpgSQL_stmt_execsql));
  	execsql->cmd_type = PLPGSQL_STMT_EXECSQL;
--- 2819,2825 ----
  	expr->ns			= plpgsql_ns_top();
  	pfree(ds.data);
  
! 	check_sql_expr(expr->query, location, 0, row);
  
  	execsql = palloc(sizeof(PLpgSQL_stmt_execsql));
  	execsql->cmd_type = PLPGSQL_STMT_EXECSQL;
***************
*** 3409,3421 **** make_scalar_list1(char *initial_name,
   * (typically "SELECT ") prefixed to the source text.  We use this assumption
   * to transpose any error cursor position back to the function source text.
   * If no error cursor is provided, we'll just point at "location".
   */
  static void
! check_sql_expr(const char *stmt, int location, int leaderlen)
  {
  	sql_error_callback_arg cbarg;
  	ErrorContextCallback  syntax_errcontext;
  	MemoryContext oldCxt;
  
  	if (!plpgsql_check_syntax)
  		return;
--- 3413,3430 ----
   * (typically "SELECT ") prefixed to the source text.  We use this assumption
   * to transpose any error cursor position back to the function source text.
   * If no error cursor is provided, we'll just point at "location".
+  *
+  * If check_row is not NULL and PLPGSQL_XCHECK_NUM_INTO_EXPRESSIONS is enabled,
+  * the statement's result column count is compared against it when the exact
+  * number of columns is known and an error or a warning is reported.
   */
  static void
! check_sql_expr(const char *stmt, int location, int leaderlen, PLpgSQL_row *check_row)
  {
  	sql_error_callback_arg cbarg;
  	ErrorContextCallback  syntax_errcontext;
  	MemoryContext oldCxt;
+ 	List *raw_parsetree_list;
  
  	if (!plpgsql_check_syntax)
  		return;
***************
*** 3429,3441 **** check_sql_expr(const char *stmt, int location, int leaderlen)
  	error_context_stack = &syntax_errcontext;
  
  	oldCxt = MemoryContextSwitchTo(compile_tmp_cxt);
! 	(void) raw_parser(stmt);
  	MemoryContextSwitchTo(oldCxt);
  
  	/* Restore former ereport callback */
  	error_context_stack = syntax_errcontext.previous;
  }
  
  static void
  plpgsql_sql_error_callback(void *arg)
  {
--- 3438,3535 ----
  	error_context_stack = &syntax_errcontext;
  
  	oldCxt = MemoryContextSwitchTo(compile_tmp_cxt);
! 	raw_parsetree_list = raw_parser(stmt);
! 
! 	if (check_row != NULL &&
! 		(plpgsql_curr_compile->extra_warnings & PLPGSQL_XCHECK_NUM_INTO_EXPRESSIONS ||
! 		 plpgsql_curr_compile->extra_errors & PLPGSQL_XCHECK_NUM_INTO_EXPRESSIONS))
! 	{
! 		Node *raw_parse_tree;
! 		int ncols;
! 		int fnum;
! 		int expected_ncols = 0;
! 
! 		for (fnum = 0; fnum < check_row->nfields; fnum++)
! 		{
! 			if (check_row->varnos[fnum] < 0)
! 				continue;
! 			expected_ncols++;
! 		}
! 
! 		raw_parse_tree = linitial(raw_parsetree_list);
! 		ncols = tlist_result_column_count(raw_parse_tree);
! 		if (ncols >= 0 && ncols != expected_ncols)
! 			ereport(plpgsql_curr_compile->extra_errors & PLPGSQL_XCHECK_NUM_INTO_EXPRESSIONS ? ERROR : WARNING,
! 				(errcode(ERRCODE_SYNTAX_ERROR),
! 				errmsg("the number of result expressions does not match that expected by the INTO target"),
! 				errdetail("INTO target expects %d expressions, but the query returns %d", expected_ncols, ncols)));
! 	}
  	MemoryContextSwitchTo(oldCxt);
  
  	/* Restore former ereport callback */
  	error_context_stack = syntax_errcontext.previous;
  }
  
+ /*
+  * Expression tree walker for tlist_result_column_count.  Returns true if the
+  * expression tree contains an A_Star node, false otherwise.
+  */
+ static bool
+ find_a_star_walker(Node *node, void *context)
+ {
+ 	if (node == NULL)
+ 		return false;
+ 	if (IsA(node, A_Star))
+ 		return true;
+ 	if (IsA(node, ColumnRef))
+ 	{
+ 		ColumnRef *ref = (ColumnRef *) node;
+ 		/* A_Star can only be the last element */
+ 		if (IsA(llast(ref->fields), A_Star))
+ 			return true;
+ 	}
+ 	return raw_expression_tree_walker((Node *) node,
+ 									  find_a_star_walker,
+ 									  context);
+ }
+ 
+ /*
+  * Find the number of columns in a raw statement's target list (if SELECT) or
+  * returningList (if INSERT, UPDATE or DELETE).  Returns -1 if the number of
+  * columns could not be determined because of an A_Star.
+  */
+ static int
+ tlist_result_column_count(Node *stmt)
+ {
+ 	List *tlist;
+ 
+ 	if (IsA(stmt, SelectStmt))
+ 	{
+ 		SelectStmt *s = (SelectStmt *) stmt;
+ 		if (s->valuesLists)
+ 			tlist = (List *) linitial(s->valuesLists);
+ 		else if (s->op == SETOP_NONE)
+ 			tlist = s->targetList;
+ 		else
+ 			return tlist_result_column_count((Node *) s->larg);
+ 	}
+ 	else if (IsA(stmt, InsertStmt))
+ 		tlist = ((InsertStmt *) stmt)->returningList;
+ 	else if (IsA(stmt, UpdateStmt))
+ 		tlist = ((UpdateStmt *) stmt)->returningList;
+ 	else if (IsA(stmt, DeleteStmt))
+ 		tlist = ((DeleteStmt *) stmt)->returningList;
+ 	else
+ 		elog(ERROR, "unknown nodeTag %d", nodeTag(stmt));
+ 
+ 	if (tlist == NIL)
+ 		return 0;
+ 
+ 	if (raw_expression_tree_walker((Node *) tlist, find_a_star_walker, NULL))
+ 		return -1;
+ 	return list_length(tlist);
+ }
+ 
  static void
  plpgsql_sql_error_callback(void *arg)
  {
*** a/src/pl/plpgsql/src/pl_handler.c
--- b/src/pl/plpgsql/src/pl_handler.c
***************
*** 87,92 **** plpgsql_extra_checks_check_hook(char **newvalue, void **extra, GucSource source)
--- 87,94 ----
  
  			if (pg_strcasecmp(tok, "shadowed_variables") == 0)
  				extrachecks |= PLPGSQL_XCHECK_SHADOWVAR;
+ 			else if (pg_strcasecmp(tok, "num_into_expressions") == 0)
+ 				extrachecks |= PLPGSQL_XCHECK_NUM_INTO_EXPRESSIONS;
  			else if (pg_strcasecmp(tok, "all") == 0 || pg_strcasecmp(tok, "none") == 0)
  			{
  				GUC_check_errdetail("Key word \"%s\" cannot be combined with other key words.", tok);
*** a/src/pl/plpgsql/src/plpgsql.h
--- b/src/pl/plpgsql/src/plpgsql.h
***************
*** 886,894 **** extern int	plpgsql_variable_conflict;
  extern bool plpgsql_print_strict_params;
  
  /* extra compile-time checks */
! #define PLPGSQL_XCHECK_NONE			0
! #define PLPGSQL_XCHECK_SHADOWVAR	1
! #define PLPGSQL_XCHECK_ALL			((int) ~0)
  
  extern int	plpgsql_extra_warnings;
  extern int	plpgsql_extra_errors;
--- 886,895 ----
  extern bool plpgsql_print_strict_params;
  
  /* extra compile-time checks */
! #define PLPGSQL_XCHECK_NONE						0
! #define PLPGSQL_XCHECK_SHADOWVAR				1
! #define PLPGSQL_XCHECK_NUM_INTO_EXPRESSIONS		2
! #define PLPGSQL_XCHECK_ALL						((int) ~0)
  
  extern int	plpgsql_extra_warnings;
  extern int	plpgsql_extra_errors;
*** a/src/test/regress/expected/plpgsql.out
--- b/src/test/regress/expected/plpgsql.out
***************
*** 5340,5342 **** NOTICE:  outer_func() done
--- 5340,5404 ----
  drop function outer_outer_func(int);
  drop function outer_func(int);
  drop function inner_func(int);
+ create temporary table somecols(a int, b int);
+ set plpgsql.extra_warnings to 'num_into_expressions';
+ -- INTO column counts
+ create or replace function into_counts()
+ returns void as $$
+ declare
+   myresult int;
+   myrec record;
+   mycols somecols;
+ begin
+   -- no warning
+   select 1 into myresult;
+   -- we don't know the column count without parse analysis, no warning
+   select * into myresult from somecols;
+   -- records are OK
+   select 1, 2 into myrec;
+   -- warning
+   select a, b into myresult from somecols;
+   -- warning
+   select into myresult from somecols;
+   -- no warning
+   select 1, 2 into mycols;
+   -- warning
+   select 1 into mycols;
+   -- warning
+   select 1, 2, 3 into mycols;
+   -- set operations; OK
+   select 1 into myresult union all select 2;
+   -- set operations; not OK
+   select 1,2 into myresult union all select 2,3;
+   -- values; OK
+   values (1), (2) into myresult;
+   -- values; not OK
+   values (1,2),(2,3) into myresult;
+ end;
+ $$ language plpgsql;
+ WARNING:  the number of result expressions does not match that expected by the INTO target
+ LINE 15:   select a, b into myresult from somecols;
+            ^
+ DETAIL:  INTO target expects 1 expressions, but the query returns 2
+ WARNING:  the number of result expressions does not match that expected by the INTO target
+ LINE 17:   select into myresult from somecols;
+            ^
+ DETAIL:  INTO target expects 1 expressions, but the query returns 0
+ WARNING:  the number of result expressions does not match that expected by the INTO target
+ LINE 21:   select 1 into mycols;
+            ^
+ DETAIL:  INTO target expects 2 expressions, but the query returns 1
+ WARNING:  the number of result expressions does not match that expected by the INTO target
+ LINE 23:   select 1, 2, 3 into mycols;
+            ^
+ DETAIL:  INTO target expects 2 expressions, but the query returns 3
+ WARNING:  the number of result expressions does not match that expected by the INTO target
+ LINE 27:   select 1,2 into myresult union all select 2,3;
+            ^
+ DETAIL:  INTO target expects 1 expressions, but the query returns 2
+ WARNING:  the number of result expressions does not match that expected by the INTO target
+ LINE 31:   values (1,2),(2,3) into myresult;
+            ^
+ DETAIL:  INTO target expects 1 expressions, but the query returns 2
+ drop table somecols;
+ reset plpgsql.extra_warnings;
*** a/src/test/regress/sql/plpgsql.sql
--- b/src/test/regress/sql/plpgsql.sql
***************
*** 4193,4195 **** select outer_outer_func(20);
--- 4193,4238 ----
  drop function outer_outer_func(int);
  drop function outer_func(int);
  drop function inner_func(int);
+ 
+ create temporary table somecols(a int, b int);
+ 
+ set plpgsql.extra_warnings to 'num_into_expressions';
+ 
+ -- INTO column counts
+ create or replace function into_counts()
+ returns void as $$
+ declare
+   myresult int;
+   myrec record;
+   mycols somecols;
+ begin
+   -- no warning
+   select 1 into myresult;
+   -- we don't know the column count without parse analysis, no warning
+   select * into myresult from somecols;
+   -- records are OK
+   select 1, 2 into myrec;
+   -- warning
+   select a, b into myresult from somecols;
+   -- warning
+   select into myresult from somecols;
+   -- no warning
+   select 1, 2 into mycols;
+   -- warning
+   select 1 into mycols;
+   -- warning
+   select 1, 2, 3 into mycols;
+   -- set operations; OK
+   select 1 into myresult union all select 2;
+   -- set operations; not OK
+   select 1,2 into myresult union all select 2,3;
+   -- values; OK
+   values (1), (2) into myresult;
+   -- values; not OK
+   values (1,2),(2,3) into myresult;
+ end;
+ $$ language plpgsql;
+ 
+ drop table somecols;
+ 
+ reset plpgsql.extra_warnings;