Support SQL/PSM-compatible CASE statement in plpgsql.

Pavel Stehule

6 files changed, 344 insertions, 18 deletions

diff --git a/src/pl/plpgsql/src/gram.y b/src/pl/plpgsql/src/gram.y
index b67a8bba776..dfd37f67717 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.112 2008/05/13 22:10:29 tgl Exp $
+ *	  $PostgreSQL: pgsql/src/pl/plpgsql/src/gram.y,v 1.113 2008/05/15 22:39:49 tgl Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -37,6 +37,8 @@ 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  PLpgSQL_stmt 	*make_case(int lineno, PLpgSQL_expr *t_expr,
+								   List *case_when_list, List *else_stmts);
 static	void			 check_assignable(PLpgSQL_datum *datum);
 static	void			 read_into_target(PLpgSQL_rec **rec, PLpgSQL_row **row,
 										  bool *strict);
@@ -102,6 +104,7 @@ static List				*read_raise_options(void);
 		PLpgSQL_nsitem			*nsitem;
 		PLpgSQL_diag_item		*diagitem;
 		PLpgSQL_stmt_fetch		*fetch;
+		PLpgSQL_case_when		*casewhen;
 }
 
 %type <declhdr> decl_sect
@@ -116,7 +119,7 @@ static List				*read_raise_options(void);
 %type <str>		decl_stmts decl_stmt
 
 %type <expr>	expr_until_semi expr_until_rightbracket
-%type <expr>	expr_until_then expr_until_loop
+%type <expr>	expr_until_then expr_until_loop opt_expr_until_when
 %type <expr>	opt_exitcond
 
 %type <ival>	assign_var
@@ -135,12 +138,16 @@ static List				*read_raise_options(void);
 %type <stmt>	stmt_return stmt_raise stmt_execsql stmt_execsql_insert
 %type <stmt>	stmt_dynexecute stmt_for stmt_perform stmt_getdiag
 %type <stmt>	stmt_open stmt_fetch stmt_move stmt_close stmt_null
+%type <stmt>	stmt_case
 
 %type <list>	proc_exceptions
 %type <exception_block> exception_sect
 %type <exception>	proc_exception
 %type <condition>	proc_conditions proc_condition
 
+%type <casewhen>	case_when
+%type <list>	case_when_list opt_case_else
+
 %type <list>	getdiag_list
 %type <diagitem> getdiag_list_item
 %type <ival>	getdiag_kind getdiag_target
@@ -157,6 +164,7 @@ static List				*read_raise_options(void);
 %token	K_ASSIGN
 %token	K_BEGIN
 %token	K_BY
+%token	K_CASE
 %token	K_CLOSE
 %token	K_CONSTANT
 %token	K_CONTINUE
@@ -581,9 +589,7 @@ decl_defkey		: K_ASSIGN
 				;
 
 proc_sect		:
-					{
-						$$ = NIL;
-					}
+					{ $$ = NIL; }
 				| proc_stmts
 					{ $$ = $1; }
 				;
@@ -598,7 +604,7 @@ proc_stmts		: proc_stmts proc_stmt
 				| proc_stmt
 						{
 							if ($1 == NULL)
-								$$ = NULL;
+								$$ = NIL;
 							else
 								$$ = list_make1($1);
 						}
@@ -610,6 +616,8 @@ proc_stmt		: pl_block ';'
 						{ $$ = $1; }
 				| stmt_if
 						{ $$ = $1; }
+				| stmt_case
+						{ $$ = $1; }
 				| stmt_loop
 						{ $$ = $1; }
 				| stmt_while
@@ -808,6 +816,67 @@ stmt_else		:
 					}
 				;
 
+stmt_case		: K_CASE lno opt_expr_until_when case_when_list opt_case_else K_END K_CASE ';'
+					{
+						$$ = make_case($2, $3, $4, $5);
+					}
+				;
+
+opt_expr_until_when	:
+					{
+						PLpgSQL_expr *expr = NULL;
+						int	tok = yylex();
+
+						if (tok != K_WHEN)
+						{
+							plpgsql_push_back_token(tok);
+							expr = plpgsql_read_expression(K_WHEN, "WHEN");
+						}
+						plpgsql_push_back_token(K_WHEN);
+						$$ = expr;
+					}
+			    ;
+
+case_when_list	: case_when_list case_when
+					{
+						$$ = lappend($1, $2);
+					}
+				| case_when
+					{
+						$$ = list_make1($1);
+					}
+				;
+
+case_when		: K_WHEN lno expr_until_then proc_sect
+					{
+						PLpgSQL_case_when *new = palloc(sizeof(PLpgSQL_case_when));
+
+						new->lineno	= $2;
+						new->expr	= $3;
+						new->stmts	= $4;
+						$$ = new;
+					}
+				;
+
+opt_case_else	:
+					{
+						$$ = NIL;
+					}
+				| K_ELSE proc_sect
+					{
+						/*
+						 * proc_sect could return an empty list, but we
+						 * must distinguish that from not having ELSE at all.
+						 * Simplest fix is to return a list with one NULL
+						 * pointer, which make_case() must take care of.
+						 */
+						if ($2 != NIL)
+							$$ = $2;
+						else
+							$$ = list_make1(NULL);
+					}
+				;
+
 stmt_loop		: opt_block_label K_LOOP lno loop_body
 					{
 						PLpgSQL_stmt_loop *new;
@@ -2804,6 +2873,103 @@ read_raise_options(void)
 	return result;
 }
 
