plpgsql: don't exit early when SELECT INTO returns no rows

Previously, a SELECT INTO statement that returned zero rows wouldn't call the continuation, since the continuation is called as a projection. This patch adds a RIGHT JOIN with a zero-column one-row Values operator to ensure that SELECT INTO statements are correctly null-extended when the underlying SQL statement returns no rows. See the postgres docs: ``` https://www.postgresql.org/docs/16/plpgsql-statements.html#PLPGSQL-STATEMENTS-SQL-ONEROW ``` Fixes #114826 Release note (bug fix): Fixed a bug that existed only in pre-release versions v23.2.0-beta.1 and v23.2.0-beta.2 which could cause PLpgSQL routines with SELECT INTO syntax to return early.
cockroachdb · Dec 5, 2023 · f83b13a · f83b13a
1 parent f7ec91b
commit f83b13a
Show file tree

Hide file tree

Showing 10 changed files with 372 additions and 67 deletions.
diff --git a/pkg/sql/logictest/testdata/logic_test/udf_plpgsql b/pkg/sql/logictest/testdata/logic_test/udf_plpgsql
@@ -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
@@ -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
diff --git a/pkg/sql/opt/norm/decorrelate_funcs.go b/pkg/sql/opt/norm/decorrelate_funcs.go
@@ -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

diff --git a/pkg/sql/opt/norm/factory.go b/pkg/sql/opt/norm/factory.go
@@ -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(

diff --git a/pkg/sql/opt/norm/factory_test.go b/pkg/sql/opt/norm/factory_test.go
@@ -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),

diff --git a/pkg/sql/opt/optbuilder/insert.go b/pkg/sql/opt/optbuilder/insert.go
@@ -572,10 +572,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
 	}

diff --git a/pkg/sql/opt/optbuilder/plpgsql.go b/pkg/sql/opt/optbuilder/plpgsql.go
@@ -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)
@@ -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()
 	}
 }
 

diff --git a/pkg/sql/opt/optbuilder/select.go b/pkg/sql/opt/optbuilder/select.go
@@ -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

diff --git a/pkg/sql/opt/optbuilder/srfs.go b/pkg/sql/opt/optbuilder/srfs.go
@@ -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

diff --git a/pkg/sql/opt/optbuilder/testdata/procedure_plpgsql b/pkg/sql/opt/optbuilder/testdata/procedure_plpgsql
@@ -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