init repo

storage/db.go

@@ -0,0 +1,75 @@
+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]
+}