plpgsql: handle EXIT/CONTINUE for root block and correct errors

This commit cleans up `EXIT` and `CONTINUE` with label handling. It is
now possible to `EXIT` from the root block or the routine itself, and
the various error messages that can result from incorrect usage have
been corrected.

Informs #115271

Release note: None

pkg/ccl/logictestccl/testdata/logic_test/procedure_plpgsql

@@ -300,8 +300,8 @@ CREATE PROCEDURE p() AS $$
   END
 $$ LANGUAGE PLpgSQL;
 
-# CONTINUE cannot apply to a block, with or without a label.
-statement error pgcode 42601 pq: CONTINUE cannot be used outside a loop
+# CONTINUE cannot apply to a block.
+statement error pgcode 42601 pq: block label \"foo\" cannot be used in CONTINUE
 CREATE PROCEDURE p() AS $$
   BEGIN
     <<foo>>
@@ -313,8 +313,86 @@ CREATE PROCEDURE p() AS $$
   END
 $$ LANGUAGE PLpgSQL;
 
+# The nested block takes precedence over the loop with the same label.
+statement error pgcode 42601 pq: block label \"foo\" cannot be used in CONTINUE
+CREATE PROCEDURE p() AS $$
+  DECLARE
+    i INT := 0;
+  BEGIN
+    <<foo>>
+    WHILE i < 5 LOOP
+      <<foo>>
+      BEGIN
+        RAISE NOTICE 'before EXIT';
+        CONTINUE foo;
+        RAISE NOTICE 'after EXIT';
+      END;
+    END LOOP;
+  END
+$$ LANGUAGE PLpgSQL;
+
+# EXIT with a nonexistent label.
+statement error pgcode 42601 pq: there is no label \"foo\" attached to any block or loop enclosing this statement
+CREATE PROCEDURE p() AS $$
+  BEGIN
+    <<bar>>
+    BEGIN
+      EXIT foo;
+    END;
+  END
+$$ LANGUAGE PLpgSQL;
+
+# CONTINUE with a nonexistent label.
+statement error pgcode 42601 pq: there is no label \"foo\" attached to any block or loop enclosing this statement
+CREATE PROCEDURE p() AS $$
+  BEGIN
+    <<bar>>
+    BEGIN
+      CONTINUE foo;
+    END;
+  END
+$$ LANGUAGE PLpgSQL;
+
+# It is possible to EXIT the root block.
+statement ok
+CREATE PROCEDURE p() AS $$
+  <<foo>>
+  DECLARE
+    i INT := 0;
+  BEGIN
+    WHILE i < 5 LOOP
+      RAISE NOTICE 'here';
+      EXIT foo;
+      RAISE NOTICE 'still here';
+    END LOOP;
+  END
+$$ LANGUAGE PLpgSQL;
+
+query T noticetrace
+CALL p();
+----
+NOTICE: here
+
+# It is possible to EXIT the routine, but this always results in an
+# end-of-function error, even for a void-returning proc.
+statement ok
+DROP PROCEDURE p;
+CREATE PROCEDURE p() AS $$
+  DECLARE
+    i INT := 0;
+  BEGIN
+    WHILE i < 5 LOOP
+      EXIT p;
+    END LOOP;
+  END
+$$ LANGUAGE PLpgSQL;
+
+statement error pgcode 2F005 control reached end of function without RETURN
+CALL p();
+
 # CONTINUE the inner loop.
 statement ok
+DROP PROCEDURE p;
 CREATE PROCEDURE p(x INT) AS $$
   <<b1>>
   DECLARE
@@ -349,12 +427,10 @@ CREATE PROCEDURE p(x INT) AS $$
           ELSIF x = 4 THEN
             IF j = 1 THEN RAISE NOTICE 'EXIT b2'; END IF;
             EXIT b2 WHEN j = 1;
-          -- TODO(drewk): uncomment these branches when it's possible to
-          -- reference the root block.
-          -- ELSIF x = 5 THEN
-          --   EXIT b1 WHEN j = 1;
-          -- ELSE
-          --   EXIT p WHEN j = 1;
+          ELSIF x = 5 THEN
+            EXIT b1 WHEN j = 1;
+          ELSE
+            EXIT p WHEN j = 1;
           END IF;
           RAISE NOTICE '<< l2 % %', i, j;
         END LOOP l2;
