bus: support multiple dispatch in one transaction

this makes it possible to run multiple DispatchCtx
in one transaction. The TransactionManager will
start/end the transaction and pass the dbsession
in the context.Context variable

pkg/services/sqlstore/shared.go

@@ -1,6 +1,7 @@
 package sqlstore
 
 import (
+	"context"
 	"reflect"
 	"time"
 
@@ -29,10 +30,35 @@ func inTransaction(callback dbTransactionFunc) error {
 	return inTransactionWithRetry(callback, 0)
 }
 
+func startSession(ctx context.Context) *DBSession {
+	value := ctx.Value(ContextSessionName)
+	var sess *xorm.Session
+	sess, ok := value.(*xorm.Session)
+
+	if !ok {
+		return newSession()
+	}
+
+	old := newSession()
+	old.Session = sess
+
+	return old
+}
+
+func withDbSession(ctx context.Context, callback dbTransactionFunc) error {
+	sess := startSession(ctx)
+
+	return callback(sess)
+}
+
 func inTransactionWithRetry(callback dbTransactionFunc, retry int) error {
+	return inTransactionWithRetryCtx(context.Background(), callback, retry)
+}
+
+func inTransactionWithRetryCtx(ctx context.Context, callback dbTransactionFunc, retry int) error {
 	var err error
 
-	sess := newSession()
+	sess := startSession(ctx)
 	defer sess.Close()
 
 	if err = sess.Begin(); err != nil {

pkg/services/sqlstore/sqlstore.go

@@ -1,6 +1,8 @@
 package sqlstore
 
 import (
+	"context"
+	"errors"
 	"fmt"
 	"net/url"
 	"os"
@@ -35,6 +37,8 @@ var (
 	sqlog log.Logger = log.New("sqlstore")
 )
 
+const ContextSessionName = "db-session"
+
 func init() {
 	registry.Register(&registry.Descriptor{
 		Name:         "SqlStore",
@@ -45,6 +49,7 @@ func init() {
 
 type SqlStore struct {
 	Cfg *setting.Cfg `inject:""`
+	Bus bus.Bus      `inject:""`
 
 	dbCfg           DatabaseConfig
 	engine          *xorm.Engine
@@ -77,6 +82,10 @@ func (ss *SqlStore) Init() error {
 	// Init repo instances
 	annotations.SetRepository(&SqlAnnotationRepo{})
 
+	ss.Bus.SetTransactionManager(&SQLTransactionManager{
+		engine: ss.engine,
+	})
+
 	// ensure admin user
 	if ss.skipEnsureAdmin {
 		return nil
@@ -85,10 +94,47 @@ func (ss *SqlStore) Init() error {
 	return ss.ensureAdminUser()
 }
 
+type SQLTransactionManager struct {
+	engine *xorm.Engine
+}
+
+func (stm *SQLTransactionManager) Begin(ctx context.Context) (context.Context, error) {
+	sess := stm.engine.NewSession()
+	err := sess.Begin()
+	if err != nil {
+		return ctx, err
+	}
+
+	withValue := context.WithValue(ctx, ContextSessionName, sess)
+
+	return withValue, nil
+}
+
+func (stm *SQLTransactionManager) End(ctx context.Context, err error) error {
+	value := ctx.Value(ContextSessionName)
+	sess, ok := value.(*xorm.Session)
+	if !ok {
+		return errors.New("context is missing transaction")
+	}
+
+	if err != nil {
+		sess.Rollback()
+		return err
+	}
+
+	defer sess.Close()
+
+	return sess.Commit()
+}
+
 func (ss *SqlStore) ensureAdminUser() error {
 	systemUserCountQuery := m.GetSystemUserCountStatsQuery{}
 
-	if err := bus.Dispatch(&systemUserCountQuery); err != nil {
+	err := bus.InTransaction(context.Background(), func(ctx context.Context) error {
+		return bus.DispatchCtx(ctx, &systemUserCountQuery)
+	})
+
+	if err != nil {
 		return fmt.Errorf("Could not determine if admin user exists: %v", err)
 	}
 
@@ -240,6 +286,7 @@ func (ss *SqlStore) readConfig() {
 func InitTestDB(t *testing.T) *SqlStore {
 	sqlstore := &SqlStore{}
 	sqlstore.skipEnsureAdmin = true
+	sqlstore.Bus = bus.New()
 
 	dbType := migrator.SQLITE
 

pkg/services/sqlstore/stats.go

@@ -1,6 +1,7 @@
 package sqlstore
 
 import (
+	"context"
 	"time"
 
 	"github.com/grafana/grafana/pkg/bus"
@@ -13,6 +14,7 @@ func init() {
 	bus.AddHandler("sql", GetDataSourceAccessStats)
 	bus.AddHandler("sql", GetAdminStats)
 	bus.AddHandler("sql", GetSystemUserCountStats)
+	bus.AddCtxHandler("sql", GetSystemUserCountStatsCtx)
 }
 
 var activeUserTimeLimit = time.Hour * 24 * 30
@@ -133,6 +135,22 @@ func GetAdminStats(query *m.GetAdminStatsQuery) error {
 	return err
 }
 
+func GetSystemUserCountStatsCtx(ctx context.Context, query *m.GetSystemUserCountStatsQuery) error {
+	return withDbSession(ctx, func(sess *DBSession) error {
+
+		var rawSql = `SELECT COUNT(id) AS Count FROM ` + dialect.Quote("user")
+		var stats m.SystemUserCountStats
+		_, err := sess.SQL(rawSql).Get(&stats)
+		if err != nil {
+			return err
+		}
+
+		query.Result = &stats
+
+		return err
+	})
+}
+
 func GetSystemUserCountStats(query *m.GetSystemUserCountStatsQuery) error {
 	var rawSql = `SELECT COUNT(id) AS Count FROM ` + dialect.Quote("user")
 	var stats m.SystemUserCountStats