diff --git a/internal/db/migrations/v1.0.13_fix_request_logs.go b/internal/db/migrations/v1.0.13_fix_request_logs.go index a17e15d..a948728 100644 --- a/internal/db/migrations/v1.0.13_fix_request_logs.go +++ b/internal/db/migrations/v1.0.13_fix_request_logs.go @@ -7,124 +7,134 @@ import ( "gorm.io/gorm" ) +// TODO: 更新迁移，待多个版本后旧版本都升级差不多之后移除。 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 + // 如果有key_id，就执行修复 + if !db.Migrator().HasColumn(&models.RequestLog{}, "key_id") { + return nil + } + + logrus.Info("Old schema detected. Starting data migration for request_logs...") + + if !db.Migrator().HasColumn(&models.RequestLog{}, "group_name") { + logrus.Info("Adding 'group_name' column to request_logs table...") + if err := db.Migrator().AddColumn(&models.RequestLog{}, "group_name"); err != nil { + return err // Column addition is critical + } + } + if !db.Migrator().HasColumn(&models.RequestLog{}, "key_value") { + logrus.Info("Adding 'key_value' column to request_logs table...") + if err := db.Migrator().AddColumn(&models.RequestLog{}, "key_value"); err != nil { + return err // Column addition is critical + } + } + + 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 := db.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 { + logrus.WithError(result.Error).Error("Failed to fetch batch of logs. Skipping to next batch.") + continue } - logrus.Info("Old schema detected. Starting data migration for request_logs...") + if len(oldLogs) == 0 { + logrus.Info("All batches processed.") + break + } - 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 + keyIDMap := make(map[uint]bool) + groupIDMap := make(map[uint]bool) + for _, logEntry := range oldLogs { + if logEntry.KeyID > 0 { + keyIDMap[logEntry.KeyID] = true } - } - 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 + if logEntry.GroupID > 0 { + groupIDMap[logEntry.GroupID] = true } } - type OldRequestLog struct { - ID string - KeyID uint `gorm:"column:key_id"` - GroupID uint + var apiKeys []models.APIKey + if len(keyIDMap) > 0 { + var keyIDs []uint + for id := range keyIDMap { + keyIDs = append(keyIDs, id) + } + if err := db.Model(&models.APIKey{}).Where("id IN ?", keyIDs).Find(&apiKeys).Error; err != nil { + logrus.WithError(err).Warn("Failed to fetch API keys for the current batch. Some logs may not be updated.") + } + } + keyValueMapping := make(map[uint]string) + for _, key := range apiKeys { + keyValueMapping[key.ID] = key.KeyValue } - batchSize := 1000 - for i := 0; ; i++ { - logrus.Infof("Processing batch %d...", i+1) - var oldLogs []OldRequestLog + var groups []models.Group + if len(groupIDMap) > 0 { + var groupIDs []uint + for id := range groupIDMap { + groupIDs = append(groupIDs, id) + } + if err := db.Model(&models.Group{}).Where("id IN ?", groupIDs).Find(&groups).Error; err != nil { + logrus.WithError(err).Warn("Failed to fetch groups for the current batch. Some logs may not be updated.") + } + } + groupNameMapping := make(map[uint]string) + for _, group := range groups { + groupNameMapping[group.ID] = group.Name + } - 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) + updateGroups := make(map[string]map[string][]string) - if result.Error != nil { - return result.Error + 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) } - if len(oldLogs) == 0 { - logrus.Info("All batches processed.") - break + 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) } - 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 - } + if _, ok := updateGroups[groupName]; !ok { + updateGroups[groupName] = make(map[string][]string) } + updateGroups[groupName][keyValue] = append(updateGroups[groupName][keyValue], logEntry.ID) + } - 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) - } - + for groupName, keyMap := range updateGroups { + for keyValue, ids := range keyMap { 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 + if err := db.Model(&models.RequestLog{}).Where("id IN ?", ids).UpdateColumns(updates).Error; err != nil { + logrus.WithError(err).Errorf("Failed to update a batch of log entries. Skipping this batch.") } } - logrus.Infof("Successfully updated %d log entries in batch %d.", len(oldLogs), i+1) } + logrus.Infof("Finished processing batch %d. Updated %d log entries.", i+1, len(oldLogs)) + } - 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("Data migration complete. Dropping 'key_id' column from request_logs table...") + if err := db.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 - }) + logrus.Info("Database migration finished!") + return nil }