gocommon/pkg/database/multi_db/cnn.go

package multi_db

import (
	"errors"
	"fmt"
	"sync"

	"git.shuncheng.lu/bigthing/gocommon/pkg/conf"
	"git.shuncheng.lu/bigthing/gocommon/pkg/database/common_db"
	"git.shuncheng.lu/bigthing/gocommon/pkg/internal/util"
)

var (
	notFoundError = func(dbName string) error {
		return errors.New(fmt.Sprintf("MultiDb not fount %s db", dbName))
	}
	lock     sync.RWMutex
	multiMap map[string]common_db.DbConn = map[string]common_db.DbConn{}
)

func Init() error {
	dbNameStr := conf.GetString("multi_db", "db_name")
	dbNames := util.SplitIgnoreSpace(dbNameStr, ",")
	for _, elem := range dbNames {
		conn, err := common_db.NewDbConn(elem)
		if err != nil {
			return err
		}
		setDb(elem, conn)
		util.Infof("[MultiDB] init db %s success !", elem)
	}
	return nil
}

func setDb(dbName string, conn common_db.DbConn) {
	lock.Lock()
	defer lock.Unlock()
	multiMap[dbName] = conn
}

func GetDb(dbName string) (common_db.DbConn, error) {
	lock.RLock()
	defer lock.RUnlock()
	db, isExist := multiMap[dbName]
	if isExist {
		return db, nil
	}
	return nil, notFoundError(dbName)
}