76 lines
1.5 KiB
Go
76 lines
1.5 KiB
Go
package storage
|
|
|
|
import (
|
|
"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 {
|
|
db.AutoMigrate(tables...)
|
|
}
|
|
}
|
|
|
|
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]
|
|
}
|