Merge #122321

122321: plpgsql: fix exception handling for nested blocks r=DrewKimball a=DrewKimball

#### plpgsql: prevent inlining for block-exit continuation

This commit prevents inlining for the continuation that transitions out
of a PL/pgSQL block with an exception handler. This is necessary to
ensure that the statements following the nested block are considered
part of the parent block, not the nested block. Otherwise, an error thrown
after the nested block might still be caught by the nested block's
exception handler, which is incorrect behavior.

#### 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.

#### plpgsql: roll back all cursors within a block that handles an exception

This commit fixes handling for cursors opened within the scope of a
block with an exception handler that has nested blocks or nested routine
calls. Previously, if a PL/pgSQL block caught an exception, only the cursors
opened directly by that block would be rolled back. Any cursors opened by
a nested block or routine call would remain open. Now, the block state
tracks the timestamp when the block's execution began. Once an exception
is caught, all cursors with a timestamp later than the block's start are
rolled back.

Fixes #122278
Fixes #121078

Release note: None

Co-authored-by: Drew Kimball <[email protected]>

pkg/ccl/logictestccl/testdata/logic_test/plpgsql_block

@@ -114,9 +114,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);
@@ -195,6 +236,83 @@ NOTICE: a2
 NOTICE: a3
 NOTICE: a4
 
+# Regression test for #122278 - all cursors within the scope of a block
+# (including those in nested blocks or routines) should be closed when the block
+# catches an exception.
+statement ok
+CREATE PROCEDURE p_nested(curs REFCURSOR) AS $$
+  BEGIN
+    OPEN curs FOR SELECT -100;
+  END;
+$$ LANGUAGE PLpgSQL;
+
+statement ok
+DROP FUNCTION f;
+CREATE FUNCTION f(n INT) RETURNS INT AS $$
+  DECLARE
+    x REFCURSOR;
+    y REFCURSOR;
+  BEGIN
+    OPEN x FOR SELECT 100;
+    BEGIN
+      OPEN y FOR SELECT 200;
+      IF n = 0 THEN
+        RETURN 1 // 0;
+      END IF;
+      CALL p_nested('foo');
+      IF n = 1 THEN
+        RETURN 1 // 0;
+      END IF;
+    EXCEPTION
+      WHEN division_by_zero THEN
+        RETURN (SELECT count(*) FROM pg_cursors);
+    END;
+    CALL p_nested('bar');
+    IF n = 2 THEN
+      RETURN 1 // 0;
+    END IF;
+    RETURN (SELECT count(*) FROM pg_cursors);
+  EXCEPTION
+    WHEN division_by_zero THEN
+      RETURN (SELECT count(*) FROM pg_cursors);
+  END
+$$ LANGUAGE PLpgSQL;
+
+statement ok
+CLOSE ALL;
+
+query I
+SELECT f(0);
+----
+1
+
+statement ok
+CLOSE ALL;
+
+query I
+SELECT f(1);
+----
+1
+
+statement ok
+CLOSE ALL;
+
+query I
+SELECT f(2);
+----
+0
+
+statement ok
+CLOSE ALL;
+
+query I
+SELECT f(3);
+----
+4
+
+statement ok
+CLOSE ALL;
+
 subtest nested_block_exceptions
 
 # Don't catch an exception thrown from the variable declarations.
@@ -341,7 +459,7 @@ CALL p(1);
 ----
 NOTICE: 1
 
-statement error pgcode 34000 pq: cursor \"<unnamed portal 11>\" does not exist
+statement error pgcode 34000 pq: cursor \"<unnamed portal 19>\" does not exist
 CALL p(2);
 
 query T noticetrace
@@ -434,6 +552,40 @@ NOTICE: inner block: 1
 NOTICE: inner handler: 2
 NOTICE: outer handler: 3
 
+# Regression test for #122278 - calling f(1) should result in the second
+# exception handler being triggered.
+statement ok
+CREATE TABLE t122278(x INT);
+
+statement ok
+DROP FUNCTION f;
+CREATE FUNCTION f(n INT) RETURNS INT AS $$
+  BEGIN
+    BEGIN
+      IF n = 0 THEN
+        RETURN 1 // 0;
+      END IF;
+    EXCEPTION
+      WHEN division_by_zero THEN
+        RETURN (SELECT 100 + count(*) FROM t122278);
+    END;
+    RETURN 1 // 0;
+  EXCEPTION
+    WHEN division_by_zero THEN
+      RETURN (SELECT 200 + count(*) FROM t122278);
+  END
+$$ LANGUAGE PLpgSQL;
+
+query I
+SELECT f(0);
+----
+100
+
+query I
+SELECT f(1);
+----
+200
+
 subtest error
 
 statement ok

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/memo/testdata/logprops/tail-calls

@@ -902,7 +902,7 @@ values
            └── body
                 └── values
                      └── tuple
-                          └── udf: exception_block_5
+                          └── udf: nested_block_5
                                ├── tail-call
                                ├── body
                                │    └── values

pkg/sql/opt/optbuilder/plpgsql.go

@@ -372,18 +372,20 @@ func (b *plpgsqlBuilder) buildBlock(astBlock *ast.Block, s *scope) *scope {
 	}
 	// Build the exception handler. This has to happen after building the variable
 	// declarations, since the exception handler can reference the block's vars.
