Add an ASSERT statement in plpgsql.

This is meant to make it easier to insert simple debugging cross-checks
in plpgsql functions.

Pavel Stehule, reviewed by Jim Nasby

src/pl/plpgsql/src/pl_exec.c

@@ -153,6 +153,8 @@ 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_assert(PLpgSQL_execstate *estate,
+				 PLpgSQL_stmt_assert *stmt);
 static int exec_stmt_execsql(PLpgSQL_execstate *estate,
 				  PLpgSQL_stmt_execsql *stmt);
 static int exec_stmt_dynexecute(PLpgSQL_execstate *estate,
@@ -363,8 +365,8 @@ plpgsql_exec_function(PLpgSQL_function *func, FunctionCallInfo fcinfo,
 		estate.err_text = NULL;
 
 		/*
-		 * Provide a more helpful message if a CONTINUE or RAISE has been used
-		 * outside the context it can work in.
+		 * Provide a more helpful message if a CONTINUE has been used outside
+		 * the context it can work in.
 		 */
 		if (rc == PLPGSQL_RC_CONTINUE)
 			ereport(ERROR,
@@ -730,8 +732,8 @@ plpgsql_exec_trigger(PLpgSQL_function *func,
 		estate.err_text = NULL;
 
 		/*
-		 * Provide a more helpful message if a CONTINUE or RAISE has been used
-		 * outside the context it can work in.
+		 * Provide a more helpful message if a CONTINUE has been used outside
+		 * the context it can work in.
 		 */
 		if (rc == PLPGSQL_RC_CONTINUE)
 			ereport(ERROR,
@@ -862,8 +864,8 @@ plpgsql_exec_event_trigger(PLpgSQL_function *func, EventTriggerData *trigdata)
 		estate.err_text = NULL;
 
 		/*
-		 * Provide a more helpful message if a CONTINUE or RAISE has been used
-		 * outside the context it can work in.
+		 * Provide a more helpful message if a CONTINUE has been used outside
+		 * the context it can work in.
 		 */
 		if (rc == PLPGSQL_RC_CONTINUE)
 			ereport(ERROR,
@@ -1027,12 +1029,14 @@ exception_matches_conditions(ErrorData *edata, PLpgSQL_condition *cond)
 		int			sqlerrstate = cond->sqlerrstate;
 
 		/*
-		 * OTHERS matches everything *except* query-canceled; if you're
-		 * foolish enough, you can match that explicitly.
+		 * OTHERS matches everything *except* query-canceled and
+		 * assert-failure.  If you're foolish enough, you can match those
+		 * explicitly.
 		 */
 		if (sqlerrstate == 0)
 		{
-			if (edata->sqlerrcode != ERRCODE_QUERY_CANCELED)
+			if (edata->sqlerrcode != ERRCODE_QUERY_CANCELED &&
+				edata->sqlerrcode != ERRCODE_ASSERT_FAILURE)
 				return true;
 		}
 		/* Exact match? */
@@ -1471,6 +1475,10 @@ exec_stmt(PLpgSQL_execstate *estate, PLpgSQL_stmt *stmt)
 			rc = exec_stmt_raise(estate, (PLpgSQL_stmt_raise *) stmt);
 			break;
 
+		case PLPGSQL_STMT_ASSERT:
+			rc = exec_stmt_assert(estate, (PLpgSQL_stmt_assert *) stmt);
+			break;
+
 		case PLPGSQL_STMT_EXECSQL:
 			rc = exec_stmt_execsql(estate, (PLpgSQL_stmt_execsql *) stmt);
 			break;
@@ -3117,6 +3125,48 @@ exec_stmt_raise(PLpgSQL_execstate *estate, PLpgSQL_stmt_raise *stmt)
 	return PLPGSQL_RC_OK;
 }
 
+/* ----------
+ * exec_stmt_assert			Assert statement
+ * ----------
+ */
+static int
+exec_stmt_assert(PLpgSQL_execstate *estate, PLpgSQL_stmt_assert *stmt)
+{
+	bool		value;
+	bool		isnull;
+
+	/* do nothing when asserts are not enabled */
+	if (!plpgsql_check_asserts)
+		return PLPGSQL_RC_OK;
+
+	value = exec_eval_boolean(estate, stmt->cond, &isnull);
+	exec_eval_cleanup(estate);
+
+	if (isnull || !value)
+	{
+		char	   *message = NULL;
+
+		if (stmt->message != NULL)
+		{
+			Datum		val;
+			Oid			typeid;
+			int32		typmod;
+
+			val = exec_eval_expr(estate, stmt->message,
+								 &isnull, &typeid, &typmod);
+			if (!isnull)
+				message = convert_value_to_string(estate, val, typeid);
+			/* we mustn't do exec_eval_cleanup here */
+		}
+
+		ereport(ERROR,
+				(errcode(ERRCODE_ASSERT_FAILURE),
+				 message ? errmsg_internal("%s", message) :
+				 errmsg("assertion failed")));
+	}
+
+	return PLPGSQL_RC_OK;
+}
 
 /* ----------
  * Initialize a mostly empty execution state

src/pl/plpgsql/src/pl_funcs.c

@@ -244,6 +244,8 @@ plpgsql_stmt_typename(PLpgSQL_stmt *stmt)
 			return "RETURN QUERY";
 		case PLPGSQL_STMT_RAISE:
 			return "RAISE";
+		case PLPGSQL_STMT_ASSERT:
+			return "ASSERT";
 		case PLPGSQL_STMT_EXECSQL:
 			return _("SQL statement");
 		case PLPGSQL_STMT_DYNEXECUTE:
@@ -330,6 +332,7 @@ 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_assert(PLpgSQL_stmt_assert *stmt);
 static void free_execsql(PLpgSQL_stmt_execsql *stmt);
 static void free_dynexecute(PLpgSQL_stmt_dynexecute *stmt);
 static void free_dynfors(PLpgSQL_stmt_dynfors *stmt);
@@ -391,6 +394,9 @@ free_stmt(PLpgSQL_stmt *stmt)
 		case PLPGSQL_STMT_RAISE:
 			free_raise((PLpgSQL_stmt_raise *) stmt);
 			break;
+		case PLPGSQL_STMT_ASSERT:
+			free_assert((PLpgSQL_stmt_assert *) stmt);
+			break;
 		case PLPGSQL_STMT_EXECSQL:
 			free_execsql((PLpgSQL_stmt_execsql *) stmt);
 			break;
@@ -610,6 +616,13 @@ free_raise(PLpgSQL_stmt_raise *stmt)
 	}
 }
 
+static void
+free_assert(PLpgSQL_stmt_assert *stmt)
+{
+	free_expr(stmt->cond);
+	free_expr(stmt->message);
+}
+
 static void
 free_execsql(PLpgSQL_stmt_execsql *stmt)
 {
@@ -732,6 +745,7 @@ 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_assert(PLpgSQL_stmt_assert *stmt);
 static void dump_execsql(PLpgSQL_stmt_execsql *stmt);
 static void dump_dynexecute(PLpgSQL_stmt_dynexecute *stmt);
 static void dump_dynfors(PLpgSQL_stmt_dynfors *stmt);
@@ -804,6 +818,9 @@ dump_stmt(PLpgSQL_stmt *stmt)
 		case PLPGSQL_STMT_RAISE:
 			dump_raise((PLpgSQL_stmt_raise *) stmt);
 			break;
+		case PLPGSQL_STMT_ASSERT:
+			dump_assert((PLpgSQL_stmt_assert *) stmt);
+			break;
 		case PLPGSQL_STMT_EXECSQL:
 			dump_execsql((PLpgSQL_stmt_execsql *) stmt);
 			break;
@@ -1353,6 +1370,25 @@ dump_raise(PLpgSQL_stmt_raise *stmt)
 	dump_indent -= 2;
 }
 
+static void
+dump_assert(PLpgSQL_stmt_assert *stmt)
+{
+	dump_ind();
+	printf("ASSERT ");
+	dump_expr(stmt->cond);
+	printf("\n");
+
+	dump_indent += 2;
+	if (stmt->message != NULL)
+	{
+		dump_ind();
+		printf("    MESSAGE = ");
+		dump_expr(stmt->message);
+		printf("\n");
+	}
+	dump_indent -= 2;
+}
+
 static void
 dump_execsql(PLpgSQL_stmt_execsql *stmt)
 {

src/pl/plpgsql/src/pl_gram.y

@@ -192,7 +192,7 @@ static	void			check_raise_parameters(PLpgSQL_stmt_raise *stmt);
 %type <loop_body>	loop_body
 %type <stmt>	proc_stmt pl_block
 %type <stmt>	stmt_assign stmt_if stmt_loop stmt_while stmt_exit
-%type <stmt>	stmt_return stmt_raise stmt_execsql
+%type <stmt>	stmt_return stmt_raise stmt_assert stmt_execsql
 %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 stmt_foreach_a
@@ -247,6 +247,7 @@ static	void			check_raise_parameters(PLpgSQL_stmt_raise *stmt);
 %token <keyword>	K_ALIAS
 %token <keyword>	K_ALL
 %token <keyword>	K_ARRAY
+%token <keyword>	K_ASSERT
 %token <keyword>	K_BACKWARD
 %token <keyword>	K_BEGIN
 %token <keyword>	K_BY
@@ -871,6 +872,8 @@ proc_stmt		: pl_block ';'
 						{ $$ = $1; }
 				| stmt_raise
 						{ $$ = $1; }
+				| stmt_assert
+						{ $$ = $1; }
 				| stmt_execsql
 						{ $$ = $1; }
 				| stmt_dynexecute
@@ -1847,6 +1850,29 @@ stmt_raise		: K_RAISE
 					}
 				;
 
+stmt_assert		: K_ASSERT
+					{
+						PLpgSQL_stmt_assert		*new;
+						int	tok;
+
+						new = palloc(sizeof(PLpgSQL_stmt_assert));
+
+						new->cmd_type	= PLPGSQL_STMT_ASSERT;
+						new->lineno		= plpgsql_location_to_lineno(@1);
+
+						new->cond = read_sql_expression2(',', ';',
+														 ", or ;",
+														 &tok);
+
+						if (tok == ',')
+							new->message = read_sql_expression(';', ";");
+						else
+							new->message = NULL;
+
+						$$ = (PLpgSQL_stmt *) new;
+					}
+				;
+
 loop_body		: proc_sect K_END K_LOOP opt_label ';'
 					{
 						$$.stmts = $1;
@@ -2315,6 +2341,7 @@ unreserved_keyword	:
 				K_ABSOLUTE
 				| K_ALIAS
 				| K_ARRAY
+				| K_ASSERT
 				| K_BACKWARD
 				| K_CLOSE
 				| K_COLLATE

src/pl/plpgsql/src/pl_handler.c

@@ -44,6 +44,8 @@ int			plpgsql_variable_conflict = PLPGSQL_RESOLVE_ERROR;
 
 bool		plpgsql_print_strict_params = false;
 
+bool		plpgsql_check_asserts = true;
+
 char	   *plpgsql_extra_warnings_string = NULL;
 char	   *plpgsql_extra_errors_string = NULL;
 int			plpgsql_extra_warnings;
@@ -160,6 +162,14 @@ _PG_init(void)
 							 PGC_USERSET, 0,
 							 NULL, NULL, NULL);
 
+	DefineCustomBoolVariable("plpgsql.check_asserts",
+				  gettext_noop("Perform checks given in ASSERT statements."),
+							 NULL,
+							 &plpgsql_check_asserts,
+							 true,
+							 PGC_USERSET, 0,
+							 NULL, NULL, NULL);
+
 	DefineCustomStringVariable("plpgsql.extra_warnings",
 							   gettext_noop("List of programming constructs that should produce a warning."),
 							   NULL,

src/pl/plpgsql/src/pl_scanner.c

@@ -98,6 +98,7 @@ static const ScanKeyword unreserved_keywords[] = {
 	PG_KEYWORD("absolute", K_ABSOLUTE, UNRESERVED_KEYWORD)
 	PG_KEYWORD("alias", K_ALIAS, UNRESERVED_KEYWORD)
 	PG_KEYWORD("array", K_ARRAY, UNRESERVED_KEYWORD)
+	PG_KEYWORD("assert", K_ASSERT, UNRESERVED_KEYWORD)
 	PG_KEYWORD("backward", K_BACKWARD, UNRESERVED_KEYWORD)
 	PG_KEYWORD("close", K_CLOSE, UNRESERVED_KEYWORD)
 	PG_KEYWORD("collate", K_COLLATE, UNRESERVED_KEYWORD)
@@ -607,8 +608,7 @@ plpgsql_scanner_errposition(int location)
  * Beware of using yyerror for other purposes, as the cursor position might
  * be misleading!
  */
-void
-pg_attribute_noreturn
+void pg_attribute_noreturn
 plpgsql_yyerror(const char *message)
 {
 	char	   *yytext = core_yy.scanbuf + plpgsql_yylloc;

src/pl/plpgsql/src/plpgsql.h

@@ -94,6 +94,7 @@ enum PLpgSQL_stmt_types
 	PLPGSQL_STMT_RETURN_NEXT,
 	PLPGSQL_STMT_RETURN_QUERY,
 	PLPGSQL_STMT_RAISE,
+	PLPGSQL_STMT_ASSERT,
 	PLPGSQL_STMT_EXECSQL,
 	PLPGSQL_STMT_DYNEXECUTE,
 	PLPGSQL_STMT_DYNFORS,
@@ -630,6 +631,13 @@ typedef struct
 	PLpgSQL_expr *expr;
 } PLpgSQL_raise_option;
 
+typedef struct
+{								/* ASSERT statement */
+	int			cmd_type;
+	int			lineno;
+	PLpgSQL_expr *cond;
+	PLpgSQL_expr *message;
+} PLpgSQL_stmt_assert;
 
 typedef struct
 {								/* Generic SQL statement to execute */
@@ -889,6 +897,8 @@ extern int	plpgsql_variable_conflict;
 
 extern bool plpgsql_print_strict_params;
 
+extern bool plpgsql_check_asserts;
+
 /* extra compile-time checks */
 #define PLPGSQL_XCHECK_NONE			0
 #define PLPGSQL_XCHECK_SHADOWVAR	1