81 lines
1.7 KiB
Go
81 lines
1.7 KiB
Go
package storage
|
|
|
|
import (
|
|
"context"
|
|
|
|
"github.com/nose7en/ToyBoomServer/common"
|
|
"gorm.io/gorm"
|
|
)
|
|
|
|
const (
|
|
DefaultDBName = "default"
|
|
)
|
|
|
|
type DBManager interface {
|
|
GetDB(dbType, dbName string) *gorm.DB
|
|
GetDefaultDB() *gorm.DB
|
|
SetDB(dbType, dbName string, db *gorm.DB)
|
|
RemoveDB(dbType, dbName string)
|
|
Init(tables ...interface{})
|
|
}
|
|
|
|
type dbManagerImpl struct {
|
|
DBs map[string]map[string]*gorm.DB // [key: dbname, value: dbs] => [key: dbtype value: db]
|
|
defaultDBType string
|
|
}
|
|
|
|
func (dbm *dbManagerImpl) Init(tables ...interface{}) {
|
|
dbs := dbm.DBs[DefaultDBName]
|
|
for _, db := range dbs {
|
|
if err := db.AutoMigrate(tables...); err != nil {
|
|
common.Logger(context.Background()).WithError(err).Fatalf("auto migrate error!!!")
|
|
}
|
|
}
|
|
}
|
|
|
|
var (
|
|
dbm *dbManagerImpl
|
|
)
|
|
|
|
func MustInitDBManager(dbs map[string]map[string]*gorm.DB, defaultDBType string) {
|
|
if dbm == nil {
|
|
dbm = NewDBManager(dbs, defaultDBType)
|
|
}
|
|
}
|
|
|
|
func NewDBManager(dbs map[string]map[string]*gorm.DB, defaultDBType string) *dbManagerImpl {
|
|
if dbs == nil {
|
|
dbs = make(map[string]map[string]*gorm.DB)
|
|
}
|
|
return &dbManagerImpl{
|
|
DBs: dbs,
|
|
defaultDBType: defaultDBType,
|
|
}
|
|
}
|
|
|
|
func GetDBManager() DBManager {
|
|
if dbm == nil {
|
|
dbm = NewDBManager(nil, "")
|
|
}
|
|
return dbm
|
|
}
|
|
|
|
func (dbm *dbManagerImpl) GetDB(dbType, dbName string) *gorm.DB {
|
|
return dbm.DBs[dbName][dbType]
|
|
}
|
|
|
|
func (dbm *dbManagerImpl) SetDB(dbType, dbName string, db *gorm.DB) {
|
|
if dbm.DBs[dbName] == nil {
|
|
dbm.DBs[dbName] = make(map[string]*gorm.DB)
|
|
}
|
|
dbm.DBs[dbName][dbType] = db
|
|
}
|
|
|
|
func (dbm *dbManagerImpl) RemoveDB(dbType, dbName string) {
|
|
delete(dbm.DBs[dbName], dbType)
|
|
}
|
|
|
|
func (dbm *dbManagerImpl) GetDefaultDB() *gorm.DB {
|
|
return dbm.DBs[DefaultDBName][dbm.defaultDBType]
|
|
}
|