isession: add savepoint support

This adds Savepoint support to the internal session. LDR needs save
points for performance reasons. It allows it to handle individual LWW
losers without aborting the entire replication batch.

Release note: none
Epic: CRDB-48647

Author: jeffswenson

pkg/sql/isession/internal_session.go

@@ -52,6 +52,25 @@ var (
 		AST: &tree.RollbackTransaction{},
 		SQL: "ROLLBACK",
 	}
+
+	savepointBegin = statements.Statement[tree.Statement]{
+		AST: &tree.Savepoint{
+			Name: tree.Name("internal_session"),
+		},
+		SQL: "SAVEPOINT internal_session",
+	}
+	savepointRelease = statements.Statement[tree.Statement]{
+		AST: &tree.ReleaseSavepoint{
+			Savepoint: tree.Name("internal_session"),
+		},
+		SQL: "RELEASE SAVEPOINT internal_session",
+	}
+	savepointRollback = statements.Statement[tree.Statement]{
+		AST: &tree.RollbackToSavepoint{
+			Savepoint: tree.Name("internal_session"),
+		},
+		SQL: "ROLLBACK TO SAVEPOINT internal_session",
+	}
 )
 
 var _ isql.Session = &InternalSession{}
@@ -115,6 +134,40 @@ func (i *InternalSession) Txn(ctx context.Context, do func(ctx context.Context)
 	return err
 }
 
+func (i *InternalSession) Savepoint(ctx context.Context, do func(ctx context.Context) error) error {
+	if i.poison != nil {
+		return i.poison
+	}
+
+	if err := i.executeStatement(ctx, savepointBegin); err != nil {
+		return errors.Wrap(err, "failed to create savepoint")
+	}
+
+	innerErr := do(ctx)
+	if innerErr != nil {
+		// Return the rollback error as primary since it indicates a more
+		// serious problem than the original error.
+		savePointErr := i.executeStatement(ctx, savepointRollback)
+		if savePointErr != nil {
+			return errors.CombineErrors(savePointErr, innerErr)
+		}
+		// NOTE: Rollback does not release the savepoint. We need to release the
+		// savepoint to avoid leaking it and causing weird behavior if the user is
+		// nesting savepoint calls.
+		releaseErr := i.executeStatement(ctx, savepointRelease)
+		if releaseErr != nil {
+			return errors.CombineErrors(releaseErr, innerErr)
+		}
+		return innerErr
+	}
+
+	if err := i.executeStatement(ctx, savepointRelease); err != nil {
+		return errors.Wrap(err, "failed to release the savepoint")
+	}
+
+	return nil
+}
+
 func (i *InternalSession) Prepare(
 	ctx context.Context, name string, stmt statements.Statement[tree.Statement], types []*types.T,
 ) (isql.PreparedStatement, error) {
@@ -233,7 +286,7 @@ func (i *InternalSession) executeStatement(
 		return errors.Wrap(err, "unable to push sync statement")
 	}
 	_, _, err = i.readResults(ctx)
-	return errors.Wrap(err, "unable to execute raw statement")
+	return errors.Wrapf(err, "unable to execute raw statement %s", stmt.SQL)
 }
 
 func (i *InternalSession) readResults(ctx context.Context) ([]tree.Datums, int, error) {

pkg/sql/isession/internal_session_test.go

@@ -538,3 +538,77 @@ func tcTableNotFound(
 	_, err = session.Prepare(ctx, "table-not-found-select", stmt, nil)
 	require.ErrorContains(t, err, "relation \"defaultdb.non_existent_table\" does not exist")
 }
+
+func TestSavepoint(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	ctx := context.Background()
+	s := serverutils.StartServerOnly(t, base.TestServerArgs{})
+	defer s.Stopper().Stop(ctx)
+
+	db := s.SQLConn(t)
+	_, err := db.Exec("CREATE TABLE test (id INT PRIMARY KEY, val INT)")
+	require.NoError(t, err)
+
+	idb := s.InternalDB().(descs.DB)
+	session, err := idb.Session(ctx, "test-session")
+	require.NoError(t, err)
+	defer session.Close(ctx)
+
+	insertStmt, err := parser.ParseOne("INSERT INTO defaultdb.test VALUES ($1, $2)")
+	require.NoError(t, err)
+
+	insertPrepared, err := session.Prepare(ctx, "insert", insertStmt, []*types.T{types.Int, types.Int})
+	require.NoError(t, err)
+
+	selectStmt, err := parser.ParseOne("SELECT id, val FROM defaultdb.test ORDER BY id")
+	require.NoError(t, err)
+
+	selectPrepared, err := session.Prepare(ctx, "select", selectStmt, nil)
+	require.NoError(t, err)
+
+	// Create a a three deep savepoint and rollback the inner two savepoints.
+	err = session.Txn(ctx, func(ctx context.Context) error {
+		return session.Savepoint(ctx, func(ctx context.Context) error {
+			_, err := session.ExecutePrepared(ctx, insertPrepared, tree.Datums{
+				tree.NewDInt(tree.DInt(1)),
+				tree.NewDInt(tree.DInt(10)),
+			})
+			if err != nil {
+				return err
+			}
+			err = session.Savepoint(ctx, func(ctx context.Context) error {
+				_, err := session.ExecutePrepared(ctx, insertPrepared, tree.Datums{
+					tree.NewDInt(tree.DInt(2)),
+					tree.NewDInt(tree.DInt(20)),
+				})
+				if err != nil {
+					return err
+				}
+				return session.Savepoint(ctx, func(ctx context.Context) error {
+					_, err := session.ExecutePrepared(ctx, insertPrepared, tree.Datums{
+						tree.NewDInt(tree.DInt(3)),
+						tree.NewDInt(tree.DInt(30)),
+					})
+					if err != nil {
+						return err
+					}
+					// Rollback this innermost savepoint
+					return errors.New("rollback innermost savepoint")
+				})
+			})
+			// The second savepoint should fail due to the nested savepoint rollback
+			require.ErrorContains(t, err, "rollback innermost savepoint")
+			return nil
+		})
+	})
+	require.NoError(t, err)
+
+	// Verify final state - only first insert should be committed
+	rows, err := session.QueryPrepared(ctx, selectPrepared, nil)
+	require.NoError(t, err)
+	require.Equal(t, []tree.Datums{
+		{tree.NewDInt(tree.DInt(1)), tree.NewDInt(tree.DInt(10))},
+	}, rows)
+}

pkg/sql/isql/isql_session.go

@@ -64,6 +64,22 @@ type Session interface {
 	// }
 	Txn(ctx context.Context, do func(context.Context) error) error
 
+	// Savepoint creates a savepoint within an existing transaction and executes
+	// the given function. If the function returns an error, the savepoint is
+	// rolled back. If the function succeeds, the savepoint is released.
+	// Savepoints must be used within a transaction.
+	//
+	// Example:
+	// err := session.Txn(ctx, func(ctx context.Context) error {
+	// 	return session.Savepoint(ctx, func(ctx context.Context) error {
+	// 		return session.ExecutePrepared(ctx, stmt, []tree.Datum{tree.NewDInt(1)})
+	// 	})
+	// })
+	// if err != nil {
+	// 	return err
+	// }
+	Savepoint(ctx context.Context, do func(context.Context) error) error
+
 	// Close closes the session and cleans up internal resources.
 	Close(ctx context.Context)
 }