model/model.go

package model

import (
	"strings"
	"time"

	"gorm.io/driver/mysql"
	"gorm.io/driver/postgres"
	"gorm.io/gorm"
	"qoobing.com/gomod/database"
	"qoobing.com/gomod/log"
	"qoobing.com/gomod/redis"
	"qoobing.com/gomod/redis/sentinel"
)

type Model struct {
	DB         *gorm.DB
	DbTxStatus DbTxStatus
	Redis      redis.Conn
	RedisPool  *redis.Pool
}

type DbTxStatus int

const (
	DBTX_STATUS_TX_NONE    DbTxStatus = 0
	DBTX_STATUS_TX_DOING   DbTxStatus = 1
	DBTX_STATUS_TX_SUCCESS DbTxStatus = 2
	DBTX_STATUS_TX_FAILED  DbTxStatus = 3
)

var (
	defaultDb         *database.Config
	defaultDsn        string
	defaultDbDebug    bool
	defaultGormDB     *gorm.DB
	defaultRds        *redis.Config
	defaultRedis      *redis.Pool
	defaultRedisDebug bool
	defaultOptions    []Option
)

func NewModel() *Model {
	if len(defaultOptions) == 0 {
		panic("No defualt Database&Redis been configed")
	}
	return NewModelWithOption(defaultOptions...)
}

func NewModelDefault() *Model {
	return NewModelWithOption(OptOpenDefaultDatabase, OptOpenDefaultRedis)
}

func NewModelDefaultRedis() *Model {
	return NewModelWithOption(OptOpenDefaultRedis)
}

func NewModelDefaultDatabase() *Model {
	return NewModelWithOption(OptOpenDefaultDatabase)
}

func NewModelWithOption(options ...Option) *Model {
	n := 0
	m := Model{}
	//log.Debugf("Start NewModelWithOption...")
	for _, option := range options {
		n++
		option(&m)
	}
	//log.Debugf("Finish NewModelWithOption(with %d options)", n)
	return &m
}

func (m *Model) Close() {
	if m.Redis != nil {
		m.Redis.Close()
	}
	if m.DB == nil {
		m.DbTxStatus = DBTX_STATUS_TX_NONE
	} else if m.DbTxStatus == DBTX_STATUS_TX_DOING {
		m.DB.Rollback()
		m.DbTxStatus = DBTX_STATUS_TX_FAILED
	}
}

func (m *Model) Begin() {
	if m.DB == nil {
		panic("unreachable code, m.DB is uninitialized")
	} else if m.DbTxStatus != DBTX_STATUS_TX_NONE {
		panic("unreachable code, begin transaction towice???")
	} else {
		m.DbTxStatus = DBTX_STATUS_TX_DOING
		m.DB = m.DB.Begin()
	}
}

func (m *Model) Commit() error {
	if m.DB == nil {
		panic("unreachable code, m.DB is uninitialized")
	} else if m.DbTxStatus != DBTX_STATUS_TX_DOING {
		return nil
	} else if err := m.DB.Commit().Error; err != nil {
		m.DbTxStatus = DBTX_STATUS_TX_FAILED
		return err
	} else {
		m.DbTxStatus = DBTX_STATUS_TX_SUCCESS
		return nil
	}
}

// Option is model's option
type Option func(*Model)

// OptOpenDefaultDatabase option function for open default database
func OptOpenDefaultDatabase(m *Model) {
	var err error
	if defaultDb == nil {
		panic("defaultDb not init")
	}

	if defaultGormDB != nil {
		m.DB = defaultGormDB
		//TODO: check cocurrent
		m.DB = m.DB.Session(&gorm.Session{QueryFields: true})
		//log.Debugf("Opt for open default database done")
		return
	}

	var dialector gorm.Dialector
	var gormconfig = &gorm.Config{
		NowFunc: func() time.Time {
			return time.Now().UTC()
		},
	}
	switch defaultDb.Type {
	case "mysql":
		dialector = mysql.New(
			mysql.Config{
				DSN:               defaultDsn,
				DefaultStringSize: 512,
			},
		)
	case "pgsql":
		dialector = postgres.New(
			postgres.Config{
				DSN:                  defaultDsn,
				PreferSimpleProtocol: true,
			},
		)
	default:
		panic("UNKNOWN DATABASE TYPE:" + defaultDb.Type)
	}

	var db *gorm.DB
	db, err = gorm.Open(dialector, gormconfig)
	if err != nil {
		panic("DATABASE_OPEN_ERROR")
	}

	if defaultDbDebug {
		m.DB = db.Debug()
	} else {
		m.DB = db
	}
	if sqlDB, err := m.DB.DB(); err != nil {
		panic("DATABASE_OPEN_ERROR")
	} else {
		sqlDB.SetMaxIdleConns(3)
		sqlDB.SetMaxOpenConns(10)
		sqlDB.SetConnMaxLifetime(time.Minute)
	}
	defaultGormDB = m.DB

	//TODO: check cocurrent
	m.DB = m.DB.Session(&gorm.Session{QueryFields: true})
	//log.Debugf("Opt for open default database done")
	return
}

// OptOpenDefaultRedisSentinelPool option function for open default redis sentinel
func OptOpenDefaultRedis(m *Model) {
	if defaultRedis != nil {
		m.Redis = defaultRedis.Get()
		m.RedisPool = defaultRedis
	} else if pool := sentinel.NewPool(*defaultRds); pool == nil {
		panic("REDIS_ERROR")
	} else {
		defaultRedis = pool
		m.Redis = defaultRedis.Get()
		m.RedisPool = defaultRedis
	}

	//log.Debugf("Opt for open default redis done")
	return
}

// Init init default database config & redis config
func Init(_defaultDb *database.Config, _defaultRds *redis.Config) {
	defaultDb = _defaultDb
	defaultRds = _defaultRds
	defaultOptions = []Option{}

	if defaultDb != nil {
		// database init
		defaultDsn = database.GetDsn(defaultDb)
		defaultDbDebug = defaultDb.Debug

		// database debug
		if defaultDbDebug && len(defaultDb.Password) > 5 {
			p := defaultDb.Password
			l := len(p)
			secDefaultDsn := strings.Replace(defaultDsn, p[2:l-3], "*****", 1)
			log.Infof("defaultDsn='%s'", secDefaultDsn)
		}

		defaultOptions = append(defaultOptions, OptOpenDefaultDatabase)
	}

	if defaultRds != nil {
		// redis init
		defaultRds = defaultRds
		defaultRedisDebug = defaultRds.Debug

		// redis debug
		if defaultRedisDebug && len(defaultRds.Password) > 5 {
			p := defaultRds.Password
			l := len(p)
			secDefaultRds := *defaultRds
			secDefaultRds.Password = strings.Replace(p, p[2:l-3], "*****", 1)
			log.PrintPretty("defaultRds=", secDefaultRds)
		}
		defaultOptions = append(defaultOptions, OptOpenDefaultRedis)
	}
}