-	if exceptions := b.buildExceptions(astBlock); exceptions != nil {
+	if len(astBlock.Exceptions) > 0 {
+		exceptionBlock := b.buildExceptions(astBlock)
+		block.hasExceptionHandler = true
+
 		// There is an implicit block around the body statements, with an optional
 		// exception handler. Note that the variable declarations are not in block
 		// scope, and exceptions thrown during variable declaration are not caught.
 		//
-		// The routine is volatile to prevent inlining. Only the block and
-		// variable-assignment routines need to be volatile; see the buildExceptions
+		// The routine is volatile to prevent inlining; see the buildExceptions
 		// comment for details.
-		block.hasExceptionHandler = true
-		blockCon := b.makeContinuation("exception_block")
-		blockCon.def.ExceptionBlock = exceptions
+		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)
 	}
@@ -407,7 +409,15 @@ func (b *plpgsqlBuilder) buildPLpgSQLStatements(stmts []ast.Statement, s *scope)
 			// For a nested block, push a continuation with the remaining statements
 			// before calling recursively into buildBlock. The continuation will be
 			// called when the control flow within the nested block terminates.
-			blockCon := b.makeContinuationWithTyp("nested_block", t.Label, continuationBlockExit)
+			blockCon := b.makeContinuationWithTyp("post_nested_block", t.Label, continuationBlockExit)
+			if len(t.Exceptions) > 0 {
+				// If the block has an exception handler, mark the continuation as
+				// volatile to prevent inlining. This is necessary to ensure that
+				// transitions out of a PL/pgSQL block are correctly tracked during
+				// execution. The transition *into* the block is marked volatile for the
+				// same reason; see also buildBlock and buildExceptions.
+				blockCon.def.Volatility = volatility.Volatile
+			}
 			b.appendPlpgSQLStmts(&blockCon, stmts[i+1:])
 			b.pushContinuation(blockCon)
 			return b.buildBlock(t, s)
@@ -1375,13 +1385,16 @@ func (b *plpgsqlBuilder) makeRaiseFormatMessage(
 // cannot throw an exception, and so the "i := 2" assignment will never become
 // visible.
 //
-// The block and assignment continuations must be volatile to prevent inlining.
-// The presence of an exception handler does not impose restrictions on inlining
-// for other continuations.
+// The block entry/exit and assignment continuations for a block with an
+// exception handler must be volatile to prevent inlining. The presence of an
+// exception handler does not impose restrictions on inlining for other types of
+// continuations.
+//
+// Inlining is disabled for the block-exit continuation to ensure that the
+// statements following the nested block are correctly handled as part of the
+// parent block. Otherwise, an error thrown from the parent block could
+// incorrectly be caught by the exception handler of the nested block.
 func (b *plpgsqlBuilder) buildExceptions(block *ast.Block) *memo.ExceptionBlock {
-	if len(block.Exceptions) == 0 {
-		return nil
-	}
 	codes := make([]pgcode.Code, 0, len(block.Exceptions))
 	handlers := make([]*memo.UDFDefinition, 0, len(block.Exceptions))
 	addHandler := func(codeStr string, handler *memo.UDFDefinition) {

pkg/sql/opt/optbuilder/testdata/procedure_plpgsql

@@ -368,11 +368,11 @@ call
                 │                                       │              ├── const: ''
                 │                                       │              └── const: '00000'
                 │                                       └── project
-                │                                            ├── columns: nested_block_3:25
+                │                                            ├── columns: post_nested_block_3:25
                 │                                            ├── values
                 │                                            │    └── tuple
                 │                                            └── projections
-                │                                                 └── udf: nested_block_3 [as=nested_block_3:25]
+                │                                                 └── udf: post_nested_block_3 [as=post_nested_block_3:25]
                 │                                                      ├── tail-call
                 │                                                      ├── args
                 │                                                      │    ├── variable: x:21
@@ -502,11 +502,11 @@ call
                 │                   │              ├── const: ''
                 │                   │              └── const: '00000'
                 │                   └── project
-                │                        ├── columns: exception_block_7:17
+                │                        ├── columns: nested_block_7:17
                 │                        ├── values
                 │                        │    └── tuple
                 │                        └── projections
-                │                             └── udf: exception_block_7 [as=exception_block_7:17]
+                │                             └── udf: nested_block_7 [as=nested_block_7:17]
                 │                                  ├── tail-call
                 │                                  ├── args
                 │                                  │    └── variable: x:2
@@ -532,11 +532,11 @@ call
                 │                                  │                        │              ├── const: 1
                 │                                  │                        │              └── const: 0
                 │                                  │                        └── project
-                │                                  │                             ├── columns: nested_block_3:15
+                │                                  │                             ├── columns: post_nested_block_3:15
                 │                                  │                             ├── values
                 │                                  │                             │    └── tuple
                 │                                  │                             └── projections
-                │                                  │                                  └── udf: nested_block_3 [as=nested_block_3:15]
+                │                                  │                                  └── udf: post_nested_block_3 [as=post_nested_block_3:15]
                 │                                  │                                       ├── tail-call
                 │                                  │                                       ├── args
                 │                                  │                                       │    └── variable: x:13
@@ -581,7 +581,7 @@ call
                 │                                  └── exception-handler
                 │                                       └── SQLSTATE '22012'
                 │                                            └── project
-                │                                                 ├── columns: nested_block_3:11
+                │                                                 ├── columns: post_nested_block_3:11
                 │                                                 ├── barrier
                 │                                                 │    ├── columns: x:10!null
                 │                                                 │    └── project
@@ -591,7 +591,7 @@ call
                 │                                                 │         └── projections
                 │                                                 │              └── const: 100 [as=x:10]
                 │                                                 └── projections
-                │                                                      └── udf: nested_block_3 [as=nested_block_3:11]
+                │                                                      └── udf: post_nested_block_3 [as=post_nested_block_3:11]
                 │                                                           ├── args
                 │                                                           │    └── variable: x:10
                 │                                                           ├── params: x:4