Merge #122657

122657: plpgsql: disallow txn control in udfs and nested procedures r=DrewKimball a=DrewKimball

This patch adds validation for transaction control statements, so that attempting to use them from within a UDF results in an error. In addition, attempting to use a transaction control statement from a nested stored procedure call will now result in an "unimplemented" error.

Informs #122266

Release note: None

Co-authored-by: Drew Kimball <drewk@cockroachlabs.com>

pkg/ccl/logictestccl/testdata/logic_test/plpgsql_txn

@@ -983,4 +983,30 @@ CREATE PROCEDURE p() LANGUAGE PLpgSQL AS $$
   END
 $$;
 
+# Regression test for #122266 - functions should not be allowed to use
+# COMMIT/ROLLBACK, and COMMIT/ROLLBACK is currently unsupported in a nested
+# CALL statement.
+subtest commit_rollback
+
+statement error pgcode 2D000 pq: invalid transaction termination
+CREATE FUNCTION f() RETURNS INT LANGUAGE PLpgSQL AS $$ BEGIN COMMIT; END $$;
+
+statement error pgcode 2D000 pq: invalid transaction termination
+CREATE FUNCTION f() RETURNS INT LANGUAGE PLpgSQL AS $$ BEGIN ROLLBACK; END $$;
+
+statement error pgcode 2D000 pq: invalid transaction termination
+CREATE FUNCTION f() RETURNS INT LANGUAGE PLpgSQL AS $$
+  DECLARE i INT := 0; BEGIN WHILE i < 10 LOOP COMMIT; i := i + 1; END LOOP; END
+$$;
+
+statement ok
+CREATE PROCEDURE p_nested_commit() LANGUAGE PLpgSQL AS $$ BEGIN COMMIT; END $$;
+CREATE PROCEDURE p_nested_rollback() LANGUAGE PLpgSQL AS $$ BEGIN ROLLBACK; END $$;
+
+statement error pgcode 0A000 pq: unimplemented: transaction control statements in nested routines
+CREATE PROCEDURE p() LANGUAGE PLpgSQL AS $$ BEGIN CALL p_nested_commit(); END $$;
+
+statement error pgcode 0A000 pq: unimplemented: transaction control statements in nested routines
+CREATE PROCEDURE p() LANGUAGE PLpgSQL AS $$ BEGIN CALL p_nested_rollback(); END $$;
+
 subtest end

pkg/sql/opt/optbuilder/builder.go

@@ -153,6 +153,10 @@ type Builder struct {
 	// insideDataSource is true when we are processing a data source.
 	insideDataSource bool
 
+	// insideNestedPLpgSQLCall is true when we are processing a nested PLpgSQL
+	// CALL statement.
+	insideNestedPLpgSQLCall bool
+
 	// If set, we are collecting view dependencies in schemaDeps. This can only
 	// happen inside view/function definitions.
 	//

pkg/sql/opt/optbuilder/create_function.go

@@ -26,6 +26,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/plpgsql"
 	plpgsqlparser "github.com/cockroachdb/cockroach/pkg/sql/plpgsql/parser"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/cast"
+	"github.com/cockroachdb/cockroach/pkg/sql/sem/plpgsqltree"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
 	"github.com/cockroachdb/cockroach/pkg/util/errorutil/unimplemented"
@@ -376,6 +377,18 @@ func (b *Builder) buildCreateFunction(cf *tree.CreateRoutine, inScope *scope) (o
 			panic(err)
 		}
 
+		// Check for transaction control statements in UDFs.
+		if !cf.IsProcedure {
+			var tc transactionControlVisitor
+			plpgsqltree.Walk(&tc, stmt.AST)
+			if tc.foundTxnControlStatement {
+				panic(errors.WithDetailf(
+					pgerror.Newf(pgcode.InvalidTransactionTermination, "invalid transaction termination"),
+					"transaction control statements are only allowed in procedures",
+				))
+			}
+		}
+
 		// We need to disable stable function folding because we want to catch the
 		// volatility of stable functions. If folded, we only get a scalar and lose
 		// the volatility.

pkg/sql/opt/optbuilder/plpgsql.go

@@ -285,6 +285,15 @@ func (b *plpgsqlBuilder) buildRootBlock(
 		var tc transactionControlVisitor
 		ast.Walk(&tc, astBlock)
 		if tc.foundTxnControlStatement {
+			if b.ob.insideNestedPLpgSQLCall {
+				// Disallow transaction control statements in nested routines for now.
+				// TODO(#122266): once we support this, make sure to validate that
+				// transaction control statements are only allowed in a nested procedure
+				// when all ancestors are procedures or DO blocks.
+				panic(unimplemented.NewWithIssue(122266,
+					"transaction control statements in nested routines",
+				))
+			}
 			// Disable stable folding, since different parts of the routine can be run
 			// in different transactions.
 			b.ob.factory.FoldingControl().TemporarilyDisallowStableFolds(func() {
@@ -939,8 +948,9 @@ func (b *plpgsqlBuilder) buildPLpgSQLStatements(stmts []ast.Statement, s *scope)
 			procTyp := proc.ResolvedType()
 			colName := scopeColName("").WithMetadataName(b.makeIdentifier("stmt_call"))
 			col := b.ob.synthesizeColumn(callScope, colName, procTyp, nil /* expr */, nil /* scalar */)
-			procScalar := b.ob.buildRoutine(proc, def, callCon.s, callScope, b.colRefs)
-			col.scalar = procScalar
+			b.ob.withinNestedPLpgSQLCall(func() {
+				col.scalar = b.ob.buildRoutine(proc, def, callCon.s, callScope, b.colRefs)
+			})
 			b.ob.constructProjectForScope(callCon.s, callScope)
 
 			// Collect any target variables in OUT-parameter position. The result of

pkg/sql/opt/optbuilder/routine.go

@@ -677,3 +677,11 @@ func (b *Builder) maybeAddRoutineAssignmentCasts(
 	}
 	return b.constructProject(expr, stmtScope.cols), stmtScope.makePhysicalProps()
 }
+
+func (b *Builder) withinNestedPLpgSQLCall(fn func()) {
+	defer func(origValue bool) {
+		b.insideNestedPLpgSQLCall = origValue
+	}(b.insideNestedPLpgSQLCall)
+	b.insideNestedPLpgSQLCall = true
+	fn()
+}