u-desk/internal/service/connection_service.go

package service

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"u-desk/internal/crypto"
	"u-desk/internal/dbclient"
	"u-desk/internal/storage/models"
	"u-desk/internal/storage/repository"

	"gorm.io/gorm"
)

// ConnectionService 连接管理服务
type ConnectionService struct {
	repo repository.ConnectionRepository
}

// NewConnectionService 创建连接服务
func NewConnectionService() (*ConnectionService, error) {
	repo, err := repository.NewConnectionRepository()
	if err != nil {
		return nil, fmt.Errorf("创建连接仓库失败: %v", err)
	}
	return &ConnectionService{repo: repo}, nil
}

// SaveConnection 保存连接配置
func (s *ConnectionService) SaveConnection(conn *models.DbConnection) error {
	// 验证
	if conn.Name == "" {
		return fmt.Errorf("连接名称不能为空")
	}
	if conn.Type == "" {
		return fmt.Errorf("数据库类型不能为空")
	}
	if conn.Host == "" {
		return fmt.Errorf("主机地址不能为空")
	}

	// 检查名称是否重复
	existing, err := s.repo.FindByName(conn.Name, conn.ID)
	if err != nil {
		return fmt.Errorf("检查连接名称失败: %v", err)
	}
	if existing != nil {
		return fmt.Errorf("连接名称已存在")
	}

	// 处理密码
	if conn.ID > 0 {
		if conn.Password == "" {
			// 更新模式：保留原密码
			conn.Password, err = s.getPassword(conn.ID)
			if err != nil {
				return err
			}
		} else {
			// 加密新密码
			conn.Password, err = crypto.EncryptPassword(conn.Password)
			if err != nil {
				return fmt.Errorf("密码加密失败: %v", err)
			}
		}
	} else {
		// 新增模式：加密密码
		conn.Password, err = crypto.EncryptPassword(conn.Password)
		if err != nil {
			return fmt.Errorf("密码加密失败: %v", err)
		}
	}

	return s.repo.Save(conn)
}

// getPassword 获取原始密码
func (s *ConnectionService) getPassword(id uint) (string, error) {
	existing, err := s.repo.FindByID(id)
	if err != nil {
		return "", fmt.Errorf("获取原连接配置失败: %v", err)
	}
	return existing.Password, nil
}

// ListConnections 获取连接列表
func (s *ConnectionService) ListConnections() ([]models.DbConnection, error) {
	return s.repo.FindAll()
}

// GetConnection 获取连接详情
func (s *ConnectionService) GetConnection(id uint) (*models.DbConnection, error) {
	return s.repo.FindByID(id)
}

// DeleteConnection 删除连接配置（含关联数据和连接池清理）
func (s *ConnectionService) DeleteConnection(id uint) error {
	conn, err := s.repo.FindByID(id)
	if err != nil {
		if errors.Is(err, gorm.ErrRecordNotFound) {
			return nil // 连接不存在视为成功
		}
		return fmt.Errorf("获取连接配置失败: %v", err)
	}

	// 关闭连接池中的连接
	dbclient.GetPool().CloseConnection(id, conn.Type)

	// 删除连接记录
	return s.repo.Delete(id)
}

// TestConnection 测试连接（通过已保存的连接ID）
func (s *ConnectionService) TestConnection(id uint) error {
	conn, err := s.repo.FindByID(id)
	if err != nil {
		return fmt.Errorf("获取连接配置失败: %v", err)
	}

	// 解密密码用于测试
	password, err := crypto.DecryptPassword(conn.Password)
	if err != nil {
		return fmt.Errorf("密码解密失败: %v", err)
	}

	// 根据类型测试连接
	switch conn.Type {
	case "mysql":
		return dbclient.TestMySQLConnection(conn.Host, conn.Port, conn.Username, password, conn.Database)
	case "redis":
		return dbclient.TestRedisConnection(conn.Host, conn.Port, password)
	case "mongo":
		// 解析 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
				}
			}
		}
		return dbclient.TestMongoConnectionWithOptions(conn.Host, conn.Port, conn.Username, password, conn.Database, authSource, authMechanism)
	default:
		return fmt.Errorf("不支持的数据库类型: %s", conn.Type)
	}
}

