Merge #147923

147923: sql: do not attempt to re-execute portal after PL/pgSQL txn control r=yuzefovich,michae2 a=DrewKimball

Previously, the PL/pgSQL transaction control statements (COMMIT/ROLLBACK) did not work with the extended wire protocol. If the `CALL` statement was executed via portal, we would attempt to re-execute the portal after the original transaction ended. This would result in an error like `unknown portal ""`.

This commit fixes the bug by replacing the original command with a dummy `ExecStmt` command when resuming stored proc execution in a new transaction. This allows the portal to be cleaned up with the first transaction without attempting to resolve it after the fact. The fix is controled by the session var `use_proc_txn_control_extended_protocol_fix`, which is on by default.

Informs #147701

Release note (bug fix): Fixed a bug that would cause a CALL statement executed via a portal in the extended wire protocol to result in an error like `unknown portal ""` if the stored procedure contained `COMMIT` or `ROLLBACK` statements. The bug has existed since PL/pgSQL transaction control statements were introduced in v24.1.

Co-authored-by: Drew Kimball <[email protected]>

Author: craig[bot]

pkg/sql/conn_executor.go

@@ -1651,6 +1651,12 @@ type connExecutor struct {
 			// checks for resumeProc in buildExecMemo, and executes it as the next
 			// statement if it is non-nil.
 			resumeProc *memo.Memo
+
+			// resumeStmt is set if resumeProc is set. If the stored procedure was
+			// executed via a portal in the extended wire protocol, this statement
+			// will be used to synthesize an ExecStmt command to avoid attempting to
+			// execute the portal twice.
+			resumeStmt statements.Statement[tree.Statement]
 		}
 
 		// shouldExecuteOnTxnRestart indicates that ex.onTxnRestart will be
@@ -2300,6 +2306,15 @@ func (ex *connExecutor) execCmd() (retErr error) {
 		return err // err could be io.EOF
 	}
 
+	// Special handling for COMMIT/ROLLBACK in PL/pgSQL stored procedures. See the
+	// makeCmdForStoredProcResume comment for details.
+	if ex.extraTxnState.storedProcTxnState.resumeProc != nil {
+		cmd, err = ex.makeCmdForStoredProcResume(cmd)
+		if err != nil {
+			return err
+		}
+	}
+
 	if log.ExpensiveLogEnabled(ctx, 2) {
 		ex.sessionEventf(ctx, "[%s pos:%d] executing %s",
 			ex.machine.CurState(), pos, cmd)
@@ -2716,6 +2731,7 @@ func (ex *connExecutor) execCmd() (retErr error) {
 		// and transaction modes. The stored procedure has finished execution,
 		// either successfully or with an error.
 		ex.extraTxnState.storedProcTxnState.resumeProc = nil
+		ex.extraTxnState.storedProcTxnState.resumeStmt = statements.Statement[tree.Statement]{}
 		ex.extraTxnState.storedProcTxnState.txnModes = nil
 	}
 
@@ -2892,6 +2908,35 @@ func stateToTxnStatusIndicator(s fsm.State) TransactionStatusIndicator {
 	}
 }
 
+// makeCmdForStoredProcResume creates a Command that can be used to resume
+// execution of a stored procedure that has ended the previous transaction via
+// COMMIT or ROLLBACK. It is not enough to just process the original command
+// again, because it may have been an ExecPortal statement, and the portal will
+// already have been closed after the first phase of execution.
+func (ex *connExecutor) makeCmdForStoredProcResume(curCmd Command) (Command, error) {
+	if !ex.sessionData().UseProcTxnControlExtendedProtocolFix {
+		// The fix is not enabled, so return the original command.
+		return curCmd, nil
+	}
+	var timeReceived crtime.Mono
+	switch t := curCmd.(type) {
+	case ExecStmt:
+		// NOTE: it is not strictly necessary to replace ExecStmt. However, it seems
+		// best to handle the commands in a consistent way.
+		timeReceived = t.TimeReceived
+	case ExecPortal:
+		timeReceived = t.TimeReceived
+	default:
+		return nil, errors.AssertionFailedf(
+			"unexpected command type %T for stored procedure resume", t,
+		)
+	}
+	return ExecStmt{
+		Statement:    ex.extraTxnState.storedProcTxnState.resumeStmt,
+		TimeReceived: timeReceived,
+	}, nil
+}
+
 // isCopyToExternalStorage returns true if the CopyIn command is writing to an
 // ExternalStorage such as nodelocal or userfile. It does so by checking the
 // target table/schema names against the sentinel, internal table/schema names

