package store import ( "fmt" "path/filepath" "time" "github.com/rs/zerolog" "github.com/spf13/viper" "gorm.io/driver/mysql" "gorm.io/driver/postgres" "gorm.io/driver/sqlite" "gorm.io/gorm" "gorm.io/gorm/logger" "kumoly.io/kumoly/app/util" ) type DBTYPE string const ( MYSQL DBTYPE = "mysql" SQLITE DBTYPE = "sqlite" POSTGRES DBTYPE = "postgres" ) type Store struct { TYPE DBTYPE DB *gorm.DB User string Password string Host string Port string Name string AutoMigrate bool Prod bool Path string config *gorm.Config } var DB *gorm.DB var std *Store func init() { // Database // type [mysql,sqlite] viper.SetDefault("db.type", string(MYSQL)) // mysql default viper.SetDefault("db.user", "root") viper.SetDefault("db.passwd", "admin") viper.SetDefault("db.host", "127.0.0.1") viper.SetDefault("db.port", "3306") viper.SetDefault("db.name", "app") viper.SetDefault("db.reconnect_interval", 5) viper.SetDefault("db.automigrate", true) } var l zerolog.Logger // Init initialize default db using New() followed by Connect() func Setup() { l = util.Klog.With().Str("mod", "store").Logger() var err error std = New(DBTYPE(viper.GetString("db.type"))) DB, err = std.Connect() if err != nil { l.Error().Err(err).Msg("std connection error") panic(err) } } func Migrate(dst ...interface{}) error { return std.Migrate(dst...) } func New(t DBTYPE) *Store { s := &Store{ TYPE: t, Prod: viper.GetBool("prod"), Name: viper.GetString("db.name"), AutoMigrate: viper.GetBool("db.automigrate"), } if t == MYSQL || t == POSTGRES { s.User = viper.GetString("db.user") s.Password = viper.GetString("db.passwd") s.Host = viper.GetString("db.host") s.Port = viper.GetString("db.port") } else if t == SQLITE { s.Path = viper.GetString("data") } else { err := fmt.Errorf("unknown db type %s", t) l.Error().Err(err).Msg("unknown db type") panic(err) } return s } // Connect to db and store db connection to var DB func (s *Store) Connect() (db *gorm.DB, err error) { glog := logger.Default l.Trace(). // Int("level", viper.GetInt("log.level")). // Bool("tog_info", (zerolog.Level(viper.GetInt("log.level")) <= zerolog.DebugLevel)). Msgf("log level %v", viper.GetInt("log.level")) if zerolog.Level(viper.GetInt("log.level")) <= zerolog.DebugLevel { l.Debug().Msg("set gorm log level to info") glog.LogMode(logger.Info) } s.config = &gorm.Config{ Logger: glog, SkipDefaultTransaction: true, } if s.Prod { s.config.Logger = logger.Discard } switch s.TYPE { case MYSQL: return s.DB, s.mysqlConnector() case POSTGRES: return s.DB, s.postgresConnector() case SQLITE: return s.DB, s.sqliteConnector() default: return nil, fmt.Errorf("unknown db type %s", s.TYPE) } } func (s *Store) postgresConnector() error { l.Info().Msg("Connecting to postgres...") dsn := fmt.Sprintf( "host=%v user=%v password=%v dbname=%v port=%v sslmode=disable TimeZone=Asia/Taipei", s.Host, s.User, s.Password, s.Name, s.Port, ) for { db, err := gorm.Open(postgres.New(postgres.Config{ DSN: dsn, // data source name // PreferSimpleProtocol: true, // disables implicit prepared statement usage }), s.config) if err == nil { s.DB = db break } inter := viper.GetDuration("db.reconnect_interval") if inter <= 0 { inter = 5 } l.Warn().Err(err).Msg("Unable to connect to database") l.Warn().Msgf("Retrying in %v second.", inter) time.Sleep(time.Second * inter) } l.Info().Msg("Connection to postgres, ok.") return nil } //mysqlConnector connection func (s *Store) mysqlConnector() error { l.Info().Msg("Connecting to mysql...") dsn := fmt.Sprintf( "%v:%v@tcp(%v:%v)/%v?charset=utf8mb4&parseTime=True&loc=Local", s.User, s.Password, s.Host, s.Port, s.Name, ) // l.Debug(dsn) for { db, err := gorm.Open(mysql.New(mysql.Config{ DSN: dsn, // data source name DefaultStringSize: 256, // default size for string fields }), s.config) if err == nil { s.DB = db break } inter := viper.GetDuration("db.reconnect_interval") if inter <= 0 { inter = 5 } l.Warn().Err(err).Msg("Unable to connect to database") l.Warn().Msgf("Retrying in %v second.", inter) time.Sleep(time.Second * inter) } l.Info().Msg("Connection to mysql, ok.") return nil } //sqliteConnector connection func (s *Store) sqliteConnector() error { util.Mkdir(s.Path) dbPath := filepath.Join(s.Path, s.Name+".db") l.Info().Str("path", dbPath).Msg("Connecting to sqlite...") db, err := gorm.Open(sqlite.Open(dbPath), s.config) if err != nil { l.Error().Err(err).Msg("failed to connect database") return err } s.DB = db return nil } func (s *Store) Migrate(dst ...interface{}) error { if !s.AutoMigrate { l.Debug().Msg("AutoMigration is set to off, migration skipped") return nil } return s.DB.AutoMigrate(dst...) }