Chore: Deduplicate sqlstore transaction code (#17069)

* Deduplicate transaction code

* More deduplication

pkg/services/sqlstore/session.go

@@ -18,6 +18,11 @@ func (sess *DBSession) publishAfterCommit(msg interface{}) {
 	sess.events = append(sess.events, msg)
 }
 
+// NewSession returns a new DBSession
+func (ss *SqlStore) NewSession() *DBSession {
+	return &DBSession{Session: ss.engine.NewSession()}
+}
+
 func newSession() *DBSession {
 	return &DBSession{Session: x.NewSession()}
 }
@@ -41,6 +46,16 @@ func startSession(ctx context.Context, engine *xorm.Engine, beginTran bool) (*DB
 	return newSess, nil
 }
 
+// WithDbSession calls the callback with an session attached to the context.
+func (ss *SqlStore) WithDbSession(ctx context.Context, callback dbTransactionFunc) error {
+	sess, err := startSession(ctx, ss.engine, false)
+	if err != nil {
+		return err
+	}
+
+	return callback(sess)
+}
+
 func withDbSession(ctx context.Context, callback dbTransactionFunc) error {
 	sess, err := startSession(ctx, x, false)
 	if err != nil {

pkg/services/sqlstore/sqlstore.go

@@ -26,7 +26,6 @@ import (
 	_ "github.com/grafana/grafana/pkg/tsdb/mssql"
 	"github.com/grafana/grafana/pkg/util"
 	_ "github.com/lib/pq"
-	sqlite3 "github.com/mattn/go-sqlite3"
 )
 
 var (
@@ -58,64 +57,6 @@ type SqlStore struct {
 	skipEnsureAdmin bool
 }
 
-// NewSession returns a new DBSession
-func (ss *SqlStore) NewSession() *DBSession {
-	return &DBSession{Session: ss.engine.NewSession()}
-}
-
-// WithDbSession calls the callback with an session attached to the context.
-func (ss *SqlStore) WithDbSession(ctx context.Context, callback dbTransactionFunc) error {
-	sess, err := startSession(ctx, ss.engine, false)
-	if err != nil {
-		return err
-	}
-
-	return callback(sess)
-}
-
-// WithTransactionalDbSession calls the callback with an session within a transaction
-func (ss *SqlStore) WithTransactionalDbSession(ctx context.Context, callback dbTransactionFunc) error {
-	return ss.inTransactionWithRetryCtx(ctx, callback, 0)
-}
-
-func (ss *SqlStore) inTransactionWithRetryCtx(ctx context.Context, callback dbTransactionFunc, retry int) error {
-	sess, err := startSession(ctx, ss.engine, true)
-	if err != nil {
-		return err
-	}
-
-	defer sess.Close()
-
-	err = callback(sess)
-
-	// special handling of database locked errors for sqlite, then we can retry 3 times
-	if sqlError, ok := err.(sqlite3.Error); ok && retry < 5 {
-		if sqlError.Code == sqlite3.ErrLocked {
-			sess.Rollback()
-			time.Sleep(time.Millisecond * time.Duration(10))
-			sqlog.Info("Database table locked, sleeping then retrying", "retry", retry)
-			return ss.inTransactionWithRetryCtx(ctx, callback, retry+1)
-		}
-	}
-
-	if err != nil {
-		sess.Rollback()
-		return err
-	} else if err = sess.Commit(); err != nil {
-		return err
-	}
-
-	if len(sess.events) > 0 {
-		for _, e := range sess.events {
-			if err = bus.Publish(e); err != nil {
-				log.Error(3, "Failed to publish event after commit. error: %v", err)
-			}
-		}
-	}
-
-	return nil
-}
-
 func (ss *SqlStore) Init() error {
 	ss.log = log.New("sqlstore")
 	ss.readConfig()

pkg/services/sqlstore/transactions.go

@@ -4,63 +4,34 @@ import (
 	"context"
 	"time"
 
+	"github.com/go-xorm/xorm"
 	"github.com/grafana/grafana/pkg/bus"
 	"github.com/grafana/grafana/pkg/infra/log"
 	sqlite3 "github.com/mattn/go-sqlite3"
 )
 
+// WithTransactionalDbSession calls the callback with an session within a transaction
+func (ss *SqlStore) WithTransactionalDbSession(ctx context.Context, callback dbTransactionFunc) error {
+	return inTransactionWithRetryCtx(ss.engine, ctx, callback, 0)
+}
+
 func (ss *SqlStore) InTransaction(ctx context.Context, fn func(ctx context.Context) error) error {
 	return ss.inTransactionWithRetry(ctx, fn, 0)
 }
 
 func (ss *SqlStore) inTransactionWithRetry(ctx context.Context, fn func(ctx context.Context) error, retry int) error {
-	sess, err := startSession(ctx, ss.engine, true)
-	if err != nil {
-		return err
-	}
-
-	defer sess.Close()
-
-	withValue := context.WithValue(ctx, ContextSessionName, sess)
-
-	err = fn(withValue)
-
-	// special handling of database locked errors for sqlite, then we can retry 3 times
-	if sqlError, ok := err.(sqlite3.Error); ok && retry < 5 {
-		if sqlError.Code == sqlite3.ErrLocked {
-			sess.Rollback()
-			time.Sleep(time.Millisecond * time.Duration(10))
-			ss.log.Info("Database table locked, sleeping then retrying", "retry", retry)
-			return ss.inTransactionWithRetry(ctx, fn, retry+1)
-		}
-	}
-
-	if err != nil {
-		sess.Rollback()
-		return err
-	}
-
-	if err = sess.Commit(); err != nil {
-		return err
-	}
-
-	if len(sess.events) > 0 {
-		for _, e := range sess.events {
-			if err = bus.Publish(e); err != nil {
-				ss.log.Error("Failed to publish event after commit", err)
-			}
-		}
-	}
-
-	return nil
+	return inTransactionWithRetryCtx(ss.engine, ctx, func(sess *DBSession) error {
+		withValue := context.WithValue(ctx, ContextSessionName, sess)
+		return fn(withValue)
+	}, retry)
 }
 
 func inTransactionWithRetry(callback dbTransactionFunc, retry int) error {
-	return inTransactionWithRetryCtx(context.Background(), callback, retry)
+	return inTransactionWithRetryCtx(x, context.Background(), callback, retry)
 }
 
-func inTransactionWithRetryCtx(ctx context.Context, callback dbTransactionFunc, retry int) error {
-	sess, err := startSession(ctx, x, true)
+func inTransactionWithRetryCtx(engine *xorm.Engine, ctx context.Context, callback dbTransactionFunc, retry int) error {
+	sess, err := startSession(ctx, engine, true)
 	if err != nil {
 		return err
 	}
@@ -102,5 +73,5 @@ func inTransaction(callback dbTransactionFunc) error {
 }
 
 func inTransactionCtx(ctx context.Context, callback dbTransactionFunc) error {
-	return inTransactionWithRetryCtx(ctx, callback, 0)
+	return inTransactionWithRetryCtx(x, ctx, callback, 0)
 }