diff --git a/service/singleton/clean_monitor_history_test.go b/service/singleton/clean_monitor_history_test.go new file mode 100644 index 00000000..e01c4610 --- /dev/null +++ b/service/singleton/clean_monitor_history_test.go @@ -0,0 +1,57 @@ +package singleton + +import ( + "path/filepath" + "testing" + + "github.com/stretchr/testify/require" + "gorm.io/gorm" + + "github.com/nezhahq/nezha/model" +) + +func setupCleanMonitorHistoryTestDB(t *testing.T) { + t.Helper() + + previousDB := DB + var err error + DB, err = gorm.Open(openSQLiteDialector(filepath.Join(t.TempDir(), "dashboard.sqlite")), &gorm.Config{}) + require.NoError(t, err) + sqlDB, err := DB.DB() + require.NoError(t, err) + t.Cleanup(func() { + DB = previousDB + if err := sqlDB.Close(); err != nil { + t.Errorf("close transfer cleanup test database: %v", err) + } + }) + + require.NoError(t, DB.AutoMigrate(&model.Server{}, &model.Transfer{}, &model.AlertRule{})) + require.NoError(t, DB.Exec("INSERT INTO servers (id, name, uuid) VALUES (1, 'server', 'clean-monitor-history-test')").Error) +} + +func TestCleanMonitorHistoryWithoutRulesDeletesAllTransfers(t *testing.T) { + setupCleanMonitorHistoryTestDB(t) + require.NoError(t, DB.Create(&model.Transfer{ServerID: 1, In: 1}).Error) + + CleanMonitorHistory() + + var count int64 + require.NoError(t, DB.Model(&model.Transfer{}).Count(&count).Error) + require.Zero(t, count) +} + +func TestCleanMonitorHistoryPreservesTransfersWhenAlertRulesCannotBeLoaded(t *testing.T) { + setupCleanMonitorHistoryTestDB(t) + require.NoError(t, DB.Create(&model.Transfer{ServerID: 1, In: 1}).Error) + require.NoError(t, DB.Exec("INSERT INTO alert_rules (id, name, rules_raw, fail_trigger_tasks_raw, recover_trigger_tasks_raw) VALUES (1, 'broken', '{', '[]', '[]')").Error) + + var alerts []model.AlertRule + require.Error(t, DB.Find(&alerts).Error, "precondition: malformed rules_raw must fail AlertRule.AfterFind") + + CleanMonitorHistory() + + var count int64 + require.NoError(t, DB.Model(&model.Transfer{}).Count(&count).Error) + require.EqualValues(t, 1, count) +} diff --git a/service/singleton/singleton.go b/service/singleton/singleton.go index e2a11da2..bc906afd 100644 --- a/service/singleton/singleton.go +++ b/service/singleton/singleton.go @@ -160,7 +160,10 @@ func CleanMonitorHistory() { specialServerKeep := make(map[uint64]time.Time) var specialServerIDs []uint64 var alerts []model.AlertRule - DB.Find(&alerts) + if err := DB.Find(&alerts).Error; err != nil { + log.Printf("NEZHA>> Failed to load alert rules while cleaning transfer history: %v", err) + return + } for _, alert := range alerts { for _, rule := range alert.Rules { // 是不是流量记录规则 @@ -188,6 +191,14 @@ func CleanMonitorHistory() { for id, couldRemove := range specialServerKeep { DB.Unscoped().Delete(&model.Transfer{}, "server_id = ? AND datetime(`created_at`) < datetime(?)", id, couldRemove) } + if len(specialServerIDs) == 0 { + if allServerKeep.IsZero() { + DB.Unscoped().Session(&gorm.Session{AllowGlobalUpdate: true}).Delete(&model.Transfer{}) + } else { + DB.Unscoped().Delete(&model.Transfer{}, "datetime(`created_at`) < datetime(?)", allServerKeep) + } + return + } if allServerKeep.IsZero() { DB.Unscoped().Delete(&model.Transfer{}, "server_id NOT IN (?)", specialServerIDs) } else {