新增：连接管理、数据查询等功能

2026-01-22 18:34:59 +08:00
parent 95d3a20292
commit 652f5e5d60
87 changed files with 15082 additions and 162 deletions
--- a/internal/dbclient/pool.go
+++ b/internal/dbclient/pool.go
@@ -0,0 +1,236 @@
+package dbclient
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"sync"
+	"time"
+
+	"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(), 2*time.Second)
+		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(), 2*time.Second)
+		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)
+}