From dc72e5850ab80b63d2a64cd04fa8c5aacc1c789e Mon Sep 17 00:00:00 2001 From: tbphp Date: Wed, 16 Jul 2025 22:17:13 +0800 Subject: [PATCH] =?UTF-8?q?fix:=20=E4=BF=AE=E5=A4=8D=E6=97=A5=E5=BF=97?= =?UTF-8?q?=E8=A1=A8=E6=95=B0=E6=8D=AE?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/app/app.go | 3 + internal/db/migrations/migration.go | 10 ++ .../db/migrations/v1.0.13_fix_request_logs.go | 130 ++++++++++++++++++ 3 files changed, 143 insertions(+) create mode 100644 internal/db/migrations/migration.go create mode 100644 internal/db/migrations/v1.0.13_fix_request_logs.go diff --git a/internal/app/app.go b/internal/app/app.go index 13d32e1..a0b2704 100644 --- a/internal/app/app.go +++ b/internal/app/app.go @@ -9,6 +9,7 @@ import ( "time" "gpt-load/internal/config" + db "gpt-load/internal/db/migrations" "gpt-load/internal/keypool" "gpt-load/internal/models" "gpt-load/internal/proxy" @@ -88,6 +89,8 @@ func (a *App) Start() error { ); err != nil { return fmt.Errorf("database auto-migration failed: %w", err) } + // 数据修复 + db.MigrateDatabase(a.db) logrus.Info("Database auto-migration completed.") // 初始化系统设置 diff --git a/internal/db/migrations/migration.go b/internal/db/migrations/migration.go new file mode 100644 index 0000000..fac2eb6 --- /dev/null +++ b/internal/db/migrations/migration.go @@ -0,0 +1,10 @@ +package db + +import ( + "gorm.io/gorm" +) + +func MigrateDatabase(db *gorm.DB) error { + // v1.0.13 修复请求日志数据 + return V1_0_13_FixRequestLogs(db) +} diff --git a/internal/db/migrations/v1.0.13_fix_request_logs.go b/internal/db/migrations/v1.0.13_fix_request_logs.go new file mode 100644 index 0000000..a17e15d --- /dev/null +++ b/internal/db/migrations/v1.0.13_fix_request_logs.go @@ -0,0 +1,130 @@ +package db + +import ( + "gpt-load/internal/models" + + "github.com/sirupsen/logrus" + "gorm.io/gorm" +) + +func V1_0_13_FixRequestLogs(db *gorm.DB) error { + return db.Transaction(func(tx *gorm.DB) error { + // 如果有key_id，就执行修复 + if !tx.Migrator().HasColumn(&models.RequestLog{}, "key_id") { + return nil + } + + logrus.Info("Old schema detected. Starting data migration for request_logs...") + + if !tx.Migrator().HasColumn(&models.RequestLog{}, "group_name") { + logrus.Info("Adding 'group_name' column to request_logs table...") + if err := tx.Migrator().AddColumn(&models.RequestLog{}, "group_name"); err != nil { + return err + } + } + if !tx.Migrator().HasColumn(&models.RequestLog{}, "key_value") { + logrus.Info("Adding 'key_value' column to request_logs table...") + if err := tx.Migrator().AddColumn(&models.RequestLog{}, "key_value"); err != nil { + return err + } + } + + type OldRequestLog struct { + ID string + KeyID uint `gorm:"column:key_id"` + GroupID uint + } + + batchSize := 1000 + for i := 0; ; i++ { + logrus.Infof("Processing batch %d...", i+1) + var oldLogs []OldRequestLog + + result := tx.Model(&models.RequestLog{}). + Select("id", "key_id", "group_id"). + Where("key_value IS NULL OR group_name IS NULL"). + Limit(batchSize). + Find(&oldLogs) + + if result.Error != nil { + return result.Error + } + + if len(oldLogs) == 0 { + logrus.Info("All batches processed.") + break + } + + keyIDMap := make(map[uint]bool) + groupIDMap := make(map[uint]bool) + for _, logEntry := range oldLogs { + if logEntry.KeyID > 0 { + keyIDMap[logEntry.KeyID] = true + } + if logEntry.GroupID > 0 { + groupIDMap[logEntry.GroupID] = true + } + } + + var apiKeys []models.APIKey + if len(keyIDMap) > 0 { + var keyIDs []uint + for id := range keyIDMap { + keyIDs = append(keyIDs, id) + } + if err := tx.Model(&models.APIKey{}).Where("id IN ?", keyIDs).Find(&apiKeys).Error; err != nil { + return err + } + } + keyValueMapping := make(map[uint]string) + for _, key := range apiKeys { + keyValueMapping[key.ID] = key.KeyValue + } + + var groups []models.Group + if len(groupIDMap) > 0 { + var groupIDs []uint + for id := range groupIDMap { + groupIDs = append(groupIDs, id) + } + if err := tx.Model(&models.Group{}).Where("id IN ?", groupIDs).Find(&groups).Error; err != nil { + return err + } + } + groupNameMapping := make(map[uint]string) + for _, group := range groups { + groupNameMapping[group.ID] = group.Name + } + + for _, logEntry := range oldLogs { + groupName, gExists := groupNameMapping[logEntry.GroupID] + if !gExists { + logrus.Warnf("Log ID %s: Could not find Group for group_id %d. Setting group_name to empty string.", logEntry.ID, logEntry.GroupID) + } + + keyValue, kExists := keyValueMapping[logEntry.KeyID] + if !kExists { + logrus.Warnf("Log ID %s: Could not find APIKey for key_id %d. Setting key_value to empty string.", logEntry.ID, logEntry.KeyID) + } + + updates := map[string]any{ + "group_name": groupName, + "key_value": keyValue, + } + if err := tx.Model(&models.RequestLog{}).Where("id = ?", logEntry.ID).UpdateColumns(updates).Error; err != nil { + logrus.WithError(err).Errorf("Failed to update log entry with ID: %s", logEntry.ID) + continue + } + } + logrus.Infof("Successfully updated %d log entries in batch %d.", len(oldLogs), i+1) + } + + logrus.Info("Data migration complete. Dropping 'key_id' column from request_logs table...") + if err := tx.Migrator().DropColumn(&models.RequestLog{}, "key_id"); err != nil { + logrus.WithError(err).Warn("Failed to drop 'key_id' column. This can be done manually.") + } + + logrus.Info("Database migration successful!") + return nil + }) +}