plpgsql: keep track of the subroutine that begins a PL/pgSQL block

This commit adds logic to keep track of the PL/pgSQL sub-routine that
logically transitions into a PL/pgSQL block with an exception handler.
This is necessary to ensure that the state shared between sub-routines
within the same block is correctly initialized. Previously, the block
state was only initialized once, but this is incorrect for loops, which
need to re-initialize the state on each iteration.

Fixes #122278

Release note: None

pkg/ccl/logictestccl/testdata/logic_test/plpgsql_block

@@ -116,9 +116,50 @@ NOTICE: 2, 1
 NOTICE: final j: 2
 NOTICE: final i: 3
 
-subtest nested_block_cursors
+# Regression test for #122278 - a nested block with an exception handler inside
+# a loop should only rollback mutations from the current iteration.
+statement ok
+CREATE TABLE t122278(x INT);
+
 statement ok
 DROP PROCEDURE p;
+CREATE PROCEDURE p() AS $$
+  DECLARE
+    i INT := 0;
+  BEGIN
+    WHILE i < 5 LOOP
+      i := i + 1;
+      BEGIN
+        INSERT INTO t122278 VALUES (i);
+        IF i = 3 THEN
+          SELECT 1 // 0;
+        END IF;
+      EXCEPTION WHEN division_by_zero THEN
+        RAISE NOTICE 'saw exception';
+      END;
+    END LOOP;
+  END;
+$$ LANGUAGE PLpgSQL;
+
+query T noticetrace
+CALL p();
+----
+NOTICE: saw exception
+
+query I rowsort
+SELECT * FROM t122278;
+----
+1
+2
+4
+5
+
+statement ok
+DROP TABLE t122278 CASCADE;
+
+subtest nested_block_cursors
+
+statement ok
 CREATE PROCEDURE p() AS $$
   DECLARE
     curs1 CURSOR FOR SELECT 1 FROM generate_series(1, 10);

pkg/sql/opt/exec/execbuilder/relational.go

@@ -3392,6 +3392,7 @@ func (b *Builder) buildCall(c *memo.CallExpr) (_ execPlan, outputCols colOrdMap,
 		udf.Def.SetReturning,
 		false, /* tailCall */
 		true,  /* procedure */
+		false, /* blockStart */
 		nil,   /* blockState */
 		nil,   /* cursorDeclaration */
 	)

pkg/sql/opt/exec/execbuilder/scalar.go

@@ -702,6 +702,7 @@ func (b *Builder) buildExistsSubquery(
 				false, /* generator */
 				false, /* tailCall */
 				false, /* procedure */
+				false, /* blockStart */
 				nil,   /* blockState */
 				nil,   /* cursorDeclaration */
 			),
@@ -822,6 +823,7 @@ func (b *Builder) buildSubquery(
 			false, /* generator */
 			false, /* tailCall */
 			false, /* procedure */
+			false, /* blockStart */
 			nil,   /* blockState */
 			nil,   /* cursorDeclaration */
 		), nil
@@ -881,6 +883,7 @@ func (b *Builder) buildSubquery(
 			false, /* generator */
 			false, /* tailCall */
 			false, /* procedure */
+			false, /* blockStart */
 			nil,   /* blockState */
 			nil,   /* cursorDeclaration */
 		), nil
@@ -992,6 +995,7 @@ func (b *Builder) buildUDF(ctx *buildScalarCtx, scalar opt.ScalarExpr) (tree.Typ
 		udf.Def.SetReturning,
 		tailCall,
 		false, /* procedure */
+		udf.Def.BlockStart,
 		blockState,
 		udf.Def.CursorDeclaration,
 	), nil
@@ -1047,6 +1051,7 @@ func (b *Builder) initRoutineExceptionHandler(
 			action.SetReturning,
 			false, /* tailCall */
 			false, /* procedure */
+			false, /* blockStart */
 			nil,   /* blockState */
 			nil,   /* cursorDeclaration */
 		)

pkg/sql/opt/memo/expr.go

@@ -709,10 +709,16 @@ type UDFDefinition struct {
 	// data source.
 	MultiColDataSource bool
 
-	// IsRecursive indicates whether the UDF recursively calls itself. This
+	// IsRecursive indicates whether the routine recursively calls itself. This
 	// applies to direct as well as indirect recursive calls (mutual recursion).
 	IsRecursive bool
 
+	// BlockStart indicates whether the routine marks the start of a PL/pgSQL
+	// block with an exception handler. This is used to determine when to
+	// initialize the common state held between sub-routines within the same
+	// block.
+	BlockStart bool
+
 	// RoutineType indicates whether this routine is a UDF, stored procedure, or
 	// builtin function.
 	RoutineType tree.RoutineType

pkg/sql/opt/optbuilder/plpgsql.go

@@ -385,6 +385,7 @@ func (b *plpgsqlBuilder) buildBlock(astBlock *ast.Block, s *scope) *scope {
 		blockCon := b.makeContinuation("nested_block")
 		blockCon.def.ExceptionBlock = exceptionBlock
 		blockCon.def.Volatility = volatility.Volatile
+		blockCon.def.BlockStart = true
 		b.appendPlpgSQLStmts(&blockCon, astBlock.Body)
 		return b.callContinuation(&blockCon, s)
 	}

pkg/sql/routine.go

@@ -439,19 +439,17 @@ func (g *routineGenerator) closeCursors(blockState *tree.BlockState) error {
 	return err
 }
 
-// maybeInitBlockState creates a savepoint if all the following are true:
-//  1. The current routine is within a PLpgSQL exception block.
-//  2. The current block has an exception handler
-//  3. The savepoint hasn't already been created for this block.
+// maybeInitBlockState creates a savepoint for a routine that marks a transition
+// into a PL/pgSQL block with an exception handler.
 //
 // Note that it is not necessary to explicitly release the savepoint at any
 // point, because it does not add any overhead.
 func (g *routineGenerator) maybeInitBlockState(ctx context.Context) error {
 	blockState := g.expr.BlockState
-	if blockState == nil {
+	if blockState == nil || !g.expr.BlockStart {
 		return nil
 	}
-	if blockState.ExceptionHandler != nil && blockState.SavepointTok == nil {
+	if blockState.ExceptionHandler != nil {
 		// Drop down a savepoint for the current scope.
 		var err error
 		if blockState.SavepointTok, err = g.p.Txn().CreateSavepoint(ctx); err != nil {

pkg/sql/sem/tree/routine.go

@@ -128,6 +128,11 @@ type RoutineExpr struct {
 	// Procedure is true if the routine is a procedure being invoked by CALL.
 	Procedure bool
 
+	// BlockStart is true if this routine marks the start of a PL/pgSQL block with
+	// an exception handler. It determines when to initialize the state shared
+	// between sub-routines for the block.
+	BlockStart bool
+
 	// BlockState holds the information needed to coordinate error-handling
 	// between the sub-routines that make up a PLpgSQL exception block.
 	BlockState *BlockState
@@ -149,6 +154,7 @@ func NewTypedRoutineExpr(
 	generator bool,
 	tailCall bool,
 	procedure bool,
+	blockStart bool,
 	blockState *BlockState,
 	cursorDeclaration *RoutineOpenCursor,
 ) *RoutineExpr {
@@ -163,6 +169,7 @@ func NewTypedRoutineExpr(
 		Generator:         generator,
 		TailCall:          tailCall,
 		Procedure:         procedure,
+		BlockStart:        blockStart,
 		BlockState:        blockState,
 		CursorDeclaration: cursorDeclaration,
 	}