app/store/db.go

package store

import (
	"fmt"
	"path/filepath"
	"time"

	"github.com/rs/zerolog/log"
	"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)
}

// Init initialize default db using New() followed by Connect()
func Setup() {
	var err error
	std = New(DBTYPE(viper.GetString("db.type")))
	DB, err = std.Connect()
	if err != nil {
		log.Error().Str("mod", "store").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)
		log.Error().Str("mod", "store").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) {
	s.config = &gorm.Config{
		Logger:                 logger.Default,
		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 {
	log.Info().Str("mod", "store").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
		}
		log.Warn().Str("mod", "store").Err(err).Msg("Unable to connect to database")
		log.Warn().Str("mod", "store").Msgf("Retrying in %v second.", inter)
		time.Sleep(time.Second * inter)
	}
	log.Info().Str("mod", "store").Msg("Connection to postgres, ok.")
	return nil
}

//mysqlConnector connection
func (s *Store) mysqlConnector() error {
	log.Info().Str("mod", "store").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,
	)
	// log.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
		}
		log.Warn().Str("mod", "store").Err(err).Msg("Unable to connect to database")
		log.Warn().Str("mod", "store").Msgf("Retrying in %v second.", inter)
		time.Sleep(time.Second * inter)
	}
	log.Info().Str("mod", "store").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")
	log.Info().Str("mod", "store").Str("path", dbPath).Msg("Connecting to sqlite...")
	db, err := gorm.Open(sqlite.Open(dbPath), s.config)
	if err != nil {
		log.Error().Str("mod", "store").Err(err).Msg("failed to connect database")
		return err
	}
	s.DB = db
	return nil
}

func (s *Store) Migrate(dst ...interface{}) error {
	if !s.AutoMigrate {
		log.Debug().Str("mod", "store").Msg("AutoMigration is set to off, migration skipped")
		return nil
	}
	return s.DB.AutoMigrate(dst...)
}