u-desk/internal/dbclient/pool.go

package dbclient

import (
	"context"
	"encoding/json"
	"fmt"
	"sync"

	"u-desk/internal/common"
	"u-desk/internal/crypto"
	"u-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)
}