From 3479cf4b396e84a13cd766a5adde14ba338da515 Mon Sep 17 00:00:00 2001
From: bergquist <carl.bergquist@gmail.com>
Date: Mon, 18 Jun 2018 14:50:36 +0200
Subject: [PATCH] expose functions to use sessions

---
 pkg/services/sqlstore/session.go  | 20 +++++------
 pkg/services/sqlstore/sqlstore.go | 60 ++++++++++++++++++++++++++++++-
 2 files changed, 69 insertions(+), 11 deletions(-)

diff --git a/pkg/services/sqlstore/session.go b/pkg/services/sqlstore/session.go
index c85346231e4..29d7392678f 100644
--- a/pkg/services/sqlstore/session.go
+++ b/pkg/services/sqlstore/session.go
@@ -27,18 +27,18 @@ func startSession(ctx context.Context, engine *xorm.Engine, beginTran bool) (*DB
 	var sess *DBSession
 	sess, ok := value.(*DBSession)
 
-	if !ok {
-		newSess := &DBSession{Session: engine.NewSession()}
-		if beginTran {
-			err := newSess.Begin()
-			if err != nil {
-				return nil, err
-			}
-		}
-		return newSess, nil
+	if ok {
+		return sess, nil
 	}
 
-	return sess, nil
+	newSess := &DBSession{Session: engine.NewSession()}
+	if beginTran {
+		err := newSess.Begin()
+		if err != nil {
+			return nil, err
+		}
+	}
+	return newSess, nil
 }
 
 func withDbSession(ctx context.Context, callback dbTransactionFunc) error {
diff --git a/pkg/services/sqlstore/sqlstore.go b/pkg/services/sqlstore/sqlstore.go
index 40101528df5..b0edc1676e0 100644
--- a/pkg/services/sqlstore/sqlstore.go
+++ b/pkg/services/sqlstore/sqlstore.go
@@ -26,7 +26,7 @@ import (
 
 	_ "github.com/grafana/grafana/pkg/tsdb/mssql"
 	_ "github.com/lib/pq"
-	_ "github.com/mattn/go-sqlite3"
+	sqlite3 "github.com/mattn/go-sqlite3"
 )
 
 var (
@@ -56,6 +56,64 @@ 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", err)
+			}
+		}
+	}
+
+	return nil
+}
+
 func (ss *SqlStore) Init() error {
 	ss.log = log.New("sqlstore")
 	ss.readConfig()