kpt-temporary-mqtt/store/kptstore/rw_store.go

package kptstore

import (
	"kpt-temporary-mqtt/config"
	"time"

	"gitee.com/xuyiping_admin/pkg/logger/logrus"
	"gitee.com/xuyiping_admin/pkg/xerr"

	"gorm.io/driver/mysql"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"

	"gitee.com/xuyiping_admin/pkg/di"
)

var Module = di.Options(
	di.Provide(MustNewStore),
)

type DB struct {
	*gorm.DB
}

func NewStore(engine *gorm.DB) *DB {
	return &DB{engine}
}

type goRmLog struct {
}

func (g goRmLog) Printf(s string, i ...interface{}) {
	logrus.Infof(s, i...)
}

var newLogger = logger.New(
	goRmLog{},
	logger.Config{
		SlowThreshold: 5 * time.Second,
		LogLevel:      logger.Info,
	},
)

func MustNewStore(cfg *config.AppConfig) *DB {
	db, err := gorm.Open(mysql.New(mysql.Config{
		DriverName: cfg.StoreSetting.DriverName,
		DSN:        cfg.StoreSetting.KptRW}),
		&gorm.Config{Logger: newLogger},
	)
	if err != nil {
		panic(xerr.WithStack(err))
	}

	if cfg.StoreSetting.ShowSQL {
		db.Logger.LogMode(logger.Info)
	}

	return NewStore(db)
}

func MustMigrateStore(cfg *config.AppConfig) *gorm.DB {
	db, err := gorm.Open(mysql.New(mysql.Config{
		DriverName: cfg.StoreSetting.DriverName,
		DSN:        cfg.StoreSetting.KptMigr}),
		&gorm.Config{Logger: newLogger},
	)
	if err != nil {
		panic(xerr.WithStack(err))
	}

	if cfg.StoreSetting.ShowSQL {
		db.Logger.LogMode(logger.Info)
	}
	return db
}