Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

plpgsql: fix exception handling for nested blocks #122321

Merged
merged 3 commits into from
Apr 16, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
156 changes: 154 additions & 2 deletions pkg/ccl/logictestccl/testdata/logic_test/plpgsql_block
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
1 change: 1 addition & 0 deletions pkg/sql/opt/exec/execbuilder/relational.go
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
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
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
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
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
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
Loading