plpgsql: add support for COMMIT and ROLLBACK statements

pkg/ccl/logictestccl/tests/fakedist-disk/BUILD.bazel

@@ -12,7 +12,7 @@ go_test(
         "//build/toolchains:is_heavy": {"test.Pool": "heavy"},
         "//conditions:default": {"test.Pool": "large"},
     }),
-    shard_count = 19,
+    shard_count = 20,
     tags = [
         "ccl_test",
         "cpu:2",

pkg/ccl/logictestccl/tests/fakedist-vec-off/BUILD.bazel

@@ -12,7 +12,7 @@ go_test(
         "//build/toolchains:is_heavy": {"test.Pool": "heavy"},
         "//conditions:default": {"test.Pool": "large"},
     }),
-    shard_count = 19,
+    shard_count = 20,
     tags = [
         "ccl_test",
         "cpu:2",

pkg/ccl/logictestccl/tests/fakedist/BUILD.bazel

@@ -12,7 +12,7 @@ go_test(
         "//build/toolchains:is_heavy": {"test.Pool": "heavy"},
         "//conditions:default": {"test.Pool": "large"},
     }),
-    shard_count = 20,
+    shard_count = 21,
     tags = [
         "ccl_test",
         "cpu:2",

pkg/ccl/logictestccl/tests/local-legacy-schema-changer/BUILD.bazel

@@ -9,7 +9,7 @@ go_test(
         "//pkg/ccl/logictestccl:testdata",  # keep
     ],
     exec_properties = {"test.Pool": "large"},
-    shard_count = 19,
+    shard_count = 20,
     tags = [
         "ccl_test",
         "cpu:1",

pkg/ccl/logictestccl/tests/local-mixed-23.2/BUILD.bazel

@@ -9,7 +9,7 @@ go_test(
         "//pkg/ccl/logictestccl:testdata",  # keep
     ],
     exec_properties = {"test.Pool": "large"},
-    shard_count = 18,
+    shard_count = 19,
     tags = [
         "ccl_test",
         "cpu:1",

pkg/ccl/logictestccl/tests/local-read-committed/BUILD.bazel

@@ -9,7 +9,7 @@ go_test(
         "//pkg/ccl/logictestccl:testdata",  # keep
     ],
     exec_properties = {"test.Pool": "large"},
-    shard_count = 19,
+    shard_count = 20,
     tags = [
         "ccl_test",
         "cpu:1",

pkg/ccl/logictestccl/tests/local-vec-off/BUILD.bazel

@@ -9,7 +9,7 @@ go_test(
         "//pkg/ccl/logictestccl:testdata",  # keep
     ],
     exec_properties = {"test.Pool": "large"},
-    shard_count = 19,
+    shard_count = 20,
     tags = [
         "ccl_test",
         "cpu:1",

pkg/ccl/logictestccl/tests/local/BUILD.bazel

@@ -9,7 +9,7 @@ go_test(
         "//pkg/ccl/logictestccl:testdata",  # keep
     ],
     exec_properties = {"test.Pool": "large"},
-    shard_count = 34,
+    shard_count = 35,
     tags = [
         "ccl_test",
         "cpu:1",

pkg/sql/conn_executor.go

@@ -48,6 +48,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/execstats"
 	"github.com/cockroachdb/cockroach/pkg/sql/idxrecommendations"
 	"github.com/cockroachdb/cockroach/pkg/sql/idxusage"
+	"github.com/cockroachdb/cockroach/pkg/sql/opt/memo"
 	"github.com/cockroachdb/cockroach/pkg/sql/parser"
 	"github.com/cockroachdb/cockroach/pkg/sql/parser/statements"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
@@ -1507,6 +1508,26 @@ type connExecutor struct {
 			implicit bool
 		}
 
+		// storedProcTxnState contains fields that are used for explicit transaction
+		// management in stored procedures. Both fields are set during execution of
+		// a stored procedure through the planner. Both fields are reset by the
+		// connExecutor in execCmd, depending on the execution status (e.g. success,
+		// error, retry).
+		storedProcTxnState struct {
+			// txnOp, if not set to StoredProcTxnNoOp, indicates that the current txn
+			// should be committed or aborted. The connExecutor uses it in
+			// execStmtInOpenState to decide whether to commit or rollback the current
+			// transaction.
+			txnOp tree.StoredProcTxnOp
+
+			// resumeProc, if non-nil, contains the plan for a CALL statement that
+			// will continue execution of a stored procedure that previously suspended
+			// execution in order to commit or abort its transaction. The planner
+			// checks for resumeProc in buildExecMemo, and executes it as the next
+			// statement if it is non-nil.
+			resumeProc *memo.Memo
+		}
+
 		// shouldExecuteOnTxnRestart indicates that ex.onTxnRestart will be
 		// called when txn is being retried. It is true when txn is started but
 		// can remain false when txn is executed within another higher-level
@@ -2592,6 +2613,26 @@ func (ex *connExecutor) execCmd() (retErr error) {
 		panic(errors.AssertionFailedf("unexpected advance code: %s", advInfo.code))
 	}
 
+	// Special handling for COMMIT/ROLLBACK in PL/pgSQL stored procedures. We
+	// unconditionally reset the StoredProcTxnOp because it has either
+	// successfully set in motion the commit/rollback of the current transaction,
+	// or execution failed and the PL/pgSQL command should be ignored.
+	ex.extraTxnState.storedProcTxnState.txnOp = tree.StoredProcTxnNoOp
+	switch advInfo.code {
+	case stayInPlace, rewind:
+		// Do not reset the "resume" plan. This will allow a sub-transaction
+		// within a stored procedure to be retried individually from any previous or
+		// following transactions within the stored procedure.
+		//
+		// NOTE: we will eventually reach the default case below when the retries
+		// terminate.
+	default:
+		// Finish cleaning up the stored proc state by resetting the "resume" plan.
+		// The stored procedure has finished execution, either successfully or with
+		// an error.
+		ex.extraTxnState.storedProcTxnState.resumeProc = nil
+	}
+
 	if err := ex.updateTxnRewindPosMaybe(ctx, cmd, pos, advInfo); err != nil {
 		return err
 	}
@@ -3734,6 +3775,7 @@ func (ex *connExecutor) initPlanner(ctx context.Context, p *planner) {
 	p.noticeSender = nil
 	p.preparedStatements = ex.getPrepStmtsAccessor()
 	p.sqlCursors = ex.getCursorAccessor()
+	p.storedProcTxnState = ex.getStoredProcTxnStateAccessor()
 	p.createdSequences = ex.getCreatedSequencesAccessor()
 
 	p.queryCacheSession.Init()
@@ -4228,6 +4270,12 @@ func (ex *connExecutor) getCursorAccessor() sqlCursors {
 	}
 }
 
+func (ex *connExecutor) getStoredProcTxnStateAccessor() storedProcTxnStateAccessor {
+	return storedProcTxnStateAccessor{
+		ex: ex,
+	}
+}
+
 func (ex *connExecutor) getCreatedSequencesAccessor() createdSequences {
 	return connExCreatedSequencesAccessor{
 		ex: ex,

pkg/sql/conn_executor_exec.go

@@ -977,7 +977,8 @@ func (ex *connExecutor) execStmtInOpenState(
 	}
 
 	// For regular statements (the ones that get to this point), we
-	// don't return any event unless an error happens.
+	// don't return any event unless an error happens, or a CALL statement
+	// performs a nested transaction COMMIT or ROLLBACK.
 
 	// For a portal (prepared stmt), since handleAOST() is called when preparing
 	// the statement, and this function is idempotent, we don't need to
@@ -1240,6 +1241,29 @@ func (ex *connExecutor) execStmtInOpenState(
 		log.VEventf(ctx, 2, "push detected for non-refreshable txn but auto-retry not possible")
 	}
 
+	// Special handling for explicit transaction management in CALL statements.
+	// A stored procedure has executed a COMMIT or ROLLBACK statement and
+	// suspended its execution. Direct the connExecutor to commit/rollback the
+	// current transaction, and then resume execution within the new transaction.
+	switch ex.extraTxnState.storedProcTxnState.txnOp {
+	case tree.StoredProcTxnCommit:
+		// Commit the current transaction. The connExecutor will open a new
+		// transaction, and then return to executing the same CALL statement.
+		ev, payload := ex.commitSQLTransaction(ctx, ast, ex.commitSQLTransactionInternal)
+		if payload != nil {
+			return ev, payload, nil
+		}
+		return eventTxnFinishCommittedPLpgSQL{}, nil, nil
+	case tree.StoredProcTxnRollback:
+		// Abort the current transaction. The connExecutor will open a new
+		// transaction, and then return to executing the same CALL statement.
+		ev, payload := ex.rollbackSQLTransaction(ctx, ast)
+		if payload != nil {
+			return ev, payload, nil
+		}
+		return eventTxnFinishAbortedPLpgSQL{}, nil, nil
+	}
+
 	// No event was generated.
 	return nil, nil, nil
 }