woodpecker/shared/database/database.go

package database

//go:generate go-bindata -pkg database -o database_gen.go sqlite3/ mysql/ postgres/

import (
	"database/sql"
	"os"

	"github.com/drone/drone/shared/envconfig"

	log "github.com/Sirupsen/logrus"
	_ "github.com/go-sql-driver/mysql"
	_ "github.com/lib/pq"
	_ "github.com/mattn/go-sqlite3"
	"github.com/rubenv/sql-migrate"
	"github.com/russross/meddler"
)

func Load(env envconfig.Env) *sql.DB {
	var (
		driver = env.String("DATABASE_DRIVER", "sqlite3")
		config = env.String("DATABASE_CONFIG", "drone.sqlite")
	)

	log.Infof("using database driver %s", driver)
	log.Infof("using database config %s", config)

	return Open(driver, config)
}

// Open opens a database connection, runs the database migrations, and returns
// the database connection. Any errors connecting to the database or executing
// migrations will cause the application to exit.
func Open(driver, config string) *sql.DB {
	var db, err = sql.Open(driver, config)
	if err != nil {
		log.Errorln(err)
		log.Fatalln("database connection failed")
	}
	switch driver {
	case "mysql":
		meddler.Default = meddler.MySQL
	case "postgres":
		meddler.Default = meddler.PostgreSQL
	}

	var migrations = &migrate.AssetMigrationSource{
		Asset:    Asset,
		AssetDir: AssetDir,
		Dir:      driver,
	}

	_, err = migrate.Exec(db, driver, migrations, migrate.Up)
	if err != nil {
		log.Errorln(err)
		log.Fatalln("migration failed")
	}
	return db
}

func OpenTest() *sql.DB {
	var (
		driver = "sqlite3"
		config = ":memory:"
	)
	if os.Getenv("DATABASE_DRIVER") != "" {
		driver = os.Getenv("DATABASE_DRIVER")
		config = os.Getenv("DATABASE_CONFIG")
	}
	return Open(driver, config)
}