Skip to content

Commit

Permalink
Merge pull request cockroachdb#122433 from cockroachdb/blathers/backp…
Browse files Browse the repository at this point in the history 
…ort-release-24.1-122321

release-24.1: plpgsql: fix exception handling for nested blocks
  • Loading branch information
@DrewKimball
DrewKimball authored Apr 16, 2024
2 parents b7c39a8 + 84039d1 commit 425ba75
Show file tree
Hide file tree
Showing 10 changed files with 270 additions and 89 deletions.
156 changes: 154 additions & 2 deletions pkg/ccl/logictestccl/testdata/logic_test/plpgsql_block
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -197,6 +238,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.
Expand Down Expand Up @@ -343,7 +461,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
Expand Down Expand Up @@ -436,6 +554,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
Expand Down
1 change: 1 addition & 0 deletions pkg/sql/opt/exec/execbuilder/relational.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 */
)
Expand Down
5 changes: 5 additions & 0 deletions pkg/sql/opt/exec/execbuilder/scalar.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -702,6 +702,7 @@ func (b *Builder) buildExistsSubquery(
false, /* generator */
false, /* tailCall */
false, /* procedure */
false, /* blockStart */
nil, /* blockState */
nil, /* cursorDeclaration */
),
Expand Down Expand Up @@ -822,6 +823,7 @@ func (b *Builder) buildSubquery(
false, /* generator */
false, /* tailCall */
false, /* procedure */
false, /* blockStart */
nil, /* blockState */
nil, /* cursorDeclaration */
), nil
Expand Down Expand Up @@ -881,6 +883,7 @@ func (b *Builder) buildSubquery(
false, /* generator */
false, /* tailCall */
false, /* procedure */
false, /* blockStart */
nil, /* blockState */
nil, /* cursorDeclaration */
), nil
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -1047,6 +1051,7 @@ func (b *Builder) initRoutineExceptionHandler(
action.SetReturning,
false, /* tailCall */
false, /* procedure */
false, /* blockStart */
nil, /* blockState */
nil, /* cursorDeclaration */
)
Expand Down
8 changes: 7 additions & 1 deletion pkg/sql/opt/memo/expr.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion pkg/sql/opt/memo/testdata/logprops/tail-calls
View file
Original file line number Diff line number Diff line change
Expand Up @@ -902,7 +902,7 @@ values
└── body
└── values
└── tuple
└── udf: exception_block_5
└── udf: nested_block_5
├── tail-call
├── body
│ └── values
Expand Down
39 changes: 26 additions & 13 deletions pkg/sql/opt/optbuilder/plpgsql.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
Expand All @@ -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)
Expand Down Expand Up @@ -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) {
Expand Down
16 changes: 8 additions & 8 deletions pkg/sql/opt/optbuilder/testdata/procedure_plpgsql
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down
Loading

0 comments on commit 425ba75

Please sign in to comment.