pkg/sql/conn_io.go

@@ -143,6 +143,8 @@ type ExecStmt struct {
 
 	// LastInBatch indicates if this command contains the last query in a
 	// simple protocol Query message that contains a batch of 1 or more queries.
+	// This is used to determine whether autocommit can be applied to the
+	// transaction, and need not be set for correctness.
 	LastInBatch bool
 	// LastInBatchBeforeShowCommitTimestamp indicates that this command contains
 	// the second-to-last query in a simple protocol Query message that contains
@@ -151,7 +153,9 @@ type ExecStmt struct {
 	// such that the SHOW COMMIT TIMESTAMP statement can return the timestamp of
 	// the transaction which applied to all the other statements in the batch.
 	// Note that SHOW COMMIT TIMESTAMP is not permitted in any other position in
-	// such a multi-statement implicit transaction.
+	// such a multi-statement implicit transaction. This is used to determine
+	// whether autocommit can be applied to the transaction, and need not be set
+	// for correctness.
 	LastInBatchBeforeShowCommitTimestamp bool
 }
 

pkg/sql/exec_util.go

@@ -4221,6 +4221,10 @@ func (m *sessionDataMutator) SetDistSQLUseReducedLeafWriteSets(val bool) {
 	m.data.DistSQLUseReducedLeafWriteSets = val
 }
 
+func (m *sessionDataMutator) SetUseProcTxnControlExtendedProtocolFix(val bool) {
+	m.data.UseProcTxnControlExtendedProtocolFix = val
+}
+
 // Utility functions related to scrubbing sensitive information on SQL Stats.
 
 // quantizeCounts ensures that the Count field in the

pkg/sql/logictest/testdata/logic_test/information_schema

@@ -4137,6 +4137,7 @@ unsafe_allow_triggers_modifying_cascades                         off
 use_cputs_on_non_unique_indexes                                  off
 use_improved_routine_dependency_tracking                         on
 use_pre_25_2_variadic_builtins                                   off
+use_proc_txn_control_extended_protocol_fix                       on
 variable_inequality_lookup_join_enabled                          on
 vector_search_beam_size                                          32
 xmloption                                                        content

pkg/sql/logictest/testdata/logic_test/pg_catalog

@@ -3132,6 +3132,7 @@ use_cputs_on_non_unique_indexes                                  off
 use_declarative_schema_changer                                   on                  NULL      NULL        NULL        string
 use_improved_routine_dependency_tracking                         on                  NULL      NULL        NULL        string
 use_pre_25_2_variadic_builtins                                   off                 NULL      NULL        NULL        string
+use_proc_txn_control_extended_protocol_fix                       on                  NULL      NULL        NULL        string
 variable_inequality_lookup_join_enabled                          on                  NULL      NULL        NULL        string
 vector_search_beam_size                                          32                  NULL      NULL        NULL        string
 vectorize                                                        on                  NULL      NULL        NULL        string
@@ -3366,6 +3367,7 @@ use_cputs_on_non_unique_indexes                                  off
 use_declarative_schema_changer                                   on                  NULL  user     NULL      on                  on
 use_improved_routine_dependency_tracking                         on                  NULL  user     NULL      on                  on
 use_pre_25_2_variadic_builtins                                   off                 NULL  user     NULL      off                 off
+use_proc_txn_control_extended_protocol_fix                       on                  NULL  user     NULL      on                  on
 variable_inequality_lookup_join_enabled                          on                  NULL  user     NULL      on                  on
 vector_search_beam_size                                          32                  NULL  user     NULL      32                  32
 vectorize                                                        on                  NULL  user     NULL      on                  on
@@ -3592,6 +3594,7 @@ use_cputs_on_non_unique_indexes                                  NULL    NULL
 use_declarative_schema_changer                                   NULL    NULL     NULL     NULL        NULL
 use_improved_routine_dependency_tracking                         NULL    NULL     NULL     NULL        NULL
 use_pre_25_2_variadic_builtins                                   NULL    NULL     NULL     NULL        NULL
+use_proc_txn_control_extended_protocol_fix                       NULL    NULL     NULL     NULL        NULL
 variable_inequality_lookup_join_enabled                          NULL    NULL     NULL     NULL        NULL
 vector_search_beam_size                                          NULL    NULL     NULL     NULL        NULL
 vectorize                                                        NULL    NULL     NULL     NULL        NULL

pkg/sql/logictest/testdata/logic_test/show_source

@@ -246,6 +246,7 @@ use_cputs_on_non_unique_indexes                                  off
 use_declarative_schema_changer                                   on
 use_improved_routine_dependency_tracking                         on
 use_pre_25_2_variadic_builtins                                   off
+use_proc_txn_control_extended_protocol_fix                       on
 variable_inequality_lookup_join_enabled                          on
 vector_search_beam_size                                          32
 vectorize                                                        on

pkg/sql/pgwire/testdata/pgtest/procedure

@@ -46,3 +46,51 @@ ReadyForQuery
 {"Type":"ReadyForQuery","TxStatus":"I"}
 {"Type":"CommandComplete","CommandTag":"DROP TABLE"}
 {"Type":"ReadyForQuery","TxStatus":"I"}
+
+# Regression test for #147701: correctly handle a PL/pgSQL procedure that
+# commits or rolls back the transaction.
+send
+Query {"String": "CREATE OR REPLACE PROCEDURE p() LANGUAGE PLpgSQL AS $$ BEGIN RAISE NOTICE 'foo'; COMMIT; RAISE NOTICE 'bar'; ROLLBACK; RAISE NOTICE 'baz'; END $$;"}
+----
+
+until
+ReadyForQuery
+----
+{"Type":"CommandComplete","CommandTag":"CREATE PROCEDURE"}
+{"Type":"ReadyForQuery","TxStatus":"I"}
+
+send
+Query {"String": "CALL p()"}
+----
+
+until
+ReadyForQuery
+----
+{"Type":"RowDescription","Fields":null}
+{"Severity":"NOTICE","SeverityUnlocalized":"NOTICE","Code":"00000","Message":"foo","Detail":"","Hint":"","Position":0,"InternalPosition":0,"InternalQuery":"","Where":"","SchemaName":"","TableName":"","ColumnName":"","DataTypeName":"","ConstraintName":"","File":"builtins.go","Line":0,"Routine":"func378","UnknownFields":null}
+{"Type":"RowDescription","Fields":null}
+{"Severity":"NOTICE","SeverityUnlocalized":"NOTICE","Code":"00000","Message":"bar","Detail":"","Hint":"","Position":0,"InternalPosition":0,"InternalQuery":"","Where":"","SchemaName":"","TableName":"","ColumnName":"","DataTypeName":"","ConstraintName":"","File":"builtins.go","Line":0,"Routine":"func378","UnknownFields":null}
+{"Type":"RowDescription","Fields":null}
+{"Severity":"NOTICE","SeverityUnlocalized":"NOTICE","Code":"00000","Message":"baz","Detail":"","Hint":"","Position":0,"InternalPosition":0,"InternalQuery":"","Where":"","SchemaName":"","TableName":"","ColumnName":"","DataTypeName":"","ConstraintName":"","File":"builtins.go","Line":0,"Routine":"func378","UnknownFields":null}
+{"Type":"CommandComplete","CommandTag":"CALL"}
+{"Type":"ReadyForQuery","TxStatus":"I"}
+
+send
+Parse {"Name": "foo", "Query": "CALL p()"}
+Bind {"DestinationPortal": "foo", "PreparedStatement": "foo"}
+Execute {"Portal": "foo"}
+Sync
+----
+
+until
+ReadyForQuery
+----
+{"Type":"ParseComplete"}
+{"Type":"BindComplete"}
+{"Severity":"NOTICE","SeverityUnlocalized":"NOTICE","Code":"00000","Message":"foo","Detail":"","Hint":"","Position":0,"InternalPosition":0,"InternalQuery":"","Where":"","SchemaName":"","TableName":"","ColumnName":"","DataTypeName":"","ConstraintName":"","File":"builtins.go","Line":0,"Routine":"func378","UnknownFields":null}
+{"Type":"RowDescription","Fields":null}
+{"Severity":"NOTICE","SeverityUnlocalized":"NOTICE","Code":"00000","Message":"bar","Detail":"","Hint":"","Position":0,"InternalPosition":0,"InternalQuery":"","Where":"","SchemaName":"","TableName":"","ColumnName":"","DataTypeName":"","ConstraintName":"","File":"builtins.go","Line":0,"Routine":"func378","UnknownFields":null}
+{"Type":"RowDescription","Fields":null}
+{"Severity":"NOTICE","SeverityUnlocalized":"NOTICE","Code":"00000","Message":"baz","Detail":"","Hint":"","Position":0,"InternalPosition":0,"InternalQuery":"","Where":"","SchemaName":"","TableName":"","ColumnName":"","DataTypeName":"","ConstraintName":"","File":"builtins.go","Line":0,"Routine":"func378","UnknownFields":null}
+{"Type":"CommandComplete","CommandTag":"CALL"}
+{"Type":"ReadyForQuery","TxStatus":"I"}

pkg/sql/routine.go

@@ -14,6 +14,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/colinfo"
 	"github.com/cockroachdb/cockroach/pkg/sql/isql"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/memo"
+	"github.com/cockroachdb/cockroach/pkg/sql/parser/statements"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/eval"
@@ -660,14 +661,18 @@ type storedProcTxnStateAccessor struct {
 }
 
 func (a *storedProcTxnStateAccessor) setStoredProcTxnState(
-	txnOp tree.StoredProcTxnOp, txnModes *tree.TransactionModes, resumeProc *memo.Memo,
+	txnOp tree.StoredProcTxnOp,
+	txnModes *tree.TransactionModes,
+	resumeProc *memo.Memo,
+	resumeStmt statements.Statement[tree.Statement],
 ) {
 	if a.ex == nil {
 		panic(errors.AssertionFailedf("setStoredProcTxnState is not supported without connExecutor"))
 	}
 	a.ex.extraTxnState.storedProcTxnState.txnOp = txnOp
 	a.ex.extraTxnState.storedProcTxnState.txnModes = txnModes
 	a.ex.extraTxnState.storedProcTxnState.resumeProc = resumeProc
+	a.ex.extraTxnState.storedProcTxnState.resumeStmt = resumeStmt
 }
 
 func (a *storedProcTxnStateAccessor) getTxnOp() tree.StoredProcTxnOp {
@@ -713,6 +718,8 @@ func (p *planner) EvalTxnControlExpr(
 	if err != nil {
 		return nil, err
 	}
-	p.storedProcTxnState.setStoredProcTxnState(expr.Op, &expr.Modes, resumeProc.(*memo.Memo))
+	p.storedProcTxnState.setStoredProcTxnState(
+		expr.Op, &expr.Modes, resumeProc.(*memo.Memo), p.stmt.Statement,
+	)
 	return tree.DNull, nil
 }

pkg/sql/sessiondatapb/local_only_session_data.proto

@@ -685,6 +685,10 @@ message LocalOnlySessionData {
   // DistSQLUseReducedLeafWriteSets, when true, indicates that the DistSQL
   // runner should use the reduced write sets when constructing LeafTxns.
   bool distsql_use_reduced_leaf_write_sets = 174 [(gogoproto.customname) = "DistSQLUseReducedLeafWriteSets"];
+  // UseProcTxnControlExtendedProtocolFix, when true, enables the fix for
+  // PL/pgSQL transaction control statements (COMMIT, ROLLBACK) when the stored
+  // procedure is executed via a portal in the extended wire protocol.
+  bool use_proc_txn_control_extended_protocol_fix = 175;
 
   ///////////////////////////////////////////////////////////////////////////
   // WARNING: consider whether a session parameter you're adding needs to  //

pkg/sql/vars.go

@@ -4143,6 +4143,23 @@ var varGen = map[string]sessionVar{
 		},
 		GlobalDefault: globalTrue,
 	},
+
+	// CockroachDB extension.
+	`use_proc_txn_control_extended_protocol_fix`: {
+		GetStringVal: makePostgresBoolGetStringValFn(`use_proc_txn_control_extended_protocol_fix`),
+		Set: func(_ context.Context, m sessionDataMutator, s string) error {
+			b, err := paramparse.ParseBoolVar("use_proc_txn_control_extended_protocol_fix", s)
+			if err != nil {
+				return err
+			}
+			m.SetUseProcTxnControlExtendedProtocolFix(b)
+			return nil
+		},
+		Get: func(evalCtx *extendedEvalContext, _ *kv.Txn) (string, error) {
+			return formatBoolAsPostgresSetting(evalCtx.SessionData().UseProcTxnControlExtendedProtocolFix), nil
+		},
+		GlobalDefault: globalTrue,
+	},
 }
 
 func ReplicationModeFromString(s string) (sessiondatapb.ReplicationMode, error) {