Merge pull request #208 from checkr/zz/add-db-connection-retry

Add retry for DB connection

Gopkg.toml

@@ -153,6 +153,10 @@
   branch = "master"
   name = "github.com/dchest/uniuri"
 
+[[constraint]]
+  name = "github.com/avast/retry-go"
+  version = "2.2.0"
+
 
 [prune]
   go-tests = true

pkg/config/env.go

@@ -42,6 +42,9 @@ var Config = struct {
 	// DBConnectionDebug controls whether to show the database connection debugging logs
 	// warning: it may log the credentials to the stdout
 	DBConnectionDebug bool `env:"FLAGR_DB_DBCONNECTION_DEBUG" envDefault:"true"`
+	// DBConnectionRetryAttempts controls how we are going to retry on db connection when start the flagr server
+	DBConnectionRetryAttempts uint          `env:"FLAGR_DB_DBCONNECTION_RETRY_ATTEMPTS" envDefault:"9"`
+	DBConnectionRetryDelay    time.Duration `env:"FLAGR_DB_DBCONNECTION_RETRY_DELAY" envDefault:"100ms"`
 
 	// CORSEnabled - enable CORS
 	CORSEnabled bool `env:"FLAGR_CORS_ENABLED" envDefault:"true"`

pkg/entity/db.go

@@ -8,6 +8,7 @@ import (
 	_ "github.com/jinzhu/gorm/dialects/postgres" // postgres driver
 	_ "github.com/jinzhu/gorm/dialects/sqlite"   // sqlite driver
 
+	retry "github.com/avast/retry-go"
 	"github.com/checkr/flagr/pkg/config"
 	"github.com/jinzhu/gorm"
 	"github.com/sirupsen/logrus"
@@ -30,10 +31,22 @@ var AutoMigrateTables = []interface{}{
 	FlagEntityType{},
 }
 
+func connectDB() (db *gorm.DB, err error) {
+	err = retry.Do(
+		func() error {
+			db, err = gorm.Open(config.Config.DBDriver, config.Config.DBConnectionStr)
+			return err
+		},
+		retry.Attempts(config.Config.DBConnectionRetryAttempts),
+		retry.Delay(config.Config.DBConnectionRetryDelay),
+	)
+	return db, err
+}
+
 // GetDB gets the db singleton
 func GetDB() *gorm.DB {
 	singletonOnce.Do(func() {
-		db, err := gorm.Open(config.Config.DBDriver, config.Config.DBConnectionStr)
+		db, err := connectDB()
 		if err != nil {
 			if config.Config.DBConnectionDebug {
 				logrus.WithField("err", err).Fatal("failed to connect to db")

pkg/entity/db_test.go

@@ -0,0 +1,47 @@
+package entity
+
+import (
+	"testing"
+
+	"github.com/checkr/flagr/pkg/config"
+	"github.com/stretchr/testify/assert"
+)
+
+func setTestDBConfig(driver string, connectionStr string) (reset func()) {
+	old := config.Config
+
+	config.Config.DBDriver = driver
+	config.Config.DBConnectionStr = connectionStr
+	config.Config.DBConnectionRetryAttempts = 2
+
+	return func() {
+		config.Config = old
+	}
+}
+
+func TestConnectDB(t *testing.T) {
+	t.Run("happy code path", func(t *testing.T) {
+		reset := setTestDBConfig("sqlite3", ":memory:")
+		defer reset()
+
+		db, err := connectDB()
+		assert.NotNil(t, db)
+		assert.NoError(t, err)
+	})
+
+	t.Run("error code path", func(t *testing.T) {
+		reset := setTestDBConfig("mysql", "invalid")
+		defer reset()
+
+		_, err := connectDB()
+		assert.Error(t, err)
+	})
+}
+
+func TestGetDB(t *testing.T) {
+	reset := setTestDBConfig("sqlite3", ":memory:")
+	defer reset()
+
+	db := GetDB()
+	assert.NotNil(t, db)
+}