From 47c90ee64ce08b4ffcb69c61fea53882f724972e Mon Sep 17 00:00:00 2001 From: Drew Kimball Date: Fri, 12 Apr 2024 16:00:21 -0600 Subject: [PATCH] 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 --- .../testdata/logic_test/plpgsql_block | 43 ++++++++++++++++++- pkg/sql/opt/exec/execbuilder/relational.go | 1 + pkg/sql/opt/exec/execbuilder/scalar.go | 5 +++ pkg/sql/opt/memo/expr.go | 8 +++- pkg/sql/opt/optbuilder/plpgsql.go | 1 + pkg/sql/routine.go | 10 ++--- pkg/sql/sem/tree/routine.go | 7 +++ 7 files changed, 67 insertions(+), 8 deletions(-) diff --git a/pkg/ccl/logictestccl/testdata/logic_test/plpgsql_block b/pkg/ccl/logictestccl/testdata/logic_test/plpgsql_block index 2143df8e9869..5db10d8f53ec 100644 --- a/pkg/ccl/logictestccl/testdata/logic_test/plpgsql_block +++ b/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); diff --git a/pkg/sql/opt/exec/execbuilder/relational.go b/pkg/sql/opt/exec/execbuilder/relational.go index 7ce3389c957b..f58618db95ce 100644 --- a/pkg/sql/opt/exec/execbuilder/relational.go +++ b/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 */ ) diff --git a/pkg/sql/opt/exec/execbuilder/scalar.go b/pkg/sql/opt/exec/execbuilder/scalar.go index fb327c8599c2..7b0ec2be500d 100644 --- a/pkg/sql/opt/exec/execbuilder/scalar.go +++ b/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 */ ) diff --git a/pkg/sql/opt/memo/expr.go b/pkg/sql/opt/memo/expr.go index 8d5035c6abb3..79d13cfec5f4 100644 --- a/pkg/sql/opt/memo/expr.go +++ b/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 diff --git a/pkg/sql/opt/optbuilder/plpgsql.go b/pkg/sql/opt/optbuilder/plpgsql.go index 3299366d70f3..49991f1cbaf6 100644 --- a/pkg/sql/opt/optbuilder/plpgsql.go +++ b/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) } diff --git a/pkg/sql/routine.go b/pkg/sql/routine.go index 573da708a57c..e665dbfcfa56 100644 --- a/pkg/sql/routine.go +++ b/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 { diff --git a/pkg/sql/sem/tree/routine.go b/pkg/sql/sem/tree/routine.go index 2fef6325d766..5e3c3cf5400a 100644 --- a/pkg/sql/sem/tree/routine.go +++ b/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, }