plpgsql: implement rollback behavior for EXCEPTION blocks

When an exception occurs within the EXCEPTION block of a PLpgSQL routine,
all changes to database state are rolled back, but the updated values for
the variables are still visible to the exception handler. For example:
```
CREATE FUNCTION f() RETURNS INT AS $$
  DECLARE
    i INT;
  BEGIN
    INSERT INTO xy VALUES (-100, -100);
    i := 100;
    RETURN 1 // 0;
  EXCEPTION WHEN division_by_zero THEN
    RETURN i;
  END
$$ LANGUAGE PLpgSQL;
```
The above function would not modify `xy`, since the exception would
revert the insert. However, the function would return `100`, since the
variable assignment occurred before the error was thrown.

This patch implements the required rollback behavior through internal
savepoints; when the top-level routine of a PLpgSQL block begins
executing, it creates a savepoint (if it has an exception handler).
When the exception handler catches an error, it calls `RollbackSavepoint`
to revert any changes to database state within the scope of the block.

Note that some errors can leave the transaction in a poisoned state,
where immediately rolling back the nested transaction is not possible.
The transaction can be returned to a normal state in some cases
(as is done for the read committed retry loop), but this will be left
as a TODO in 23.2, tracked in cockroachdb#111446. Until this is fixed, catching
`40001` and `40003` errors is not permitted.

Fixes cockroachdb#105253

Release note (sql change): When a PLpgSQL exception handler catches an
error, it now rolls back any changes to database state that occurred
within the block. Exception blocks are not currently permitted to catch
`40001` and `40003` errors.

pkg/BUILD.bazel

@@ -1959,7 +1959,6 @@ GO_TARGETS = [
     "//pkg/sql/plpgsql/parser/lexbase:lexbase",
     "//pkg/sql/plpgsql/parser:parser_test",
     "//pkg/sql/plpgsql/parser:plpgparser",
-    "//pkg/sql/plpgsql:plpgsql",
     "//pkg/sql/privilege:privilege",
     "//pkg/sql/privilege:privilege_test",
     "//pkg/sql/protoreflect/test:test",

pkg/kv/kvclient/kvcoord/testdata/savepoints

@@ -22,6 +22,10 @@ get k
 ----
 "k" -> b
 
+can-use x
+----
+true
+
 release x
 ----
 2 <noignore>
@@ -71,6 +75,10 @@ get k
 ----
 "k" -> b
 
+can-use x
+----
+true
+
 rollback x
 ----
 2 [2-2]
@@ -118,13 +126,21 @@ savepoint y
 put k cr
 ----
 
+can-use y
+----
+true
+
 release y
 ----
 3 <noignore>
 
 put k dr
 ----
 
+can-use x
+----
+true
+
 rollback x
 ----
 4 [2-4]
@@ -168,6 +184,10 @@ savepoint x
 put a d2
 ----
 
+can-use x
+----
+true
+
 rollback x
 ----
 3 [3-3]
@@ -182,6 +202,10 @@ savepoint x
 put b 2
 ----
 
+can-use x
+----
+true
+
 rollback x
 ----
 5 [3-3][5-5]
@@ -248,10 +272,18 @@ savepoint x
 ----
 1 <noignore>
 
+can-use x
+----
+true
+
 rollback x
 ----
 1 <noignore>
 
+can-use x
+----
+true
+
 release x
 ----
 1 <noignore>
@@ -284,6 +316,10 @@ savepoint x
 put k da
 ----
 
+can-use x
+----
+true
+
 rollback x
 ----
 1 [1-1]
@@ -299,6 +335,10 @@ get k
 put k db
 ----
 
+can-use x
+----
+true
+
 rollback x
 ----
 2 [1-2]
@@ -338,10 +378,18 @@ reset
 txn error cleared
 txn id not changed
 
+can-use x
+----
+true
+
 release x
 ----
 0 <noignore>
 
+can-use x
+----
+true
+
 rollback x
 ----
 0 <noignore>
@@ -365,6 +413,10 @@ cput k v bogus_expected
 ----
 (*kvpb.ConditionFailedError) unexpected value
 
+can-use x
+----
+true
+
 rollback x
 ----
 1 [1-1]
@@ -398,6 +450,10 @@ get conflict-key locking nowait
 ----
 (*kvpb.LockConflictError) conflicting locks on "conflict-key" [reason=wait_policy]
 
+can-use x
+----
+true
+
 rollback x
 ----
 0 <noignore>
@@ -406,6 +462,10 @@ put conflict-key b nowait
 ----
 (*kvpb.LockConflictError) conflicting locks on "conflict-key" [reason=wait_policy]
 
+can-use x
+----
+true
+
 rollback x
 ----
 1 [1-1]
@@ -439,6 +499,10 @@ get conflict-key-2 lock-timeout
 ----
 (*kvpb.LockConflictError) conflicting locks on "conflict-key-2" [reason=lock_timeout]
 
+can-use x
+----
+true
+
 rollback x
 ----
 0 <noignore>
@@ -447,6 +511,10 @@ put conflict-key-2 b lock-timeout
 ----
 (*kvpb.LockConflictError) conflicting locks on "conflict-key-2" [reason=lock_timeout]
 
+can-use x
+----
+true
+
 rollback x
 ----
 1 [1-1]
@@ -470,6 +538,10 @@ inject-error
 ----
 injected error
 
+can-use x
+----
+false
+
 rollback x
 ----
 (*pgerror.withCandidateCode) unimplemented: cannot rollback to savepoint after error
@@ -498,10 +570,18 @@ reset
 txn error cleared
 txn id changed
 
+can-use x
+----
+true
+
 release x
 ----
 0 <noignore>
 
+can-use x
+----
+true
+
 rollback x
 ----
 0 <noignore>
@@ -527,6 +607,10 @@ retry
 synthetic error: TransactionRetryWithProtoRefreshError: forced retry
 epoch: 0 -> 1
 
+can-use x
+----
+false
+
 rollback x
 ----
 (*kvpb.TransactionRetryWithProtoRefreshError) TransactionRetryWithProtoRefreshError: cannot rollback to savepoint after a transaction restart

pkg/kv/kvclient/kvcoord/txn_coord_sender_savepoints.go

@@ -162,6 +162,22 @@ func (tc *TxnCoordSender) ReleaseSavepoint(ctx context.Context, s kv.SavepointTo
 	return tc.checkSavepointLocked(sp, "release")
 }
 
+// CanUseSavepoint is part of the kv.TxnSender interface.
+func (tc *TxnCoordSender) CanUseSavepoint(ctx context.Context, s kv.SavepointToken) bool {
+	if tc.typ != kv.RootTxn {
+		return false
+	}
+	tc.mu.Lock()
+	defer tc.mu.Unlock()
+	if tc.mu.txnState != txnPending {
+		return false
+	}
+	// We swallow the error here because we aren't actually performing any
+	// operation with the savepoint; only checking if we are allowed to do so.
+	sp := s.(*savepoint)
+	return tc.checkSavepointLocked(sp, "release") == nil
+}
+
 type errSavepointOperationInErrorTxn struct{}
 
 // ErrSavepointOperationInErrorTxn is reported when CreateSavepoint()

pkg/kv/kvclient/kvcoord/txn_coord_sender_savepoints_test.go

@@ -220,6 +220,15 @@ func TestSavepoints(t *testing.T) {
 					ptxn()
 				}
 
+			case "can-use":
+				spn := td.CmdArgs[0].Key
+				spt := sp[spn]
+				if txn.CanUseSavepoint(ctx, spt) {
+					fmt.Fprintf(&buf, "true\n")
+				} else {
+					fmt.Fprintf(&buf, "false\n")
+				}
+
 			default:
 				td.Fatalf(t, "unknown directive: %s", td.Cmd)
 			}