// TestConnectionWithParams 测试连接（直接传入参数，不保存数据）
func (s *ConnectionService) TestConnectionWithParams(connType, host string, port int, username, password, database, options string, existingId uint) error {
	// 验证必填项
	if connType == "" {
		return fmt.Errorf("数据库类型不能为空")
	}
	if host == "" {
		return fmt.Errorf("主机地址不能为空")
	}

	// 如果是编辑模式且密码为空，尝试获取已保存的密码
	actualPassword := password
	if existingId > 0 && password == "" {
		conn, err := s.repo.FindByID(existingId)
		if err != nil {
			return fmt.Errorf("获取原连接配置失败: %v", err)
		}
		// 解密原密码
		actualPassword, err = crypto.DecryptPassword(conn.Password)
		if err != nil {
			return fmt.Errorf("密码解密失败: %v", err)
		}
	}

	// 根据类型测试连接
	switch connType {
	case "mysql":
		return dbclient.TestMySQLConnection(host, port, username, actualPassword, database)
	case "redis":
		return dbclient.TestRedisConnection(host, port, actualPassword)
	case "mongo":
		// 解析 Options 获取 MongoDB 连接参数
		authSource := ""
		authMechanism := ""
		if options != "" {
			var opts map[string]interface{}
			if err := json.Unmarshal([]byte(options), &opts); err == nil {
				if as, ok := opts["authSource"].(string); ok && as != "" {
					authSource = as
				}
				if am, ok := opts["authMechanism"].(string); ok && am != "" {
					authMechanism = am
				}
			}
		}
		return dbclient.TestMongoConnectionWithOptions(host, port, username, actualPassword, database, authSource, authMechanism)
	default:
		return fmt.Errorf("不支持的数据库类型: %s", connType)
	}
}

// LoadAllDatabases 加载全部数据库列表
func (s *ConnectionService) LoadAllDatabases(dbType, host string, port int, username, password, database, options string, existingId uint) ([]string, error) {
	// 如果是编辑模式且密码为空，尝试获取已保存的密码
	actualPassword := password
	if existingId > 0 && password == "" {
		conn, err := s.repo.FindByID(existingId)
		if err != nil {
			return nil, fmt.Errorf("获取原连接配置失败: %v", err)
		}
		actualPassword, err = crypto.DecryptPassword(conn.Password)
		if err != nil {
			return nil, fmt.Errorf("密码解密失败: %v", err)
		}
	}

	// 解析 MongoDB 选项
	authSource := ""
	authMechanism := ""
	if options != "" {
		var opts map[string]interface{}
		if err := json.Unmarshal([]byte(options), &opts); err == nil {
			authSource, _ = opts["authSource"].(string)
			authMechanism, _ = opts["authMechanism"].(string)
		}
	}

	switch dbType {
	case "mysql":
		return loadDatabasesForMySQL(host, port, username, actualPassword, database)
	case "mongo":
		return loadDatabasesForMongo(host, port, username, actualPassword, database, authSource, authMechanism)
	case "redis":
		return []string{}, nil
	default:
		return nil, fmt.Errorf("不支持的数据库类型: %s", dbType)
	}
}

func loadDatabasesForMySQL(host string, port int, username, password, defaultDatabase string) ([]string, error) {
	config := &dbclient.MySQLConfig{
		Host: host, Port: port, Username: username,
		Password: password, Database: defaultDatabase,
	}
	client, err := dbclient.NewMySQLClient(config)
	if err != nil {
		return nil, err
	}
	defer client.Close()
	return client.ListDatabases(context.Background())
}

func loadDatabasesForMongo(host string, port int, username, password, defaultDatabase, authSource, authMechanism string) ([]string, error) {
	config := &dbclient.MongoConfig{
		Host: host, Port: port, Username: username,
		Password: password, Database: defaultDatabase,
		AuthSource: authSource, AuthMechanism: authMechanism,
	}
	client, err := dbclient.NewMongoClient(config)
	if err != nil {
		return nil, err
	}
	defer client.Close()
	return client.ListDatabases(context.Background())
}