energylink/pkg/services/sqlstore/shared.go

package sqlstore

import (
	"context"
	"reflect"
	"time"

	"github.com/go-xorm/xorm"
	"github.com/grafana/grafana/pkg/bus"
	"github.com/grafana/grafana/pkg/log"
	sqlite3 "github.com/mattn/go-sqlite3"
)

type DBSession struct {
	*xorm.Session
	events []interface{}
}

type dbTransactionFunc func(sess *DBSession) error

func (sess *DBSession) publishAfterCommit(msg interface{}) {
	sess.events = append(sess.events, msg)
}

func newSession() *DBSession {
	return &DBSession{Session: x.NewSession()}
}

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 := startSession(ctx)
	defer sess.Close()

	if err = sess.Begin(); err != nil {
		return err
	}

	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 inTransactionWithRetry(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 (sess *DBSession) InsertId(bean interface{}) (int64, error) {
	table := sess.DB().Mapper.Obj2Table(getTypeName(bean))

	dialect.PreInsertId(table, sess.Session)

	id, err := sess.Session.InsertOne(bean)

	dialect.PostInsertId(table, sess.Session)

	return id, err
}

func getTypeName(bean interface{}) (res string) {
	t := reflect.TypeOf(bean)
	for t.Kind() == reflect.Ptr {
		t = t.Elem()
	}
	return t.Name()
}