mayfly-go/server/internal/db/dbm/dbm.go

package dbm

import (
	"context"
	"mayfly-go/internal/db/dbm/dbi"
	_ "mayfly-go/internal/db/dbm/dm"
	_ "mayfly-go/internal/db/dbm/mssql"
	_ "mayfly-go/internal/db/dbm/mysql"
	_ "mayfly-go/internal/db/dbm/oracle"
	_ "mayfly-go/internal/db/dbm/postgres"
	_ "mayfly-go/internal/db/dbm/sqlite"
	"mayfly-go/internal/machine/mcm"
	"mayfly-go/pkg/logx"
	"mayfly-go/pkg/pool"
)

var (
	poolGroup = pool.NewPoolGroup[*dbi.DbConn]()
)

func init() {
	mcm.AddCheckSshTunnelMachineUseFunc(func(machineId int) bool {
		items := poolGroup.AllPool()
		for _, v := range items {
			if v.Stats().TotalConns == 0 {
				continue // 连接池中没有连接，跳过
			}
			conn, err := v.Get(context.Background())
			if err != nil {
				continue // 获取连接失败，跳过
			}
			if conn.Info.SshTunnelMachineId == machineId {
				return true
			}
		}
		return false
	})
}

// GetDbConn 从连接池中获取连接信息
func GetDbConn(ctx context.Context, dbId uint64, database string, getDbInfo func() (*dbi.DbInfo, error)) (*dbi.DbConn, error) {
	connId := dbi.GetDbConnId(dbId, database)

	pool, err := poolGroup.GetCachePool(connId, func() (*dbi.DbConn, error) {
		// 若缓存中不存在，则从回调函数中获取DbInfo
		dbInfo, err := getDbInfo()
		if err != nil {
			return nil, err
		}
		logx.Debugf("dbm - conn create, connId: %s, dbInfo: %v", connId, dbInfo)
		// 连接数据库
		return Conn(context.Background(), dbInfo)
	})

	if err != nil {
		return nil, err
	}
	// 从连接池中获取一个可用的连接
	return pool.Get(ctx)
}

// 使用指定dbInfo信息进行连接
func Conn(ctx context.Context, di *dbi.DbInfo) (*dbi.DbConn, error) {
	return di.Conn(ctx, dbi.GetMeta(di.Type))
}

// 根据实例id获取连接
func GetDbConnByInstanceId(ctx context.Context, instanceId uint64) *dbi.DbConn {
	for _, pool := range poolGroup.AllPool() {
		conn, err := pool.Get(ctx)
		if err != nil {
			continue
		}
		if conn.Info.InstanceId == instanceId {
			return conn
		}
	}
	return nil
}

// 删除db缓存并关闭该数据库所有连接
func CloseDb(dbId uint64, db string) {
	poolGroup.Close(dbi.GetDbConnId(dbId, db))
}