plpgsql: implement CONTINUE/EXIT with condition

PL/pgSQL allows adding a `WHEN <condition>` clause to `EXIT` and
`CONTINUE` statements. This is syntactic sugar for `EXIT` or `CONTINUE`
within an `IF` statement. This commit add support for this syntax.

Informs #115271

Release note (sql change): It is now possible to specify a condition
for the PL/pgSQL statemenets `EXIT` and `CONTINUE`.

pkg/ccl/logictestccl/testdata/logic_test/procedure_plpgsql

@@ -178,4 +178,66 @@ CALL get_rows(3);
 ----
 NOTICE: (1,)
 
+subtest exit_continue_cond
+
+statement ok
+CREATE PROCEDURE p(a INT, b INT, c INT) AS $$
+  DECLARE
+    i INT := 0;
+  BEGIN
+    LOOP
+      RAISE NOTICE 'iteration %', i;
+      i := i + 1;
+      CONTINUE WHEN i = a;
+      RAISE NOTICE 'new value of i: %', i;
+      EXIT WHEN i >= b;
+      CONTINUE WHEN i <> c;
+      RAISE NOTICE 'returning';
+      RETURN;
+    END LOOP;
+    RAISE NOTICE 'exited loop';
+  END
+$$ LANGUAGE PLpgSQL;
+
+query T noticetrace
+CALL p(2, 1, 4);
+----
+NOTICE: iteration 0
+NOTICE: new value of i: 1
+NOTICE: exited loop
+
+query T noticetrace
+CALL p(2, 3, 4);
+----
+NOTICE: iteration 0
+NOTICE: new value of i: 1
+NOTICE: iteration 1
+NOTICE: iteration 2
+NOTICE: new value of i: 3
+NOTICE: exited loop
+
+query T noticetrace
+CALL p(2, 4, 4);
+----
+NOTICE: iteration 0
+NOTICE: new value of i: 1
+NOTICE: iteration 1
+NOTICE: iteration 2
+NOTICE: new value of i: 3
+NOTICE: iteration 3
+NOTICE: new value of i: 4
+NOTICE: exited loop
+
+query T noticetrace
+CALL p(2, 5, 4);
+----
+NOTICE: iteration 0
+NOTICE: new value of i: 1
+NOTICE: iteration 1
+NOTICE: iteration 2
+NOTICE: new value of i: 3
+NOTICE: iteration 3
+NOTICE: new value of i: 4
+NOTICE: returning
+
 subtest end

pkg/sql/opt/optbuilder/plpgsql.go

@@ -564,17 +564,19 @@ func (b *plpgsqlBuilder) buildPLpgSQLStatements(stmts []ast.Statement, s *scope)
 					ElseBody:  []ast.Statement{&ast.Exit{}},
 				}},
 			}
-			newStmts := make([]ast.Statement, 0, len(stmts))
-			newStmts = append(newStmts, loop)
-			newStmts = append(newStmts, stmts[i+1:]...)
-			return b.buildPLpgSQLStatements(newStmts, s)
+			return b.buildPLpgSQLStatements(b.prependStmt(loop, stmts[i+1:]), s)
 
 		case *ast.Exit:
 			if t.Label != "" {
 				panic(exitLabelErr)
 			}
 			if t.Condition != nil {
-				panic(exitCondErr)
+				// EXIT with a condition is syntactic sugar for EXIT inside an IF stmt.
+				ifStmt := &ast.If{
+					Condition: t.Condition,
+					ThenBody:  []ast.Statement{&ast.Exit{Label: t.Label}},
+				}
+				return b.buildPLpgSQLStatements(b.prependStmt(ifStmt, stmts[i+1:]), s)
 			}
 			// EXIT statements are handled by calling the function that executes the
 			// statements after a loop. Errors if used outside a loop.
@@ -589,7 +591,13 @@ func (b *plpgsqlBuilder) buildPLpgSQLStatements(stmts []ast.Statement, s *scope)
 				panic(continueLabelErr)
 			}
 			if t.Condition != nil {
-				panic(continueCondErr)
+				// CONTINUE with a condition is syntactic sugar for CONTINUE inside an
+				// IF stmt.
+				ifStmt := &ast.If{
+					Condition: t.Condition,
+					ThenBody:  []ast.Statement{&ast.Continue{Label: t.Label}},
+				}
+				return b.buildPLpgSQLStatements(b.prependStmt(ifStmt, stmts[i+1:]), s)
 			}
 			// CONTINUE statements are handled by calling the function that executes
 			// the loop body. Errors if used outside a loop.
@@ -1747,6 +1755,12 @@ func (b *plpgsqlBuilder) resolveVariableForAssign(name ast.Variable) *types.T {
 	panic(pgerror.Newf(pgcode.Syntax, "\"%s\" is not a known variable", name))
 }
 
+func (b *plpgsqlBuilder) prependStmt(stmt ast.Statement, stmts []ast.Statement) []ast.Statement {
+	newStmts := make([]ast.Statement, 0, len(stmts)+1)
+	newStmts = append(newStmts, stmt)
+	return append(newStmts, stmts...)
+}
+
 func (b *plpgsqlBuilder) ensureScopeHasExpr(s *scope) {
 	if s.expr == nil {
 		s.expr = b.ob.factory.ConstructNoColsRow()
@@ -2003,15 +2017,9 @@ var (
 	exitLabelErr = unimplemented.New("EXIT label",
 		"EXIT statement labels are not yet supported",
 	)
-	exitCondErr = unimplemented.New("EXIT WHEN",
-		"conditional EXIT statements are not yet supported",
-	)
 	continueLabelErr = unimplemented.New("CONTINUE label",
 		"CONTINUE statement labels are not yet supported",
 	)
-	continueCondErr = unimplemented.New("CONTINUE WHEN",
-		"conditional CONTINUE statements are not yet supported",
-	)
 	dupIntoErr = unimplemented.New("duplicate INTO target",
 		"assigning to a variable more than once in the same INTO statement is not supported",
 	)