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

release-23.2: plpgsql: don't exit early when SELECT INTO returns no rows #115676

Merged
merged 1 commit into from
Dec 6, 2023
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.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
88 changes: 88 additions & 0 deletions pkg/sql/logictest/testdata/logic_test/udf_plpgsql
Original file line number Diff line number Diff line change
Expand Up @@ -1907,6 +1907,32 @@ SELECT f();
NOTICE: i = 104
NOTICE: i = 103

# When the SQL statement returns zero rows, the target variables are all set
# to NULL.
statement ok
CREATE OR REPLACE FUNCTION f() RETURNS INT AS $$
DECLARE
i INT := 0;
BEGIN
RAISE NOTICE 'i = %', i;
SELECT x INTO i FROM xy WHERE False;
RAISE NOTICE 'i = %', i;
i := 1;
RAISE NOTICE 'i = %', i;
INSERT INTO xy (SELECT 1, 2 FROM generate_series(1, 0)) RETURNING x INTO i;
RAISE NOTICE 'i = %', i;
RETURN 0;
END
$$ LANGUAGE PLpgSQL;

query T noticetrace
SELECT f();
----
NOTICE: i = 0
NOTICE: i = <NULL>
NOTICE: i = 1
NOTICE: i = <NULL>

statement error pgcode 0A000 pq: unimplemented: assigning to a variable more than once in the same INTO statement is not supported
CREATE OR REPLACE FUNCTION f() RETURNS INT AS $$
DECLARE
Expand Down Expand Up @@ -2730,4 +2756,66 @@ SELECT f();
----
1

# Regression test for #114826 - PLpgSQL routines should not prematurely exit
# when a SELECT INTO statement returns no rows.
subtest no_row_select_into

statement ok
CREATE TABLE t114826 (a INT);

statement ok
CREATE FUNCTION f114826() RETURNS INT AS $$
DECLARE
found INT;
BEGIN
RAISE NOTICE 'HERE 1';
SELECT a INTO found FROM t114826 WHERE a = 3;
RAISE NOTICE 'HERE 2';
INSERT INTO t114826 (SELECT * FROM t114826) RETURNING a INTO found;
RAISE NOTICE 'HERE 3';
RETURN 10;
END
$$ LANGUAGE PLpgSQL;

query T noticetrace
SELECT f114826();
----
NOTICE: HERE 1
NOTICE: HERE 2
NOTICE: HERE 3

query I
SELECT f114826();
----
10

statement ok
INSERT INTO t114826 VALUES (1);

query T noticetrace
SELECT f114826();
----
NOTICE: HERE 1
NOTICE: HERE 2
NOTICE: HERE 3

query I
SELECT * FROM t114826;
----
1
1

query I
SELECT f114826();
----
10

query I
SELECT * FROM t114826;
----
1
1
1
1

subtest end
5 changes: 1 addition & 4 deletions pkg/sql/opt/norm/decorrelate_funcs.go
Original file line number Diff line number Diff line change
Expand Up @@ -742,10 +742,7 @@ func (c *CustomFuncs) ConstructBinary(op opt.Operator, left, right opt.ScalarExp
// ConstructNoColsRow returns a Values operator having a single row with zero
// columns.
func (c *CustomFuncs) ConstructNoColsRow() memo.RelExpr {
return c.f.ConstructValues(memo.ScalarListWithEmptyTuple, &memo.ValuesPrivate{
Cols: opt.ColList{},
ID: c.f.Metadata().NextUniqueID(),
})
return c.f.ConstructNoColsRow()
}

// referenceSingleColumn returns a Variable operator that refers to the one and
Expand Down
8 changes: 8 additions & 0 deletions pkg/sql/opt/norm/factory.go
Original file line number Diff line number Diff line change
Expand Up @@ -464,6 +464,14 @@ func (f *Factory) ConstructZeroValues() memo.RelExpr {
})
}