+/*
+ * Fix up CASE statement
+ */
+static PLpgSQL_stmt *
+make_case(int lineno, PLpgSQL_expr *t_expr,
+		  List *case_when_list, List *else_stmts)
+{
+	PLpgSQL_stmt_case 	*new;
+
+	new = palloc(sizeof(PLpgSQL_stmt_case));
+	new->cmd_type = PLPGSQL_STMT_CASE;
+	new->lineno = lineno;
+	new->t_expr = t_expr;
+	new->t_varno = 0;
+	new->case_when_list = case_when_list;
+	new->have_else = (else_stmts != NIL);
+	/* Get rid of list-with-NULL hack */
+	if (list_length(else_stmts) == 1 && linitial(else_stmts) == NULL)
+		new->else_stmts = NIL;
+	else
+		new->else_stmts = else_stmts;
+
+	/*
+	 * When test expression is present, we create a var for it and then
+	 * convert all the WHEN expressions to "VAR IN (original_expression)".
+	 * This is a bit klugy, but okay since we haven't yet done more than
+	 * read the expressions as text.  (Note that previous parsing won't
+	 * have complained if the WHEN ... THEN expression contained multiple
+	 * comma-separated values.)
+	 */
+	if (t_expr)
+	{
+		ListCell *l;
+		PLpgSQL_var *t_var;
+		int		t_varno;
+
+		/*
+		 * We don't yet know the result datatype of t_expr.  Build the
+		 * variable as if it were INT4; we'll fix this at runtime if needed.
+		 */
+		t_var = (PLpgSQL_var *)
+			plpgsql_build_variable("*case*", lineno,
+								   plpgsql_build_datatype(INT4OID, -1),
+								   false);
+		t_varno = t_var->varno;
+		new->t_varno = t_varno;
+
+		foreach(l, case_when_list)
+		{
+			PLpgSQL_case_when *cwt = (PLpgSQL_case_when *) lfirst(l);
+			PLpgSQL_expr *expr = cwt->expr;
+			int		nparams = expr->nparams;
+			PLpgSQL_expr *new_expr;
+			PLpgSQL_dstring ds;
+			char	buff[32];
+
+			/* Must add the CASE variable as an extra param to expression */
+			if (nparams >= MAX_EXPR_PARAMS)
+			{
+				plpgsql_error_lineno = cwt->lineno;
+				ereport(ERROR,
+					    (errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
+					     errmsg("too many variables specified in SQL statement")));
+			}
+
+			new_expr = palloc(sizeof(PLpgSQL_expr) + sizeof(int) * (nparams + 1) - sizeof(int));
+			memcpy(new_expr, expr,
+				   sizeof(PLpgSQL_expr) + sizeof(int) * nparams - sizeof(int));
+			new_expr->nparams = nparams + 1;
+			new_expr->params[nparams] = t_varno;
+
+			/* And do the string hacking */
+			plpgsql_dstring_init(&ds);
+
+			plpgsql_dstring_append(&ds, "SELECT $");
+			snprintf(buff, sizeof(buff), "%d", nparams + 1);
+			plpgsql_dstring_append(&ds, buff);
+			plpgsql_dstring_append(&ds, " IN (");
+
+			/* copy expression query without SELECT keyword */
+			Assert(strncmp(expr->query, "SELECT ", 7) == 0);
+			plpgsql_dstring_append(&ds, expr->query + 7);
+			plpgsql_dstring_append_char(&ds, ')');
+
+			new_expr->query = pstrdup(plpgsql_dstring_get(&ds));
+
+			plpgsql_dstring_free(&ds);
+			pfree(expr->query);
+			pfree(expr);
+
+			cwt->expr = new_expr;
+		}
+	}
+
+	return (PLpgSQL_stmt *) new;
+}
+
 
 /* Needed to avoid conflict between different prefix settings: */
 #undef yylex
