speed-test/go-speed-test/internal/database/database.go

package database

import (
	"fmt"
	"time"

	"clash-speed-test/internal/config"
	applogger "clash-speed-test/internal/logger"

	"gorm.io/driver/mysql"
	"gorm.io/gorm"
	gormlogger "gorm.io/gorm/logger"
)

var DB *gorm.DB

type Node struct {
	ID          uint      `gorm:"primarykey" json:"id"`
	Name        string    `gorm:"size:255;not null" json:"name"`
	Type        string    `gorm:"size:50;not null" json:"type"`
	Server      string    `gorm:"size:255;not null" json:"server"`
	Port        int       `gorm:"not null" json:"port"`
	Password    string    `gorm:"size:255" json:"password"`
	Method      string    `gorm:"size:50" json:"method"`
	UUID        string    `gorm:"size:255" json:"uuid"`
	AlterID     int       `json:"alter_id"`
	Network     string    `gorm:"size:50" json:"network"`
	TLS         bool      `json:"tls"`
	SNI         string    `gorm:"size:255" json:"sni"`
	WSPath      string    `gorm:"size:255" json:"ws_path"`
	WSHeaders   string    `gorm:"type:text" json:"ws_headers"`
	Username    string    `gorm:"size:255" json:"username"`
	Protocol    string    `gorm:"size:50" json:"protocol"`
	Obfs        string    `gorm:"size:50" json:"obfs"`
	Group       string    `gorm:"size:255" json:"group"`
	IsActive    bool      `gorm:"default:true" json:"is_active"`
	Status      string    `gorm:"size:20;default:'offline'" json:"status"`
	LastTestTime *time.Time `json:"last_test_time"`
	LastTestResult *bool   `json:"last_test_result"`
	AverageLatency *int    `json:"average_latency"`
	AverageSpeed   *float64 `json:"average_speed"`
	FailureCount   int     `gorm:"default:0" json:"failure_count"`
	CreatedAt   time.Time `json:"created_at"`
	UpdatedAt   time.Time `json:"updated_at"`
}

type TestResult struct {
	ID           uint      `gorm:"primarykey" json:"id"`
	NodeID       uint      `gorm:"not null" json:"node_id"`
	Node         Node      `gorm:"foreignKey:NodeID" json:"node"`
	TestTime     time.Time `gorm:"not null" json:"test_time"`
	IsSuccess    bool      `gorm:"not null" json:"is_success"`
	Latency      *int      `json:"latency"`
	DownloadSpeed *float64 `json:"download_speed"`
	UploadSpeed   *float64 `json:"upload_speed"`
	PacketLoss    *float64 `json:"packet_loss"`
	TestURL       string   `gorm:"size:500" json:"test_url"`
	ErrorMessage  string   `gorm:"type:text" json:"error_message"`
	TestDuration  *int     `json:"test_duration"`
	IPAddress     string   `gorm:"size:100" json:"ip_address"`
	Location      string   `gorm:"size:255" json:"location"`
	CreatedAt     time.Time `json:"created_at"`
	UpdatedAt     time.Time `json:"updated_at"`
}

type Notification struct {
	ID        uint      `gorm:"primarykey" json:"id"`
	Type      string    `gorm:"size:50;not null" json:"type"`
	Title     string    `gorm:"size:255;not null" json:"title"`
	Message   string    `gorm:"type:text;not null" json:"message"`
	IsRead    bool      `gorm:"default:false" json:"is_read"`
	CreatedAt time.Time `json:"created_at"`
	UpdatedAt time.Time `json:"updated_at"`
}

func Init(cfg config.DatabaseConfig) (*gorm.DB, error) {
	dsn := fmt.Sprintf("%s:%s@tcp(%s:%d)/%s?charset=utf8mb4&parseTime=True&loc=Local",
		cfg.Username,
		cfg.Password,
		cfg.Host,
		cfg.Port,
		cfg.Database,
	)

	// 配置GORM日志
	gormLogger := gormlogger.New(
		&gormlogger.Writer{},
		gormlogger.Config{
			SlowThreshold:             time.Second,
			LogLevel:                  gormlogger.Warn,
			IgnoreRecordNotFoundError: true,
			Colorful:                  false,
		},
	)

	db, err := gorm.Open(mysql.Open(dsn), &gorm.Config{
		Logger: gormLogger,
	})
	if err != nil {
		return nil, fmt.Errorf("连接数据库失败: %w", err)
	}

	// 自动迁移表结构
	if err := db.AutoMigrate(&Node{}, &TestResult{}, &Notification{}); err != nil {
		return nil, fmt.Errorf("数据库迁移失败: %w", err)
	}

	DB = db
	applogger.Info("数据库连接成功", map[string]interface{}{
		"host": cfg.Host,
		"port": cfg.Port,
		"database": cfg.Database,
	})

	return db, nil
}

// 获取所有活跃节点
func GetActiveNodes() ([]Node, error) {
	var nodes []Node
	err := DB.Where("is_active = ?", true).Find(&nodes).Error
	return nodes, err
}

// 保存测试结果
func SaveTestResult(result *TestResult) error {
	return DB.Create(result).Error
}

// 获取最近的测试结果
func GetRecentTestResults(limit int) ([]TestResult, error) {
	var results []TestResult
	err := DB.Preload("Node").Order("test_time DESC").Limit(limit).Find(&results).Error
	return results, err
}

// 获取节点的测试历史
func GetNodeTestHistory(nodeID uint, limit int) ([]TestResult, error) {
	var results []TestResult
	err := DB.Where("node_id = ?", nodeID).Order("test_time DESC").Limit(limit).Find(&results).Error
	return results, err
}

// 更新节点信息
func UpdateNode(nodeID uint, updateData map[string]interface{}) error {
	return DB.Model(&Node{}).Where("id = ?", nodeID).Updates(updateData).Error
}