model/model.go

package model

import (
	"strconv"
	"strings"
	"time"

	"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) {
	if defaultGormDB != nil {
		m.DB = defaultGormDB
	} else if db, err := gorm.Open(
		postgres.New(postgres.Config{
			DSN:                  defaultDsn,
			PreferSimpleProtocol: true,
		}), &gorm.Config{}); err != nil {
		panic("DATABASE_OPEN_ERROR")
	} else {
		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) {
	defaultOptions = []Option{}

	if defaultDb != nil{
		// database init
		arrConfStr := []string{
			"host=" + defaultDb.Host,
			"port=" + strconv.Itoa(defaultDb.Port),
			"user=" + defaultDb.Username,
			"password=" + defaultDb.Password,
			"dbname=" + defaultDb.Dbname,
			defaultDb.ExtraParameters,
		}
		defaultDsn = strings.Join(arrConfStr, " ")
		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.Debugf("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)
	}
}