diff --git a/src/pl/plpgsql/src/pl_exec.c b/src/pl/plpgsql/src/pl_exec.c
index 2ba45befb71..aeb5d365b49 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.214 2008/05/13 22:10:30 tgl Exp $
+ *	  $PostgreSQL: pgsql/src/pl/plpgsql/src/pl_exec.c,v 1.215 2008/05/15 22:39:49 tgl Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -94,6 +94,8 @@ static int exec_stmt_getdiag(PLpgSQL_execstate *estate,
 				  PLpgSQL_stmt_getdiag *stmt);
 static int exec_stmt_if(PLpgSQL_execstate *estate,
 			 PLpgSQL_stmt_if *stmt);
+static int exec_stmt_case(PLpgSQL_execstate *estate,
+						  PLpgSQL_stmt_case *stmt);
 static int exec_stmt_loop(PLpgSQL_execstate *estate,
 			   PLpgSQL_stmt_loop *stmt);
 static int exec_stmt_while(PLpgSQL_execstate *estate,
@@ -1229,7 +1231,7 @@ exec_stmt(PLpgSQL_execstate *estate, PLpgSQL_stmt *stmt)
 
 	CHECK_FOR_INTERRUPTS();
 
-	switch (stmt->cmd_type)
+	switch ((enum PLpgSQL_stmt_types) stmt->cmd_type)
 	{
 		case PLPGSQL_STMT_BLOCK:
 			rc = exec_stmt_block(estate, (PLpgSQL_stmt_block *) stmt);
@@ -1251,6 +1253,10 @@ exec_stmt(PLpgSQL_execstate *estate, PLpgSQL_stmt *stmt)
 			rc = exec_stmt_if(estate, (PLpgSQL_stmt_if *) stmt);
 			break;
 
+		case PLPGSQL_STMT_CASE:
+			rc = exec_stmt_case(estate, (PLpgSQL_stmt_case *) stmt);
+			break;
+
 		case PLPGSQL_STMT_LOOP:
 			rc = exec_stmt_loop(estate, (PLpgSQL_stmt_loop *) stmt);
 			break;
@@ -1442,6 +1448,91 @@ exec_stmt_if(PLpgSQL_execstate *estate, PLpgSQL_stmt_if *stmt)
 }
 
 
+/*-----------
+ * exec_stmt_case
+ *-----------
+ */
+static int
+exec_stmt_case(PLpgSQL_execstate *estate, PLpgSQL_stmt_case *stmt)
+{
+	PLpgSQL_var *t_var = NULL;
+	bool		isnull;
+	ListCell   *l;
+
+	if (stmt->t_expr != NULL)
+	{
+		/* simple case */
+		Datum	t_val;
+		Oid		t_oid;
+
+		t_val = exec_eval_expr(estate, stmt->t_expr, &isnull, &t_oid);
+
+		t_var = (PLpgSQL_var *) estate->datums[stmt->t_varno];
+
+		/*
+		 * When expected datatype is different from real, change it.
+		 * Note that what we're modifying here is an execution copy
+		 * of the datum, so this doesn't affect the originally stored
+		 * function parse tree.
+		 */
+		if (t_var->datatype->typoid != t_oid)
+			t_var->datatype = plpgsql_build_datatype(t_oid, -1);
+
+		/* now we can assign to the variable */
+		exec_assign_value(estate,
+						  (PLpgSQL_datum *) t_var,
+						  t_val,
+						  t_oid,
+						  &isnull);
+
+		exec_eval_cleanup(estate);
+	}
+
+	/* Now search for a successful WHEN clause */
+	foreach(l, stmt->case_when_list)
+	{
+		PLpgSQL_case_when *cwt = (PLpgSQL_case_when *) lfirst(l);
+		bool	value;
+
+		value = exec_eval_boolean(estate, cwt->expr, &isnull);
+		exec_eval_cleanup(estate);
+		if (!isnull && value)
+		{
+			/* Found it */
+
+			/* We can now discard any value we had for the temp variable */
+			if (t_var != NULL)
+			{
+				free_var(t_var);
+				t_var->value = (Datum) 0;
+				t_var->isnull = true;
+			}
+
+			/* Evaluate the statement(s), and we're done */
+			return exec_stmts(estate, cwt->stmts);
+		}
+	}
+
+	/* We can now discard any value we had for the temp variable */
+	if (t_var != NULL)
+	{
+		free_var(t_var);
+		t_var->value = (Datum) 0;
+		t_var->isnull = true;
+	}
+
+	/* SQL2003 mandates this error if there was no ELSE clause */
+	if (!stmt->have_else)
+		ereport(ERROR,
+				(errcode(ERRCODE_CASE_NOT_FOUND),
+				 errmsg("case not found"),
+				 errhint("CASE statement is missing ELSE part.")));
+
+	/* Evaluate the ELSE statements, and we're done */
+	return exec_stmts(estate, stmt->else_stmts);
+}
+
+
 /* ----------
  * exec_stmt_loop			Loop over statements until
  *					an exit occurs.
diff --git a/src/pl/plpgsql/src/pl_funcs.c b/src/pl/plpgsql/src/pl_funcs.c
index cb9e9c99ec3..aab349bdb2c 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.71 2008/05/13 22:10:30 tgl Exp $
+ *	  $PostgreSQL: pgsql/src/pl/plpgsql/src/pl_funcs.c,v 1.72 2008/05/15 22:39:49 tgl Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -466,7 +466,7 @@ plpgsql_convert_ident(const char *s, char **output, int numidents)
 const char *
 plpgsql_stmt_typename(PLpgSQL_stmt *stmt)
 {
-	switch (stmt->cmd_type)
+	switch ((enum PLpgSQL_stmt_types) stmt->cmd_type)
 	{
 		case PLPGSQL_STMT_BLOCK:
 			return _("statement block");
@@ -474,6 +474,8 @@ plpgsql_stmt_typename(PLpgSQL_stmt *stmt)
 			return _("assignment");
 		case PLPGSQL_STMT_IF:
 			return "IF";
+		case PLPGSQL_STMT_CASE:
+			return "CASE";
 		case PLPGSQL_STMT_LOOP:
 			return "LOOP";
 		case PLPGSQL_STMT_WHILE:
@@ -526,6 +528,7 @@ static void dump_stmt(PLpgSQL_stmt *stmt);
 static void dump_block(PLpgSQL_stmt_block *block);
 static void dump_assign(PLpgSQL_stmt_assign *stmt);
 static void dump_if(PLpgSQL_stmt_if *stmt);
+static void dump_case(PLpgSQL_stmt_case *stmt);
 static void dump_loop(PLpgSQL_stmt_loop *stmt);
 static void dump_while(PLpgSQL_stmt_while *stmt);
 static void dump_fori(PLpgSQL_stmt_fori *stmt);
@@ -561,7 +564,7 @@ static void
 dump_stmt(PLpgSQL_stmt *stmt)
 {
 	printf("%3d:", stmt->lineno);
-	switch (stmt->cmd_type)
+	switch ((enum PLpgSQL_stmt_types) stmt->cmd_type)
 	{
 		case PLPGSQL_STMT_BLOCK:
 			dump_block((PLpgSQL_stmt_block *) stmt);
@@ -572,6 +575,9 @@ dump_stmt(PLpgSQL_stmt *stmt)
 		case PLPGSQL_STMT_IF:
 			dump_if((PLpgSQL_stmt_if *) stmt);
 			break;
+		case PLPGSQL_STMT_CASE:
+			dump_case((PLpgSQL_stmt_case *) stmt);
+			break;
 		case PLPGSQL_STMT_LOOP:
 			dump_loop((PLpgSQL_stmt_loop *) stmt);
 			break;
@@ -715,6 +721,44 @@ dump_if(PLpgSQL_stmt_if *stmt)
 }
 
 static void
+dump_case(PLpgSQL_stmt_case *stmt)
+{
+	ListCell	*l;
+
+	dump_ind();
+	printf("CASE %d ", stmt->t_varno);
+	if (stmt->t_expr)
+		dump_expr(stmt->t_expr);
+	printf("\n");
+	dump_indent += 6;
+	foreach(l, stmt->case_when_list)
+	{
+		PLpgSQL_case_when *cwt = (PLpgSQL_case_when *) lfirst(l);
+
+		dump_ind();
+		printf("WHEN ");
+		dump_expr(cwt->expr);
+		printf("\n");
+		dump_ind();
+		printf("THEN\n");
+		dump_indent += 2;
+		dump_stmts(cwt->stmts);
+		dump_indent -= 2;
+	}
+	if (stmt->have_else)
+	{
+		dump_ind();
+		printf("ELSE\n");
+		dump_indent += 2;
+		dump_stmts(stmt->else_stmts);
+		dump_indent -= 2;
+	}
+	dump_indent -= 6;
+	dump_ind();
+	printf("    ENDCASE\n");
+}
+
+static void
 dump_loop(PLpgSQL_stmt_loop *stmt)
 {
 	dump_ind();
@@ -1025,7 +1069,7 @@ dump_raise(PLpgSQL_stmt_raise *stmt)
 		foreach(lc, stmt->options)
 		{
 			PLpgSQL_raise_option *opt = (PLpgSQL_raise_option *) lfirst(lc);
-		
+
 			dump_ind();
 			switch (opt->opt_type)
 			{
@@ -1034,7 +1078,7 @@ dump_raise(PLpgSQL_stmt_raise *stmt)
 					break;
 				case PLPGSQL_RAISEOPTION_MESSAGE:
 					printf("    MESSAGE = ");
-					break;	
+					break;
 				case PLPGSQL_RAISEOPTION_DETAIL:
 					printf("    DETAIL = ");
 					break;
@@ -1044,7 +1088,7 @@ dump_raise(PLpgSQL_stmt_raise *stmt)
 			}
 			dump_expr(opt->expr);
 			printf("\n");
-		}		
+		}
 		dump_indent -= 2;
 	}
 	dump_indent -= 2;
diff --git a/src/pl/plpgsql/src/plerrcodes.h b/src/pl/plpgsql/src/plerrcodes.h
index 5fbdf255d82..e6ab24dfd0e 100644
--- a/src/pl/plpgsql/src/plerrcodes.h
+++ b/src/pl/plpgsql/src/plerrcodes.h
@@ -9,7 +9,7 @@
  *
  * Copyright (c) 2003-2008, PostgreSQL Global Development Group
  *
- * $PostgreSQL: pgsql/src/pl/plpgsql/src/plerrcodes.h,v 1.13 2008/01/15 01:36:53 tgl Exp $
+ * $PostgreSQL: pgsql/src/pl/plpgsql/src/plerrcodes.h,v 1.14 2008/05/15 22:39:49 tgl Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -80,6 +80,10 @@
 },
 
 {
+	"case_not_found", ERRCODE_CASE_NOT_FOUND
+},
+
+{
 	"cardinality_violation", ERRCODE_CARDINALITY_VIOLATION
 },
 
diff --git a/src/pl/plpgsql/src/plpgsql.h b/src/pl/plpgsql/src/plpgsql.h
index 53d691a596a..735a6810053 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.99 2008/05/13 22:10:30 tgl Exp $
+ *	  $PostgreSQL: pgsql/src/pl/plpgsql/src/plpgsql.h,v 1.100 2008/05/15 22:39:49 tgl Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -71,11 +71,12 @@ enum
  * Execution tree node types
  * ----------
  */
-enum
+enum PLpgSQL_stmt_types
 {
 	PLPGSQL_STMT_BLOCK,
 	PLPGSQL_STMT_ASSIGN,
 	PLPGSQL_STMT_IF,
+	PLPGSQL_STMT_CASE,
 	PLPGSQL_STMT_LOOP,
 	PLPGSQL_STMT_WHILE,
 	PLPGSQL_STMT_FORI,
@@ -390,6 +391,25 @@ typedef struct
 } PLpgSQL_stmt_if;
 
 
+typedef struct					/* CASE statement */
+{
+	int			cmd_type;
+	int			lineno;
+	PLpgSQL_expr *t_expr;		/* test expression, or NULL if none */
+	int			t_varno;		/* var to store test expression value into */
+	List	   *case_when_list;	/* List of PLpgSQL_case_when structs */
+	bool		have_else;		/* flag needed because list could be empty */
+	List	   *else_stmts;		/* List of statements */
+} PLpgSQL_stmt_case;
+
+typedef struct					/* one arm of CASE statement */
+{
+	int			lineno;
+	PLpgSQL_expr *expr;			/* boolean expression for this case */
+	List	   *stmts;			/* List of statements */
+} PLpgSQL_case_when;
+
+
 typedef struct
 {								/* Unconditional LOOP statement		*/
 	int			cmd_type;
diff --git a/src/pl/plpgsql/src/scan.l b/src/pl/plpgsql/src/scan.l
index 73258ec3648..77e9335a052 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.63 2008/05/13 22:10:30 tgl Exp $
+ *	  $PostgreSQL: pgsql/src/pl/plpgsql/src/scan.l,v 1.64 2008/05/15 22:39:49 tgl Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -116,6 +116,7 @@ dolqinside		[^$]+
 alias			{ return K_ALIAS;			}
 begin			{ return K_BEGIN;			}
 by				{ return K_BY;   			}
+case			{ return K_CASE;			}
 close			{ return K_CLOSE;			}
 constant		{ return K_CONSTANT;		}
 continue		{ return K_CONTINUE;		}