server/components/mongo/component.go

package mhayaMongo

import (
	"context"
	"fmt"
	"time"

	cfacade "github.com/mhaya/facade"
	clog "github.com/mhaya/logger"
	cprofile "github.com/mhaya/profile"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
	"go.mongodb.org/mongo-driver/mongo/readpref"
)

const (
	Name = "mongo_component"
)

type (
	Component struct {
		cfacade.Component
		dbMap map[string]map[string]*mongo.Database
	}

	// HashDb hash by group id
	HashDb func(dbMaps map[string]*mongo.Database) string
)

func NewComponent() *Component {
	return &Component{
		dbMap: make(map[string]map[string]*mongo.Database),
	}
}

func (*Component) Name() string {
	return Name
}

func (s *Component) Init() {
	// load only the database contained in the `db_id_list`
	mongoIdList := s.App().Settings().Get("db_id_list")
	if mongoIdList.LastError() != nil || mongoIdList.Size() < 1 {
		clog.Warnf("[nodeId = %s] `mongo_id_list` property not exists.", s.App().NodeId())
		return
	}

	mongoConfig := cprofile.GetConfig("mongo")
	if mongoConfig.LastError() != nil {
		panic("`mongo` property not exists in profile file.")
	}

	for _, groupId := range mongoConfig.Keys() {
		s.dbMap[groupId] = make(map[string]*mongo.Database)

		dbGroup := mongoConfig.GetConfig(groupId)
		for i := 0; i < dbGroup.Size(); i++ {
			item := dbGroup.GetConfig(i)

			var (
				enable  = item.GetBool("enable", true)
				id      = item.GetString("db_id")
				dbName  = item.GetString("db_name")
				uri     = item.GetString("uri")
				timeout = time.Duration(item.GetInt64("timeout", 3)) * time.Second
			)

			for _, key := range mongoIdList.Keys() {
				dbId := mongoIdList.Get(key).ToString()
				if id != dbId {
					continue
				}

				if !enable {
					panic(fmt.Sprintf("[dbName = %s] is disabled!", dbName))
				}

				db, err := CreateDatabase(uri, dbName, timeout)
				if err != nil {
					panic(fmt.Sprintf("[dbName = %s] create mongodb fail. error = %s", dbName, err))
				}

				s.dbMap[groupId][id] = db
				clog.Infof("[dbGroup =%s, dbName = %s] is connected.", groupId, id)
			}
		}
	}
}

func CreateDatabase(uri, dbName string, timeout ...time.Duration) (*mongo.Database, error) {
	tt := 5 * time.Second

	if len(timeout) > 0 && timeout[0].Seconds() > 3 {
		tt = timeout[0]
	}

	// todo tls开启
	//tlsConfig := &tls.Config{
	//	MinVersion:               tls.VersionTLS12,
	//	PreferServerCipherSuites: true,
	//}
	//o := options.Client().ApplyURI(uri).SetTLSConfig(tlsConfig)

	o := options.Client().ApplyURI(uri).SetMaxPoolSize(500)
	if err := o.Validate(); err != nil {
		return nil, err
	}

	ctx, cancel := context.WithTimeout(context.Background(), tt)
	defer cancel()

	client, err := mongo.Connect(ctx, o)
	if err != nil {
		return nil, err
	}

	err = client.Ping(context.Background(), readpref.Primary())
	if err != nil {
		return nil, err
	}

	clog.Infof("ping db [uri = %s] is ok", uri)

	return client.Database(dbName), nil
}

func (s *Component) GetDb(id string) *mongo.Database {
	for _, group := range s.dbMap {
		for k, v := range group {
			if k == id {
				return v
			}
		}
	}
	return nil
}

func (s *Component) GetHashDb(groupId string, hashFn HashDb) (*mongo.Database, bool) {
	dbGroup, found := s.GetDbMap(groupId)
	if !found {
		clog.Warnf("groupId = %s not found.", groupId)
		return nil, false
	}

	dbId := hashFn(dbGroup)
	db, found := dbGroup[dbId]
	return db, found
}

func (s *Component) GetDbMap(groupId string) (map[string]*mongo.Database, bool) {
	dbGroup, found := s.dbMap[groupId]
	return dbGroup, found
}