Reuse opened session in the context (#44939)

pkg/services/sqlstore/session.go

@@ -5,11 +5,16 @@ import (
 	"reflect"
 
 	"xorm.io/xorm"
+
+	"github.com/grafana/grafana/pkg/infra/log"
 )
 
+var sessionLogger = log.New("sqlstore.session")
+
 type DBSession struct {
 	*xorm.Session
-	events []interface{}
+	transactionOpen bool
+	events          []interface{}
 }
 
 type DBTransactionFunc func(sess *DBSession) error
@@ -32,26 +37,27 @@ func newSession(ctx context.Context) *DBSession {
 	return sess
 }
 
-func startSession(ctx context.Context, engine *xorm.Engine, beginTran bool) (*DBSession, error) {
+func startSessionOrUseExisting(ctx context.Context, engine *xorm.Engine, beginTran bool) (*DBSession, bool, error) {
 	value := ctx.Value(ContextSessionKey{})
 	var sess *DBSession
 	sess, ok := value.(*DBSession)
 
 	if ok {
+		sessionLogger.Debug("reusing existing session", "transaction", sess.transactionOpen)
 		sess.Session = sess.Session.Context(ctx)
-		return sess, nil
+		return sess, false, nil
 	}
 
-	newSess := &DBSession{Session: engine.NewSession()}
+	newSess := &DBSession{Session: engine.NewSession(), transactionOpen: beginTran}
 	if beginTran {
 		err := newSess.Begin()
 		if err != nil {
-			return nil, err
+			return nil, false, err
 		}
 	}
 
 	newSess.Session = newSess.Session.Context(ctx)
-	return newSess, nil
+	return newSess, true, nil
 }
 
 // WithDbSession calls the callback with a session.
@@ -60,10 +66,13 @@ func (ss *SQLStore) WithDbSession(ctx context.Context, callback DBTransactionFun
 }
 
 func withDbSession(ctx context.Context, engine *xorm.Engine, callback DBTransactionFunc) error {
-	sess := &DBSession{Session: engine.NewSession()}
-	sess.Session = sess.Session.Context(ctx)
-	defer sess.Close()
-
+	sess, isNew, err := startSessionOrUseExisting(ctx, engine, false)
+	if err != nil {
+		return err
+	}
+	if isNew {
+		defer sess.Close()
+	}
 	return callback(sess)
 }
 

pkg/services/sqlstore/transactions.go

@@ -3,13 +3,15 @@ package sqlstore
 import (
 	"context"
 	"errors"
+	"fmt"
 	"time"
 
+	"github.com/mattn/go-sqlite3"
+	"xorm.io/xorm"
+
 	"github.com/grafana/grafana/pkg/bus"
 	"github.com/grafana/grafana/pkg/infra/log"
 	"github.com/grafana/grafana/pkg/util/errutil"
-	"github.com/mattn/go-sqlite3"
-	"xorm.io/xorm"
 )
 
 var tsclogger = log.New("sqlstore.transactions")
@@ -35,15 +37,28 @@ func inTransactionWithRetry(callback DBTransactionFunc, retry int) error {
 }
 
 func inTransactionWithRetryCtx(ctx context.Context, engine *xorm.Engine, callback DBTransactionFunc, retry int) error {
-	sess, err := startSession(ctx, engine, true)
+	sess, isNew, err := startSessionOrUseExisting(ctx, engine, true)
 	if err != nil {
 		return err
 	}
 
-	defer sess.Close()
+	if !sess.transactionOpen && !isNew {
+		// this should not happen because the only place that creates reusable session begins a new transaction.
+		return fmt.Errorf("cannot reuse existing session that did not start transaction")
+	}
+
+	if isNew { // if this call initiated the session, it should be responsible for closing it.
+		defer sess.Close()
+	}
 
 	err = callback(sess)
 
+	if !isNew {
+		tsclogger.Debug("skip committing the transaction because it belongs to a session created in the outer scope")
+		// Do not commit the transaction if the session was reused.
+		return err
+	}
+
 	// special handling of database locked errors for sqlite, then we can retry 5 times
 	var sqlError sqlite3.Error
 	if errors.As(err, &sqlError) && retry < 5 && (sqlError.Code == sqlite3.ErrLocked || sqlError.Code == sqlite3.ErrBusy) {

pkg/services/sqlstore/transactions_test.go

@@ -8,8 +8,9 @@ import (
 	"errors"
 	"testing"
 
-	"github.com/grafana/grafana/pkg/models"
 	"github.com/stretchr/testify/require"
+
+	"github.com/grafana/grafana/pkg/models"
 )
 
 var ErrProvokedError = errors.New("testing error")
@@ -54,3 +55,59 @@ func TestTransaction(t *testing.T) {
 		require.Equal(t, cmd.Result.Id, query.Result.Id)
 	})
 }
+
+func TestReuseSessionWithTransaction(t *testing.T) {
+	ss := InitTestDB(t)
+
+	t.Run("top level transaction", func(t *testing.T) {
+		var outerSession *DBSession
+		err := ss.InTransaction(context.Background(), func(ctx context.Context) error {
+			value := ctx.Value(ContextSessionKey{})
+			var ok bool
+			outerSession, ok = value.(*DBSession)
+
+			require.True(t, ok, "Session should be available in the context but it does not exist")
+			require.True(t, outerSession.transactionOpen, "Transaction should be open")
+
+			require.NoError(t, ss.WithDbSession(ctx, func(sess *DBSession) error {
+				require.Equal(t, outerSession, sess)
+				require.False(t, sess.IsClosed(), "Session is closed but it should not be")
+				return nil
+			}))
+
+			require.NoError(t, ss.WithTransactionalDbSession(ctx, func(sess *DBSession) error {
+				require.Equal(t, outerSession, sess)
+				require.False(t, sess.IsClosed(), "Session is closed but it should not be")
+				return nil
+			}))
+
+			require.False(t, outerSession.IsClosed(), "Session is closed but it should not be")
+			return nil
+		})
+
+		require.NoError(t, err)
+		require.True(t, outerSession.IsClosed())
+	})
+
+	t.Run("fails if reuses session without transaction", func(t *testing.T) {
+		require.NoError(t, ss.WithDbSession(context.Background(), func(outerSession *DBSession) error {
+			require.NotNil(t, outerSession)
+			require.NotNil(t, outerSession.DB()) // init the session
+			require.False(t, outerSession.IsClosed(), "Session is closed but it should not be")
+
+			ctx := context.WithValue(context.Background(), ContextSessionKey{}, outerSession)
+
+			require.NoError(t, ss.WithDbSession(ctx, func(sess *DBSession) error {
+				require.Equal(t, outerSession, sess)
+				require.False(t, sess.IsClosed(), "Session is closed but it should not be")
+				return nil
+			}))
+
+			require.Error(t, ss.WithTransactionalDbSession(ctx, func(sess *DBSession) error {
+				require.FailNow(t, "WithTransactionalDbSession should not be able to reuse session that did not open the transaction ")
+				return nil
+			}))
+			return nil
+		}))
+	})
+}