@@ -449,6 +525,19 @@ NOTICE: EXIT b2
 NOTICE: << l1 2
 NOTICE: << b1 2
 
+# EXIT outer block.
+query T noticetrace
+CALL p(5);
+----
+NOTICE: >> b1 0
+NOTICE: >> l1 1
+NOTICE: >> b2 1 0
+NOTICE: >> l2 1 1
+
+# EXIT the routine.
+statement error pgcode 2F005 control reached end of function without RETURN
+CALL p(6);
+
 statement ok
 DROP PROCEDURE p;
 CREATE PROCEDURE p() AS $$

pkg/sql/opt/optbuilder/create_function.go

@@ -336,7 +336,8 @@ func (b *Builder) buildCreateFunction(cf *tree.CreateRoutine, inScope *scope) (o
 		// the volatility.
 		b.factory.FoldingControl().TemporarilyDisallowStableFolds(func() {
 			plBuilder := newPLpgSQLBuilder(
-				b, cf.Name.Object(), nil /* colRefs */, routineParams, funcReturnType, cf.IsProcedure,
+				b, cf.Name.Object(), stmt.AST.Label, nil, /* colRefs */
+				routineParams, funcReturnType, cf.IsProcedure,
 			)
 			stmtScope = plBuilder.buildRootBlock(stmt.AST, bodyScope, routineParams)
 		})

pkg/sql/opt/optbuilder/plpgsql.go

@@ -166,6 +166,7 @@ type plpgsqlBuilder struct {
 	// outParams is the set of OUT parameters for the routine.
 	outParams []ast.Variable
 
+	routineName  string
 	isProcedure  bool
 	identCounter int
 }
@@ -180,7 +181,7 @@ type routineParam struct {
 
 func newPLpgSQLBuilder(
 	ob *Builder,
-	routineName string,
+	routineName, rootBlockLabel string,
 	colRefs *opt.ColSet,
 	routineParams []routineParam,
 	returnType *types.T,
@@ -192,12 +193,13 @@ func newPLpgSQLBuilder(
 		colRefs:     colRefs,
 		returnType:  returnType,
 		blocks:      make([]plBlock, 0, initialBlocksCap),
+		routineName: routineName,
 		isProcedure: isProcedure,
 	}
 	// Build the initial block for the routine parameters, which are considered
 	// PL/pgSQL variables.
 	b.pushBlock(plBlock{
-		label:    routineName,
+		label:    rootBlockLabel,
 		vars:     make([]ast.Variable, 0, len(routineParams)),
 		varTypes: make(map[ast.Variable]*types.T),
 	})
@@ -581,7 +583,24 @@ func (b *plpgsqlBuilder) buildPLpgSQLStatements(stmts []ast.Statement, s *scope)
 			if con := b.getContinuation(conTypes, t.Label); con != nil {
 				return b.callContinuation(con, s)
 			}
-			panic(exitOutsideLoopErr)
+			if t.Label == unspecifiedLabel {
+				panic(exitOutsideLoopErr)
+			}
+			if t.Label == b.rootBlock().label {
+				// An EXIT from the root block has the same handling as when the routine
+				// ends with no RETURN statement.
+				return b.handleEndOfFunction(s)
+			}
+			if t.Label == b.routineName {
+				// An EXIT from the routine name itself results in an end-of-function
+				// error, even for a VOID-returning routine or one with OUT-parameters.
+				eofCon := b.makeContinuationWithTyp("root_exit", t.Label, continuationBlockExit)
+				b.buildEndOfFunctionRaise(&eofCon)
+				return b.callContinuation(&eofCon, s)
+			}
+			panic(pgerror.Newf(pgcode.Syntax,
+				"there is no label \"%s\" attached to any block or loop enclosing this statement", t.Label,
+			))
 
 		case *ast.Continue:
 			if t.Condition != nil {
@@ -595,10 +614,24 @@ func (b *plpgsqlBuilder) buildPLpgSQLStatements(stmts []ast.Statement, s *scope)
 			}
 			// CONTINUE statements are handled by calling the function that executes
 			// the loop body. Errors if used outside a loop.
-			if con := b.getContinuation(continuationLoopContinue, t.Label); con != nil {
+			conTypes := continuationLoopContinue
+			if t.Label != unspecifiedLabel {
+				conTypes |= continuationBlockExit
+			}
+			if con := b.getContinuation(conTypes, t.Label); con != nil {
+				if con.typ == continuationBlockExit {
+					panic(pgerror.Newf(pgcode.Syntax,
+						"block label \"%s\" cannot be used in CONTINUE", t.Label,
+					))
+				}
 				return b.callContinuation(con, s)
 			}
-			panic(continueOutsideLoopErr)
+			if t.Label == unspecifiedLabel {
+				panic(continueOutsideLoopErr)
+			}
+			panic(pgerror.Newf(pgcode.Syntax,
+				"there is no label \"%s\" attached to any block or loop enclosing this statement", t.Label,
+			))
 
 		case *ast.Raise:
 			// RAISE statements allow the PLpgSQL function to send an error or a
@@ -1364,25 +1397,33 @@ func (b *plpgsqlBuilder) handleEndOfFunction(inScope *scope) *scope {
 		return returnScope
 	}
 	// Build a RAISE statement which throws an end-of-function error if executed.
+	con := b.makeContinuation("_end_of_function")
+	b.buildEndOfFunctionRaise(&con)
+	return b.callContinuation(&con, inScope)
+}
+
+// buildEndOfFunctionRaise adds to the given continuation a RAISE statement that
+// throws an end-of-function error, as well as a typed RETURN NULL to ensure
+// that type-checking works out.
+func (b *plpgsqlBuilder) buildEndOfFunctionRaise(con *continuation) {
 	args := b.makeConstRaiseArgs(
 		"ERROR", /* severity */
 		"control reached end of function without RETURN", /* message */
 		"", /* detail */
 		"", /* hint */
 		pgcode.RoutineExceptionFunctionExecutedNoReturnStatement.String(), /* code */
 	)
-	con := b.makeContinuation("_end_of_function")
 	con.def.Volatility = volatility.Volatile
-	b.appendBodyStmt(&con, b.buildPLpgSQLRaise(con.s, args))
+	b.appendBodyStmt(con, b.buildPLpgSQLRaise(con.s, args))
+
 	// Build a dummy statement that returns NULL. It won't be executed, but
 	// ensures that the continuation routine's return type is correct.
 	eofColName := scopeColName("").WithMetadataName(b.makeIdentifier("end_of_function"))
 	eofScope := con.s.push()
 	typedNull := b.ob.factory.ConstructNull(b.returnType)
 	b.ob.synthesizeColumn(eofScope, eofColName, b.returnType, nil /* expr */, typedNull)
-	b.ob.constructProjectForScope(inScope, eofScope)
-	b.appendBodyStmt(&con, eofScope)
-	return b.callContinuation(&con, inScope)
+	b.ob.constructProjectForScope(con.s, eofScope)
+	b.appendBodyStmt(con, eofScope)
 }
 
 // addOneRowCheck handles INTO STRICT, where a SQL statement is required to
@@ -1936,6 +1977,11 @@ func (b *plpgsqlBuilder) parentBlock() *plBlock {
 	return &b.blocks[len(b.blocks)-2]
 }
 
+// rootBlock returns the root block that encapsulates the entire routine.
+func (b *plpgsqlBuilder) rootBlock() *plBlock {
+	return &b.blocks[0]
+}
+
 // pushBlock puts the given block on the stack. It is used when entering the
 // scope of a PL/pgSQL block.
 func (b *plpgsqlBuilder) pushBlock(bs plBlock) *plBlock {

pkg/sql/opt/optbuilder/routine.go

@@ -353,7 +353,9 @@ func (b *Builder) buildRoutine(
 		}
 		var expr memo.RelExpr
 		var physProps *physical.Required
-		plBuilder := newPLpgSQLBuilder(b, def.Name, colRefs, routineParams, rtyp, isProc)
+		plBuilder := newPLpgSQLBuilder(
+			b, def.Name, stmt.AST.Label, colRefs, routineParams, rtyp, isProc,
+		)
 		stmtScope := plBuilder.buildRootBlock(stmt.AST, bodyScope, routineParams)
 		finishResolveType(stmtScope)
 		expr, physProps, isMultiColDataSource =