pkg/kv/mock_transactional_sender.go

@@ -201,6 +201,11 @@ func (m *MockTransactionalSender) ReleaseSavepoint(context.Context, SavepointTok
 	panic("unimplemented")
 }
 
+// CanUseSavepoint is part of the kv.TxnSender interface.
+func (m *MockTransactionalSender) CanUseSavepoint(context.Context, SavepointToken) bool {
+	panic("unimplemented")
+}
+
 // Epoch is part of the TxnSender interface.
 func (m *MockTransactionalSender) Epoch() enginepb.TxnEpoch { panic("unimplemented") }
 

pkg/kv/sender.go

@@ -164,6 +164,10 @@ type TxnSender interface {
 	// This method is only valid when called on RootTxns.
 	ReleaseSavepoint(context.Context, SavepointToken) error
 
+	// CanUseSavepoint checks whether it would be valid to roll back or release
+	// the given savepoint in the current transaction state. It will never error.
+	CanUseSavepoint(context.Context, SavepointToken) bool
+
 	// SetFixedTimestamp makes the transaction run in an unusual way, at
 	// a "fixed timestamp": Timestamp and ReadTimestamp are set to ts,
 	// there's no clock uncertainty, and the txn's deadline is set to ts

pkg/kv/txn.go

@@ -1672,6 +1672,14 @@ func (txn *Txn) ReleaseSavepoint(ctx context.Context, s SavepointToken) error {
 	return txn.mu.sender.ReleaseSavepoint(ctx, s)
 }
 
+// CanUseSavepoint checks whether it would be valid to roll back or release
+// the given savepoint in the current transaction state. It will never error.
+func (txn *Txn) CanUseSavepoint(ctx context.Context, s SavepointToken) bool {
+	txn.mu.Lock()
+	defer txn.mu.Unlock()
+	return txn.mu.sender.CanUseSavepoint(ctx, s)
+}
+
 // DeferCommitWait defers the transaction's commit-wait operation, passing
 // responsibility of commit-waiting from the Txn to the caller of this
 // method. The method returns a function which the caller must eventually

pkg/sql/BUILD.bazel

@@ -457,7 +457,6 @@ go_library(
         "//pkg/sql/pgwire/pgwirecancel",
         "//pkg/sql/physicalplan",
         "//pkg/sql/physicalplan/replicaoracle",
-        "//pkg/sql/plpgsql",
         "//pkg/sql/privilege",
         "//pkg/sql/protoreflect",
         "//pkg/sql/querycache",