cas_tt_cloud_backend/serve/orm/orm.go

package orm

import (
	"SciencesServer/serve/orm/logic"
	"fmt"
	"log"
	"os"
	"time"

	"gorm.io/gorm/logger"
	"gorm.io/gorm/schema"

	"gorm.io/gorm"
)

var (
	orm *gorm.DB
)

type Instance struct {
	Engine *gorm.DB

	debug                                   bool
	dbMode                                  string
	tablePrefix                             string
	singularTable                           bool
	maxIdleConns, maxOpenConns, maxLifetime int
	mysql                                   *logic.Mysql
	sqlite                                  *logic.Sqlite
}

type Option func(instance *Instance)

func WithDebug(debug bool) Option {
	return func(instance *Instance) {
		instance.debug = debug
	}
}

func WithDBMode(dbMode string) Option {
	return func(instance *Instance) {
		instance.dbMode = dbMode
	}
}

func WithTablePrefix(tablePrefix string) Option {
	return func(instance *Instance) {
		instance.tablePrefix = tablePrefix
	}
}

func WithSingularTable(singularTable bool) Option {
	return func(instance *Instance) {
		instance.singularTable = singularTable
	}
}

func WithMaxIdleConns(maxIdleConns int) Option {
	return func(instance *Instance) {
		instance.maxIdleConns = maxIdleConns
	}
}

func WithMaxOpenConns(maxOpenConns int) Option {
	return func(instance *Instance) {
		instance.maxOpenConns = maxOpenConns
	}
}

func WithMaxLifetime(maxLifetime int) Option {
	return func(instance *Instance) {
		instance.maxLifetime = maxLifetime
	}
}

func WithMysqlOption(option *logic.Mysql) Option {
	return func(instance *Instance) {
		instance.mysql = option
	}
}

func WithSqliteOption(option *logic.Sqlite) Option {
	return func(instance *Instance) {
		instance.sqlite = option
	}
}

func (this *Instance) Init() *Instance {
	var engine logic.IEngine

	switch this.dbMode {
	case "mysql":
		engine = this.mysql
		break
	case "sqlite":
		engine = this.sqlite
		break
	default:
		panic(fmt.Sprintf("Unknown Engine Mode：%d", this.dbMode))
	}
	option := &gorm.Config{
		DisableForeignKeyConstraintWhenMigrating: true,
		NamingStrategy: schema.NamingStrategy{
			TablePrefix:   this.tablePrefix,
			SingularTable: this.singularTable,
		},
	}
	if this.debug {
		option.Logger = logger.New(
			log.New(os.Stdout, "\r\n", log.LstdFlags),
			logger.Config{
				SlowThreshold:             time.Second,
				LogLevel:                  logger.Info,
				Colorful:                  false,
				IgnoreRecordNotFoundError: true,
			},
		)
	}
	db, err := gorm.Open(engine.DSN(), option)

	if err != nil {
		panic("Orm Init Error：" + err.Error())
	}
	_db, _ := db.DB()
	_db.SetMaxIdleConns(this.maxIdleConns)
	_db.SetMaxOpenConns(this.maxOpenConns)
	_db.SetConnMaxLifetime(time.Duration(this.maxLifetime) * time.Second)

	this.Engine = db
	return this
}

func (this *Instance) Local() *Instance {
	orm = this.Engine
	return this
}

func NewInstance(option ...Option) *Instance {
	instance := new(Instance)

	for _, v := range option {
		v(instance)
	}
	return instance
}

func GetDB() *gorm.DB {
	return orm
}