Skip to content

Commit

Permalink
plpgsql: prevent inlining for block-exit continuation
Browse files Browse the repository at this point in the history 
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.

Informs cockroachdb#122278

Release note: None
  • Loading branch information
@DrewKimball
DrewKimball committed Apr 13, 2024
1 parent 900f8b6 commit 2f7a944
Show file tree
Hide file tree
Showing 5 changed files with 102 additions and 56 deletions.
34 changes: 34 additions & 0 deletions pkg/ccl/logictestccl/testdata/logic_test/plpgsql_block
View file
Original file line number Diff line number Diff line change
Expand Up @@ -434,6 +434,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
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
38 changes: 25 additions & 13 deletions pkg/sql/opt/optbuilder/plpgsql.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -372,17 +372,18 @@ 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
b.appendPlpgSQLStmts(&blockCon, astBlock.Body)
return b.callContinuation(&blockCon, s)
Expand All @@ -407,7 +408,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 +1384,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 2f7a944

Please sign in to comment.