refactor: db/redis/mongo连接代码包独立

2025-11-04 08:20:25 +08:00 · 2023-10-27 17:41:45 +08:00
parent a1303b52eb
commit 12f63ef3dd
45 changed files with 1112 additions and 950 deletions
--- a/server/internal/mongo/api/mongo.go
+++ b/server/internal/mongo/api/mongo.go
@@ -74,18 +74,20 @@ func (m *Mongo) DeleteMongo(rc *req.Ctx) {
 }

 func (m *Mongo) Databases(rc *req.Ctx) {
-	cli := m.MongoApp.GetMongoInst(m.GetMongoId(rc.GinCtx)).Cli
-	res, err := cli.ListDatabases(context.TODO(), bson.D{})
+	conn, err := m.MongoApp.GetMongoConn(m.GetMongoId(rc.GinCtx))
+	biz.ErrIsNil(err)
+	res, err := conn.Cli.ListDatabases(context.TODO(), bson.D{})
 	biz.ErrIsNilAppendErr(err, "获取mongo所有库信息失败: %s")
 	rc.ResData = res
 }

 func (m *Mongo) Collections(rc *req.Ctx) {
-	cli := m.MongoApp.GetMongoInst(m.GetMongoId(rc.GinCtx)).Cli
+	conn, err := m.MongoApp.GetMongoConn(m.GetMongoId(rc.GinCtx))
+	biz.ErrIsNil(err)
 	db := rc.GinCtx.Query("database")
 	biz.NotEmpty(db, "database不能为空")
 	ctx := context.TODO()
-	res, err := cli.Database(db).ListCollectionNames(ctx, bson.D{})
+	res, err := conn.Cli.Database(db).ListCollectionNames(ctx, bson.D{})
 	biz.ErrIsNilAppendErr(err, "获取库集合信息失败: %s")
 	rc.ResData = res
 }
