iot-spray-equipment/database/database.go

82 lines
1.6 KiB
Go

package database
import (
"errors"
"fmt"
"time"
"github.com/spf13/cast"
"gorm.io/driver/mysql"
"gorm.io/gorm"
)
func New(config Config) (*gorm.DB, error) {
db, err := config.Build(&gorm.Config{})
if err != nil {
return nil, err
}
sqldb, err := db.DB()
if err != nil {
return nil, err
}
if config.Has("max_open_conns") {
sqldb.SetMaxOpenConns(config.GetInt("max_open_conns"))
}
if config.Has("max_idle_conns") {
sqldb.SetMaxIdleConns(config.GetInt("max_idle_conns"))
}
if config.Has("conn_max_idle_time") {
sqldb.SetConnMaxIdleTime(config.GetDuration("conn_max_idle_time") * time.Second)
}
if config.Has("conn_max_lifetime") {
sqldb.SetConnMaxLifetime(config.GetDuration("conn_max_lifetime") * time.Second)
}
return db, nil
}
type Config map[string]interface{}
func (cfg Config) GetString(key string) string {
return cast.ToString(cfg[key])
}
func (cfg Config) GetInt(key string) int {
return cast.ToInt(cfg[key])
}
func (cfg Config) GetDuration(key string) time.Duration {
return cast.ToDuration(cfg[key])
}
func (cfg Config) Has(key string) bool {
_, ok := cfg[key]
return ok
}
func (cfg Config) Build(opts ...gorm.Option) (*gorm.DB, error) {
switch cfg.GetString("driver") {
case "mysql":
dsn := cfg.GetString("dsn")
if dsn == "" {
dsn = fmt.Sprintf(
"%v:%v@tcp(%v:%v)/%v?charset=%v&parseTime=True&loc=Local",
cfg.GetString("username"),
cfg.GetString("password"),
cfg.GetString("host"),
cfg.GetInt("port"),
cfg.GetString("database"),
cfg.GetString("charset"),
)
}
return gorm.Open(mysql.Open(dsn), opts...)
default:
return nil, errors.New("database driver not supported")
}
}