// ConstructNoColsRow constructs a Values operator with no columns and one row.
func (f *Factory) ConstructNoColsRow() memo.RelExpr {
return f.ConstructValues(memo.ScalarListWithEmptyTuple, &memo.ValuesPrivate{
Cols: opt.ColList{},
ID: f.Metadata().NextUniqueID(),
})
}

// ConstructJoin constructs the join operator that corresponds to the given join
// operator type.
func (f *Factory) ConstructJoin(
Expand Down
5 changes: 1 addition & 4 deletions pkg/sql/opt/norm/factory_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -50,10 +50,7 @@ func TestSimplifyFilters(t *testing.T) {
eq := f.ConstructEq(variable, constant)

// Filters expression evaluates to False if any operand is False.
vals := f.ConstructValues(memo.ScalarListWithEmptyTuple, &memo.ValuesPrivate{
Cols: opt.ColList{},
ID: f.Metadata().NextUniqueID(),
})
vals := f.ConstructNoColsRow()
filters := memo.FiltersExpr{
f.ConstructFiltersItem(eq),
f.ConstructFiltersItem(memo.FalseSingleton),
Expand Down
5 changes: 1 addition & 4 deletions pkg/sql/opt/optbuilder/insert.go
Original file line number Diff line number Diff line change
Expand Up @@ -566,10 +566,7 @@ func (mb *mutationBuilder) buildInputForInsert(inScope *scope, inputRows *tree.S
// Handle DEFAULT VALUES case by creating a single empty row as input.
if inputRows == nil {
mb.outScope = inScope.push()
mb.outScope.expr = mb.b.factory.ConstructValues(memo.ScalarListWithEmptyTuple, &memo.ValuesPrivate{
Cols: opt.ColList{},
ID: mb.md.NextUniqueID(),
})
mb.outScope.expr = mb.b.factory.ConstructNoColsRow()
mb.inputForInsertExpr = mb.outScope.expr
return
}
Expand Down
17 changes: 12 additions & 5 deletions pkg/sql/opt/optbuilder/plpgsql.go
Original file line number Diff line number Diff line change
Expand Up @@ -507,13 +507,23 @@ func (b *plpgsqlBuilder) buildPLpgSQLStatements(stmts []ast.Statement, s *scope)
retCon := b.makeContinuation("_stmt_exec_ret")
b.appendPlpgSQLStmts(&retCon, stmts[i+1:])

// We only need the first row from the SQL statement.
// Ensure that the SQL statement returns at most one row.
stmtScope.expr = b.ob.factory.ConstructLimit(
stmtScope.expr,
b.ob.factory.ConstructConst(tree.NewDInt(tree.DInt(1)), types.Int),
stmtScope.makeOrderingChoice(),
)

// Ensure that the SQL statement returns at least one row. The RIGHT join
// ensures that when the SQL statement returns no rows, it is extended
// with a single row of NULL values.
stmtScope.expr = b.ob.factory.ConstructRightJoin(
stmtScope.expr,
b.ob.factory.ConstructNoColsRow(),
nil, /* on */
memo.EmptyJoinPrivate,
)

// Step 2: build the INTO statement into a continuation routine that calls
// the previously built continuation.
intoScope := b.buildInto(stmtScope, t.Target)
Expand Down Expand Up @@ -1384,10 +1394,7 @@ func (b *plpgsqlBuilder) resolveVariableForAssign(name tree.Name) *types.T {

func (b *plpgsqlBuilder) ensureScopeHasExpr(s *scope) {
if s.expr == nil {
s.expr = b.ob.factory.ConstructValues(memo.ScalarListWithEmptyTuple, &memo.ValuesPrivate{
Cols: opt.ColList{},
ID: b.ob.factory.Metadata().NextUniqueID(),
})
s.expr = b.ob.factory.ConstructNoColsRow()
}
}

Expand Down
5 changes: 1 addition & 4 deletions pkg/sql/opt/optbuilder/select.go
Original file line number Diff line number Diff line change
Expand Up @@ -1289,10 +1289,7 @@ func (b *Builder) buildFrom(
outScope = b.buildFromTables(from.Tables, lockCtx, inScope)
} else {
outScope = inScope.push()
outScope.expr = b.factory.ConstructValues(memo.ScalarListWithEmptyTuple, &memo.ValuesPrivate{
Cols: opt.ColList{},
ID: b.factory.Metadata().NextUniqueID(),
})
outScope.expr = b.factory.ConstructNoColsRow()
}

return outScope
Expand Down
5 changes: 1 addition & 4 deletions pkg/sql/opt/optbuilder/srfs.go
Original file line number Diff line number Diff line change
Expand Up @@ -177,10 +177,7 @@ func (b *Builder) buildZip(exprs tree.Exprs, inScope *scope) (outScope *scope) {
}

// Construct the zip as a ProjectSet with empty input.
input := b.factory.ConstructValues(memo.ScalarListWithEmptyTuple, &memo.ValuesPrivate{
Cols: opt.ColList{},
ID: b.factory.Metadata().NextUniqueID(),
})
input := b.factory.ConstructNoColsRow()
outScope.expr = b.factory.ConstructProjectSet(input, zip)
if len(outScope.cols) == 1 {
outScope.singleSRFColumn = true
Expand Down
37 changes: 21 additions & 16 deletions pkg/sql/opt/optbuilder/testdata/procedure_plpgsql
Original file line number Diff line number Diff line change
Expand Up @@ -55,23 +55,28 @@ call
│ └── project
│ ├── columns: "_stmt_exec_ret_2":58
│ ├── project
│ │ ├── columns: c:57!null
│ │ ├── limit
│ │ │ ├── columns: count_rows:14!null
│ │ │ ├── scalar-group-by
│ │ ├── columns: c:57
│ │ ├── right-join (cross)
│ │ │ ├── columns: count_rows:14
│ │ │ ├── limit
│ │ │ │ ├── columns: count_rows:14!null
│ │ │ │ ├── project
│ │ │ │ │ └── select
│ │ │ │ │ ├── columns: k:9!null i:10 s:11 crdb_internal_mvcc_timestamp:12 tableoid:13
│ │ │ │ │ ├── scan t
│ │ │ │ │ │ └── columns: k:9!null i:10 s:11 crdb_internal_mvcc_timestamp:12 tableoid:13
│ │ │ │ │ └── filters
│ │ │ │ │ └── eq
│ │ │ │ │ ├── variable: k:9
│ │ │ │ │ └── variable: arg_k:6
│ │ │ │ └── aggregations
│ │ │ │ └── count-rows [as=count_rows:14]
│ │ │ └── const: 1
│ │ │ │ ├── scalar-group-by
│ │ │ │ │ ├── columns: count_rows:14!null
│ │ │ │ │ ├── project
│ │ │ │ │ │ └── select
│ │ │ │ │ │ ├── columns: k:9!null i:10 s:11 crdb_internal_mvcc_timestamp:12 tableoid:13
│ │ │ │ │ │ ├── scan t
│ │ │ │ │ │ │ └── columns: k:9!null i:10 s:11 crdb_internal_mvcc_timestamp:12 tableoid:13
│ │ │ │ │ │ └── filters
│ │ │ │ │ │ └── eq
│ │ │ │ │ │ ├── variable: k:9
│ │ │ │ │ │ └── variable: arg_k:6
│ │ │ │ │ └── aggregations
│ │ │ │ │ └── count-rows [as=count_rows:14]
│ │ │ │ └── const: 1
│ │ │ ├── values
│ │ │ │ └── tuple
│ │ │ └── filters (true)
│ │ └── projections
│ │ └── variable: count_rows:14 [as=c:57]
│ └── projections
Expand Down