@@ -94,8 +96,9 @@ func (m *Mongo) RunCommand(rc *req.Ctx) {
 	commandForm := new(form.MongoRunCommand)
 	ginx.BindJsonAndValid(rc.GinCtx, commandForm)

-	inst := m.MongoApp.GetMongoInst(m.GetMongoId(rc.GinCtx))
-	rc.ReqParam = collx.Kvs("mongo", inst.Info, "cmd", commandForm)
+	conn, err := m.MongoApp.GetMongoConn(m.GetMongoId(rc.GinCtx))
+	biz.ErrIsNil(err)
+	rc.ReqParam = collx.Kvs("mongo", conn.Info, "cmd", commandForm)

 	// 顺序执行
 	commands := bson.D{}
@@ -110,7 +113,7 @@ func (m *Mongo) RunCommand(rc *req.Ctx) {

 	ctx := context.TODO()
 	var bm bson.M
-	err := inst.Cli.Database(commandForm.Database).RunCommand(
+	err = conn.Cli.Database(commandForm.Database).RunCommand(
 		ctx,
 		commands,
 	).Decode(&bm)
@@ -121,10 +124,13 @@ func (m *Mongo) RunCommand(rc *req.Ctx) {

 func (m *Mongo) FindCommand(rc *req.Ctx) {
 	g := rc.GinCtx
-	cli := m.MongoApp.GetMongoInst(m.GetMongoId(g)).Cli
 	commandForm := new(form.MongoFindCommand)
 	ginx.BindJsonAndValid(g, commandForm)

+	conn, err := m.MongoApp.GetMongoConn(m.GetMongoId(g))
+	biz.ErrIsNil(err)
+	cli := conn.Cli
+
 	limit := commandForm.Limit
 	if limit != 0 {
 		biz.IsTrue(limit <= 100, "limit不能超过100")
@@ -157,8 +163,9 @@ func (m *Mongo) UpdateByIdCommand(rc *req.Ctx) {
 	commandForm := new(form.MongoUpdateByIdCommand)
 	ginx.BindJsonAndValid(g, commandForm)

-	inst := m.MongoApp.GetMongoInst(m.GetMongoId(g))
-	rc.ReqParam = collx.Kvs("mongo", inst.Info, "cmd", commandForm)
+	conn, err := m.MongoApp.GetMongoConn(m.GetMongoId(g))
+	biz.ErrIsNil(err)
+	rc.ReqParam = collx.Kvs("mongo", conn.Info, "cmd", commandForm)

 	// 解析docId文档id，如果为string类型则使用ObjectId解析，解析失败则为普通字符串
 	docId := commandForm.DocId
@@ -170,7 +177,7 @@ func (m *Mongo) UpdateByIdCommand(rc *req.Ctx) {
 		}
 	}

-	res, err := inst.Cli.Database(commandForm.Database).Collection(commandForm.Collection).UpdateByID(context.TODO(), docId, commandForm.Update)
+	res, err := conn.Cli.Database(commandForm.Database).Collection(commandForm.Collection).UpdateByID(context.TODO(), docId, commandForm.Update)
 	biz.ErrIsNilAppendErr(err, "命令执行失败: %s")

 	rc.ResData = res
@@ -181,8 +188,9 @@ func (m *Mongo) DeleteByIdCommand(rc *req.Ctx) {
 	commandForm := new(form.MongoUpdateByIdCommand)
 	ginx.BindJsonAndValid(g, commandForm)

-	inst := m.MongoApp.GetMongoInst(m.GetMongoId(g))
-	rc.ReqParam = collx.Kvs("mongo", inst.Info, "cmd", commandForm)
+	conn, err := m.MongoApp.GetMongoConn(m.GetMongoId(g))
+	biz.ErrIsNil(err)
+	rc.ReqParam = collx.Kvs("mongo", conn.Info, "cmd", commandForm)

 	// 解析docId文档id，如果为string类型则使用ObjectId解析，解析失败则为普通字符串
 	docId := commandForm.DocId
@@ -194,7 +202,7 @@ func (m *Mongo) DeleteByIdCommand(rc *req.Ctx) {
 		}
 	}

-	res, err := inst.Cli.Database(commandForm.Database).Collection(commandForm.Collection).DeleteOne(context.TODO(), bson.D{{Key: "_id", Value: docId}})
+	res, err := conn.Cli.Database(commandForm.Database).Collection(commandForm.Collection).DeleteOne(context.TODO(), bson.D{{Key: "_id", Value: docId}})
 	biz.ErrIsNilAppendErr(err, "命令执行失败: %s")
 	rc.ResData = res
 }
@@ -204,10 +212,11 @@ func (m *Mongo) InsertOneCommand(rc *req.Ctx) {
 	commandForm := new(form.MongoInsertCommand)
 	ginx.BindJsonAndValid(g, commandForm)

-	inst := m.MongoApp.GetMongoInst(m.GetMongoId(g))
-	rc.ReqParam = collx.Kvs("mongo", inst.Info, "cmd", commandForm)
+	conn, err := m.MongoApp.GetMongoConn(m.GetMongoId(g))
+	biz.ErrIsNil(err)
+	rc.ReqParam = collx.Kvs("mongo", conn.Info, "cmd", commandForm)

-	res, err := inst.Cli.Database(commandForm.Database).Collection(commandForm.Collection).InsertOne(context.TODO(), commandForm.Doc)
+	res, err := conn.Cli.Database(commandForm.Database).Collection(commandForm.Collection).InsertOne(context.TODO(), commandForm.Doc)
 	biz.ErrIsNilAppendErr(err, "命令执行失败: %s")
 	rc.ResData = res
 }
--- a/server/internal/mongo/application/mongo.go
+++ b/server/internal/mongo/application/mongo.go
@@ -1,26 +1,12 @@
 package application

 import (
-	"context"
-	"fmt"
-	"mayfly-go/internal/common/consts"
-	machineapp "mayfly-go/internal/machine/application"
-	"mayfly-go/internal/machine/infrastructure/machine"
 	"mayfly-go/internal/mongo/domain/entity"
 	"mayfly-go/internal/mongo/domain/repository"
+	"mayfly-go/internal/mongo/mgm"
 	"mayfly-go/pkg/base"
-	"mayfly-go/pkg/biz"
-	"mayfly-go/pkg/cache"
-	"mayfly-go/pkg/logx"
+	"mayfly-go/pkg/errorx"
 	"mayfly-go/pkg/model"
-	"mayfly-go/pkg/utils/netx"
-	"mayfly-go/pkg/utils/structx"
-	"net"
-	"regexp"
-	"time"
-
-	"go.mongodb.org/mongo-driver/mongo"
-	"go.mongodb.org/mongo-driver/mongo/options"
 )

 type Mongo interface {
@@ -38,7 +24,7 @@ type Mongo interface {

 	// 获取mongo连接实例
 	// @param id mongo id
-	GetMongoInst(id uint64) *MongoInstance
+	GetMongoConn(id uint64) (*mgm.MongoConn, error)
 }

 func newMongoAppImpl(mongoRepo repository.Mongo) Mongo {
@@ -61,7 +47,7 @@ func (d *mongoAppImpl) Count(condition *entity.MongoQuery) int64 {
 }

 func (d *mongoAppImpl) Delete(id uint64) error {
-	DeleteMongoCache(id)
+	mgm.CloseConn(id)
 	return d.GetRepo().DeleteById(id)
 }

@@ -71,148 +57,16 @@ func (d *mongoAppImpl) Save(m *entity.Mongo) error {
 	}

 	// 先关闭连接
-	DeleteMongoCache(m.Id)
+	mgm.CloseConn(m.Id)
 	return d.GetRepo().UpdateById(m)
 }

-func (d *mongoAppImpl) GetMongoInst(id uint64) *MongoInstance {
-	mongoInstance, err := GetMongoInstance(id, func(u uint64) (*entity.Mongo, error) {
-		mongo, err := d.GetById(new(entity.Mongo), u)
+func (d *mongoAppImpl) GetMongoConn(id uint64) (*mgm.MongoConn, error) {
+	return mgm.GetMongoConn(id, func() (*mgm.MongoInfo, error) {
+		mongo, err := d.GetById(new(entity.Mongo), id)
 		if err != nil {
-			return nil, err
+			return nil, errorx.NewBiz("mongo信息不存在")
 		}
-		return mongo, nil
-	})
-	biz.ErrIsNilAppendErr(err, "连接mongo失败: %s")
-	return mongoInstance
-}
-
-// -----------------------------------------------------------
-
-// mongo客户端连接缓存，指定时间内没有访问则会被关闭
-var mongoCliCache = cache.NewTimedCache(consts.MongoConnExpireTime, 5*time.Second).
-	WithUpdateAccessTime(true).
-	OnEvicted(func(key any, value any) {
-		logx.Infof("删除mongo连接缓存: id = %v", key)
-		value.(*MongoInstance).Close()
-	})
-
-func init() {
-	machine.AddCheckSshTunnelMachineUseFunc(func(machineId int) bool {
-		// 遍历所有mongo连接实例，若存在redis实例使用该ssh隧道机器，则返回true，表示还在使用中...
-		items := mongoCliCache.Items()
-		for _, v := range items {
-			if v.Value.(*MongoInstance).Info.SshTunnelMachineId == machineId {
-				return true
-			}
-		}
-		return false
+		return mongo.ToMongoInfo(), nil
 	})
 }
-
-// 获取mongo的连接实例
-func GetMongoInstance(mongoId uint64, getMongoEntity func(uint64) (*entity.Mongo, error)) (*MongoInstance, error) {
-	mi, err := mongoCliCache.ComputeIfAbsent(mongoId, func(_ any) (any, error) {
-		mongoEntity, err := getMongoEntity(mongoId)
-		if err != nil {
-			return nil, err
-		}
-
-		c, err := connect(mongoEntity)
-		if err != nil {
-			return nil, err
-		}
-		return c, nil
-	})
-
-	if mi != nil {
-		return mi.(*MongoInstance), err
-	}
-	return nil, err
-}
-
-func DeleteMongoCache(mongoId uint64) {
-	mongoCliCache.Delete(mongoId)
-}
-
-type MongoInfo struct {
-	Id                 uint64 `json:"id"`
-	Name               string `json:"name"`
-	TagPath            string `json:"tagPath"`
-	SshTunnelMachineId int    `json:"-"` // ssh隧道机器id
-}
-
-func (m *MongoInfo) GetLogDesc() string {
-	return fmt.Sprintf("Mongo[id=%d, tag=%s, name=%s]", m.Id, m.TagPath, m.Name)
-}
-
-type MongoInstance struct {
-	Id   uint64
-	Info *MongoInfo
-
-	Cli *mongo.Client
-}
-
-func (mi *MongoInstance) Close() {
-	if mi.Cli != nil {
-		if err := mi.Cli.Disconnect(context.Background()); err != nil {
-			logx.Errorf("关闭mongo实例[%d]连接失败: %s", mi.Id, err)
-		}
-		mi.Cli = nil
-	}
-}
-
-// 连接mongo，并返回client
-func connect(me *entity.Mongo) (*MongoInstance, error) {
-	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
-	defer cancel()
-
-	mongoInstance := &MongoInstance{Id: me.Id, Info: toMongoInfo(me)}
-
-	mongoOptions := options.Client().ApplyURI(me.Uri).
-		SetMaxPoolSize(1)
-	// 启用ssh隧道则连接隧道机器
-	if me.SshTunnelMachineId > 0 {
-		mongoOptions.SetDialer(&MongoSshDialer{machineId: me.SshTunnelMachineId})
-	}
-
-	client, err := mongo.Connect(ctx, mongoOptions)
-	if err != nil {
-		mongoInstance.Close()
-		return nil, err
-	}
-	if err = client.Ping(context.TODO(), nil); err != nil {
-		mongoInstance.Close()
-		return nil, err
-	}
-
-	logx.Infof("连接mongo: %s", func(str string) string {
-		reg := regexp.MustCompile(`(^mongodb://.+?:)(.+)(@.+$)`)
-		return reg.ReplaceAllString(str, `${1}****${3}`)
-	}(me.Uri))
-	mongoInstance.Cli = client
-	return mongoInstance, err
-}
-
-func toMongoInfo(me *entity.Mongo) *MongoInfo {
-	mi := new(MongoInfo)
-	structx.Copy(mi, me)
-	return mi
-}
-
-type MongoSshDialer struct {
-	machineId int
-}
-
-func (sd *MongoSshDialer) DialContext(ctx context.Context, network, address string) (net.Conn, error) {
-	stm, err := machineapp.GetMachineApp().GetSshTunnelMachine(sd.machineId)
-	if err != nil {
-		return nil, err
-	}
-	if sshConn, err := stm.GetDialConn(network, address); err == nil {
-		// 将ssh conn包装，否则内部部设置超时会报错,ssh conn不支持设置超时会返回错误: ssh: tcpChan: deadline not supported
-		return &netx.WrapSshConn{Conn: sshConn}, nil
-	} else {
-		return nil, err
-	}
-}
--- a/server/internal/mongo/domain/entity/mongo.go
+++ b/server/internal/mongo/domain/entity/mongo.go
@@ -1,6 +1,10 @@
 package entity

-import "mayfly-go/pkg/model"
+import (
+	"mayfly-go/internal/mongo/mgm"
+	"mayfly-go/pkg/model"
+	"mayfly-go/pkg/utils/structx"
+)

 type Mongo struct {
 	model.Model
@@ -11,3 +15,10 @@ type Mongo struct {
 	TagId              uint64 `json:"tagId"`
 	TagPath            string `json:"tagPath"`
 }
+
+// 转换为mongoInfo进行连接
+func (me *Mongo) ToMongoInfo() *mgm.MongoInfo {
+	mongoInfo := new(mgm.MongoInfo)
+	structx.Copy(mongoInfo, me)
+	return mongoInfo
+}
--- a/server/internal/mongo/mgm/conn.go
+++ b/server/internal/mongo/mgm/conn.go
@@ -0,0 +1,24 @@
+package mgm
+
+import (
+	"context"
+	"mayfly-go/pkg/logx"
+
+	"go.mongodb.org/mongo-driver/mongo"
+)
+
+type MongoConn struct {
+	Id   string
+	Info *MongoInfo
+
+	Cli *mongo.Client
+}
+
+func (mc *MongoConn) Close() {
+	if mc.Cli != nil {
+		if err := mc.Cli.Disconnect(context.Background()); err != nil {
+			logx.Errorf("关闭mongo实例[%s]连接失败: %s", mc.Id, err)
+		}
+		mc.Cli = nil
+	}
+}
--- a/server/internal/mongo/mgm/conn_cache.go
+++ b/server/internal/mongo/mgm/conn_cache.go
@@ -0,0 +1,72 @@
+package mgm
+
+import (
+	"mayfly-go/internal/common/consts"
+	"mayfly-go/internal/machine/infrastructure/machine"
+	"mayfly-go/pkg/cache"
+	"mayfly-go/pkg/logx"
+	"sync"
+	"time"
+)
+
+// mongo客户端连接缓存，指定时间内没有访问则会被关闭
+var connCache = cache.NewTimedCache(consts.MongoConnExpireTime, 5*time.Second).
+	WithUpdateAccessTime(true).
+	OnEvicted(func(key any, value any) {
+		logx.Infof("删除mongo连接缓存: id = %v", key)
+		value.(*MongoConn).Close()
+	})
+
+func init() {
+	machine.AddCheckSshTunnelMachineUseFunc(func(machineId int) bool {
+		// 遍历所有mongo连接实例，若存在redis实例使用该ssh隧道机器，则返回true，表示还在使用中...
+		items := connCache.Items()
+		for _, v := range items {
+			if v.Value.(*MongoConn).Info.SshTunnelMachineId == machineId {
+				return true
+			}
+		}
+		return false
+	})
+}
+
+var mutex sync.Mutex
+
+// 从缓存中获取mongo连接信息, 若缓存中不存在则会使用回调函数获取mongoInfo进行连接并缓存
+func GetMongoConn(mongoId uint64, getMongoInfo func() (*MongoInfo, error)) (*MongoConn, error) {
+	connId := getConnId(mongoId)
+
+	// connId不为空，则为需要缓存
+	needCache := connId != ""
+	if needCache {
+		load, ok := connCache.Get(connId)
+		if ok {
+			return load.(*MongoConn), nil
+		}
+	}
+
+	mutex.Lock()
+	defer mutex.Unlock()
+
+	// 若缓存中不存在，则从回调函数中获取MongoInfo
+	mi, err := getMongoInfo()
+	if err != nil {
+		return nil, err
+	}
+
+	// 连接mongo
+	mc, err := mi.Conn()
+	if err != nil {
+		return nil, err
+	}
+
+	if needCache {
+		connCache.Put(connId, mc)
+	}
+	return mc, nil
+}
+
+// 关闭连接，并移除缓存连接
+func CloseConn(mongoId uint64) {
+	connCache.Delete(mongoId)
+}
--- a/server/internal/mongo/mgm/info.go
+++ b/server/internal/mongo/mgm/info.go
@@ -0,0 +1,79 @@
+package mgm
+
+import (
+	"context"
+	"fmt"
+	"mayfly-go/pkg/logx"
+	"mayfly-go/pkg/utils/netx"
+	"net"
+	"regexp"
+	"time"
+
+	machineapp "mayfly-go/internal/machine/application"
+
+	"go.mongodb.org/mongo-driver/mongo"
+	"go.mongodb.org/mongo-driver/mongo/options"
+)
+
+type MongoInfo struct {
+	Id   uint64 `json:"id"`
+	Name string `json:"name"`
+
+	Uri string `json:"-"`
+
+	TagPath            string `json:"tagPath"`
+	SshTunnelMachineId int    `json:"-"` // ssh隧道机器id
+}
+
+func (mi *MongoInfo) Conn() (*MongoConn, error) {
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	mongoOptions := options.Client().ApplyURI(mi.Uri).
+		SetMaxPoolSize(1)
+	// 启用ssh隧道则连接隧道机器
+	if mi.SshTunnelMachineId > 0 {
+		mongoOptions.SetDialer(&MongoSshDialer{machineId: mi.SshTunnelMachineId})
+	}
+
+	client, err := mongo.Connect(ctx, mongoOptions)
+	if err != nil {
+		return nil, err
+	}
+	if err = client.Ping(context.TODO(), nil); err != nil {
+		client.Disconnect(ctx)
+		return nil, err
+	}
+
+	logx.Infof("连接mongo: %s", func(str string) string {
+		reg := regexp.MustCompile(`(^mongodb://.+?:)(.+)(@.+$)`)
+		return reg.ReplaceAllString(str, `${1}****${3}`)
+	}(mi.Uri))
+
+	return &MongoConn{Id: getConnId(mi.Id), Info: mi, Cli: client}, nil
+}
+
+type MongoSshDialer struct {
+	machineId int
+}
+
+func (sd *MongoSshDialer) DialContext(ctx context.Context, network, address string) (net.Conn, error) {
+	stm, err := machineapp.GetMachineApp().GetSshTunnelMachine(sd.machineId)
+	if err != nil {
+		return nil, err
+	}
+	if sshConn, err := stm.GetDialConn(network, address); err == nil {
+		// 将ssh conn包装，否则内部部设置超时会报错,ssh conn不支持设置超时会返回错误: ssh: tcpChan: deadline not supported
+		return &netx.WrapSshConn{Conn: sshConn}, nil
+	} else {
+		return nil, err
+	}
+}
+
+// 生成mongo连接id
+func getConnId(id uint64) string {
+	if id == 0 {
+		return ""
+	}
+	return fmt.Sprintf("%d", id)
+}