package dbclient import ( "context" "encoding/json" "fmt" "sync" "go-desk/internal/common" "go-desk/internal/crypto" "go-desk/internal/storage/models" ) // ConnectionPool 连接池管理器 type ConnectionPool struct { mysqlClients map[uint]*MySQLClient redisClients map[uint]*RedisClient mongoClients map[uint]*MongoClient mu sync.RWMutex } var ( globalPool *ConnectionPool poolOnce sync.Once ) // GetPool 获取全局连接池实例 func GetPool() *ConnectionPool { poolOnce.Do(func() { globalPool = &ConnectionPool{ mysqlClients: make(map[uint]*MySQLClient), redisClients: make(map[uint]*RedisClient), mongoClients: make(map[uint]*MongoClient), } }) return globalPool } // GetMySQLClient 获取或创建 MySQL 客户端 func (p *ConnectionPool) GetMySQLClient(conn *models.DbConnection) (*MySQLClient, error) { p.mu.Lock() defer p.mu.Unlock() // 检查是否已存在 if client, ok := p.mysqlClients[conn.ID]; ok { // 测试连接是否有效 if err := client.sqlDB.Ping(); err == nil { return client, nil } // 连接已断开,移除并重新创建 client.Close() delete(p.mysqlClients, conn.ID) } // 解密密码 password, err := crypto.DecryptPassword(conn.Password) if err != nil { return nil, fmt.Errorf("密码解密失败: %v", err) } // 创建新客户端 config := &MySQLConfig{ Host: conn.Host, Port: conn.Port, Username: conn.Username, Password: password, // 如果密码为空,MySQL会尝试无密码连接 Database: conn.Database, } client, err := NewMySQLClient(config) if err != nil { return nil, err } p.mysqlClients[conn.ID] = client return client, nil } // GetRedisClient 获取或创建 Redis 客户端 func (p *ConnectionPool) GetRedisClient(conn *models.DbConnection) (*RedisClient, error) { p.mu.Lock() defer p.mu.Unlock() // 检查是否已存在 if client, ok := p.redisClients[conn.ID]; ok { // 测试连接是否有效 ctx, cancel := context.WithTimeout(context.Background(), common.TimeoutPing) defer cancel() if err := client.client.Ping(ctx).Err(); err == nil { return client, nil } // 连接已断开,移除并重新创建 client.Close() delete(p.redisClients, conn.ID) } // 解密密码 password, err := crypto.DecryptPassword(conn.Password) if err != nil { return nil, fmt.Errorf("密码解密失败: %v", err) } // 解析 Redis DB 编号(从 Database 字段,默认为 0) dbNum := 0 if conn.Database != "" { // 尝试解析 Database 字段为数字 _, err := fmt.Sscanf(conn.Database, "%d", &dbNum) if err != nil { // 如果解析失败,使用默认值 0 dbNum = 0 } // 限制 DB 编号在 0-15 之间 if dbNum < 0 || dbNum > 15 { dbNum = 0 } } // 创建新客户端 config := &RedisConfig{ Host: conn.Host, Port: conn.Port, Password: password, DB: dbNum, } client, err := NewRedisClient(config) if err != nil { return nil, err } p.redisClients[conn.ID] = client return client, nil } // GetMongoClient 获取或创建 MongoDB 客户端 func (p *ConnectionPool) GetMongoClient(conn *models.DbConnection) (*MongoClient, error) { p.mu.Lock() defer p.mu.Unlock() // 检查是否已存在 if client, ok := p.mongoClients[conn.ID]; ok { // 测试连接是否有效 ctx, cancel := context.WithTimeout(context.Background(), common.TimeoutPing) defer cancel() if err := client.client.Ping(ctx, nil); err == nil { return client, nil } // 连接已断开,移除并重新创建 client.Close() delete(p.mongoClients, conn.ID) } // 解密密码 password, err := crypto.DecryptPassword(conn.Password) if err != nil { return nil, fmt.Errorf("密码解密失败: %v", err) } // 解析 Options 获取 MongoDB 连接参数 authSource := "" authMechanism := "" if conn.Options != "" { var opts map[string]interface{} if err := json.Unmarshal([]byte(conn.Options), &opts); err == nil { if as, ok := opts["authSource"].(string); ok && as != "" { authSource = as } if am, ok := opts["authMechanism"].(string); ok && am != "" { authMechanism = am } } } // 创建新客户端 config := &MongoConfig{ Host: conn.Host, Port: conn.Port, Username: conn.Username, Password: password, Database: conn.Database, AuthSource: authSource, AuthMechanism: authMechanism, } client, err := NewMongoClient(config) if err != nil { return nil, err } p.mongoClients[conn.ID] = client return client, nil } // CloseConnection 关闭指定连接 func (p *ConnectionPool) CloseConnection(connID uint, dbType string) { p.mu.Lock() defer p.mu.Unlock() switch dbType { case "mysql": if client, ok := p.mysqlClients[connID]; ok { client.Close() delete(p.mysqlClients, connID) } case "redis": if client, ok := p.redisClients[connID]; ok { client.Close() delete(p.redisClients, connID) } case "mongo": if client, ok := p.mongoClients[connID]; ok { client.Close() delete(p.mongoClients, connID) } } } // CloseAll 关闭所有连接 func (p *ConnectionPool) CloseAll() { p.mu.Lock() defer p.mu.Unlock() for _, client := range p.mysqlClients { client.Close() } for _, client := range p.redisClients { client.Close() } for _, client := range p.mongoClients { client.Close() } p.mysqlClients = make(map[uint]*MySQLClient) p.redisClients = make(map[uint]*RedisClient) p.mongoClients = make(map[uint]*MongoClient) }