MTP20_IF/db/mongodb.go

package db

import (
	"context"
	"encoding/hex"
	"fmt"
	"mtp2_if/config"
	"mtp2_if/packet"

	"github.com/qiniu/qmgo"
)

var (
	client *qmgo.Client
	db     *qmgo.Database
	ctx    context.Context
)

// InitMongoDB 初始化连接MongoDB
func InitMongoDB() error {
	// 创建链接
	// var err error
	// session, err = mgo.Dial(fmt.Sprintf("%s:%d", config.SerCfg.GetMongoDBHostname(), config.SerCfg.GetMongoDBPort()))
	// if err != nil {
	// 	return err
	// }

	// // 选择DB
	// mongodb = session.DB(config.SerCfg.GetMongoDBDBName())

	// 尝试解密
	var dbUser []byte
	ciphertext, _ := hex.DecodeString(config.SerCfg.GetMongoDBUsername())
	if len(ciphertext) > 8 {
		ciphertext = ciphertext[4 : len(ciphertext)-8]
		dbUser, _ = packet.Decrypt(ciphertext, packet.AESKey, true)
		if dbUser == nil {
			dbUser = []byte(config.SerCfg.GetMongoDBUsername())
		}
	} else {
		dbUser = []byte(config.SerCfg.GetMongoDBUsername())
	}
	var dbPwd []byte
	ciphertext, _ = hex.DecodeString(config.SerCfg.GetMongoDBPassword())
	if len(ciphertext) > 8 {
		ciphertext = ciphertext[4 : len(ciphertext)-8]
		dbPwd, _ = packet.Decrypt(ciphertext, packet.AESKey, true)
		if dbPwd == nil {
			dbPwd = []byte(config.SerCfg.GetMongoDBPassword())
		}
	} else {
		dbPwd = []byte(config.SerCfg.GetMongoDBPassword())
	}

	// // 登陆
	// if err := mongodb.Login(string(dbUser), string(dbPwd)); err != nil {
	// 	return err
	// }

	var err error

	ctx := context.Background()
	u := fmt.Sprintf("mongodb://%s:%s@%s:%d/?authSource=%s",
		string(dbUser),
		string(dbPwd),
		config.SerCfg.GetMongoDBHostname(),
		config.SerCfg.GetMongoDBPort(),
		config.SerCfg.GetMongoDBDBName(),
	)
	c, err := qmgo.NewClient(ctx, &qmgo.Config{Uri: u})
	if err != nil {
		return err
	}

	if err = c.Ping(10); err != nil {
		return err
	}

	client = c
	db = client.Database(config.SerCfg.GetMongoDBDBName())

	return nil
}

// GetMongoDB 获取MongoDB Database
func GetMongoDB() *qmgo.Database {
	return db
}

func GetMongoDBClient() *qmgo.Client {
	return client
}

// CloseMongoDB 关闭MongoDB连接
func CloseMongoDB() {
	client.Close(ctx)
}