From 356b620ffb35b6cfc74d8caba3c51cccde60c3b9 Mon Sep 17 00:00:00 2001 From: tbphp Date: Tue, 8 Jul 2025 10:00:14 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E6=89=B9=E9=87=8F=E5=88=86=E9=85=8D?= =?UTF-8?q?=E5=A4=84=E7=90=86=E6=95=B0=E6=8D=AE?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/handler/handler.go | 4 -- internal/handler/key_handler.go | 29 +++++---- internal/services/key_service.go | 108 +++++++++++++++++++++++++------ 3 files changed, 106 insertions(+), 35 deletions(-) diff --git a/internal/handler/handler.go b/internal/handler/handler.go index 24e37ad..cdfbb29 100644 --- a/internal/handler/handler.go +++ b/internal/handler/handler.go @@ -6,7 +6,6 @@ import ( "time" "gpt-load/internal/config" - "gpt-load/internal/keypool" "gpt-load/internal/models" "gpt-load/internal/services" "gpt-load/internal/types" @@ -21,7 +20,6 @@ type Server struct { DB *gorm.DB config types.ConfigManager SettingsManager *config.SystemSettingsManager - KeyValidator *keypool.KeyValidator KeyManualValidationService *services.KeyManualValidationService TaskService *services.TaskService KeyService *services.KeyService @@ -34,7 +32,6 @@ type NewServerParams struct { DB *gorm.DB Config types.ConfigManager SettingsManager *config.SystemSettingsManager - KeyValidator *keypool.KeyValidator KeyManualValidationService *services.KeyManualValidationService TaskService *services.TaskService KeyService *services.KeyService @@ -47,7 +44,6 @@ func NewServer(params NewServerParams) *Server { DB: params.DB, config: params.Config, SettingsManager: params.SettingsManager, - KeyValidator: params.KeyValidator, KeyManualValidationService: params.KeyManualValidationService, TaskService: params.TaskService, KeyService: params.KeyService, diff --git a/internal/handler/key_handler.go b/internal/handler/key_handler.go index 0344829..c9500db 100644 --- a/internal/handler/key_handler.go +++ b/internal/handler/key_handler.go @@ -84,7 +84,9 @@ func (s *Server) AddMultipleKeys(c *gin.Context) { result, err := s.KeyService.AddMultipleKeys(req.GroupID, req.KeysText) if err != nil { - if err.Error() == "no valid keys found in the input text" { + if strings.Contains(err.Error(), "batch size exceeds the limit") { + response.Error(c, app_errors.NewAPIError(app_errors.ErrValidation, err.Error())) + } else if err.Error() == "no valid keys found in the input text" { response.Error(c, app_errors.NewAPIError(app_errors.ErrValidation, err.Error())) } else { response.Error(c, app_errors.ParseDBError(err)) @@ -146,7 +148,9 @@ func (s *Server) DeleteMultipleKeys(c *gin.Context) { result, err := s.KeyService.DeleteMultipleKeys(req.GroupID, req.KeysText) if err != nil { - if err.Error() == "no valid keys found in the input text" { + if strings.Contains(err.Error(), "batch size exceeds the limit") { + response.Error(c, app_errors.NewAPIError(app_errors.ErrValidation, err.Error())) + } else if err.Error() == "no valid keys found in the input text" { response.Error(c, app_errors.NewAPIError(app_errors.ErrValidation, err.Error())) } else { response.Error(c, app_errors.ParseDBError(err)) @@ -176,7 +180,9 @@ func (s *Server) RestoreMultipleKeys(c *gin.Context) { result, err := s.KeyService.RestoreMultipleKeys(req.GroupID, req.KeysText) if err != nil { - if err.Error() == "no valid keys found in the input text" { + if strings.Contains(err.Error(), "batch size exceeds the limit") { + response.Error(c, app_errors.NewAPIError(app_errors.ErrValidation, err.Error())) + } else if err.Error() == "no valid keys found in the input text" { response.Error(c, app_errors.NewAPIError(app_errors.ErrValidation, err.Error())) } else { response.Error(c, app_errors.ParseDBError(err)) @@ -205,16 +211,15 @@ func (s *Server) TestMultipleKeys(c *gin.Context) { return } - // Re-use the parsing logic from the key service - keysToTest := s.KeyService.ParseKeysFromText(req.KeysText) - if len(keysToTest) == 0 { - response.Error(c, app_errors.NewAPIError(app_errors.ErrValidation, "no valid keys found in the input text")) - return - } - - results, err := s.KeyValidator.TestMultipleKeys(c.Request.Context(), group, keysToTest) + results, err := s.KeyService.TestMultipleKeys(c.Request.Context(), group, req.KeysText) if err != nil { - response.Error(c, app_errors.ParseDBError(err)) + if strings.Contains(err.Error(), "batch size exceeds the limit") { + response.Error(c, app_errors.NewAPIError(app_errors.ErrValidation, err.Error())) + } else if err.Error() == "no valid keys found in the input text" { + response.Error(c, app_errors.NewAPIError(app_errors.ErrValidation, err.Error())) + } else { + response.Error(c, app_errors.ParseDBError(err)) + } return } diff --git a/internal/services/key_service.go b/internal/services/key_service.go index d0ee85a..e145730 100644 --- a/internal/services/key_service.go +++ b/internal/services/key_service.go @@ -1,6 +1,7 @@ package services import ( + "context" "encoding/json" "fmt" "gpt-load/internal/keypool" @@ -11,6 +12,11 @@ import ( "gorm.io/gorm" ) +const ( + maxRequestKeys = 5000 + chunkSize = 1000 +) + // AddKeysResult holds the result of adding multiple keys. type AddKeysResult struct { AddedCount int `json:"added_count"` @@ -34,15 +40,17 @@ type RestoreKeysResult struct { // KeyService provides services related to API keys. type KeyService struct { - DB *gorm.DB - KeyProvider *keypool.KeyProvider + DB *gorm.DB + KeyProvider *keypool.KeyProvider + KeyValidator *keypool.KeyValidator } // NewKeyService creates a new KeyService. -func NewKeyService(db *gorm.DB, keyProvider *keypool.KeyProvider) *KeyService { +func NewKeyService(db *gorm.DB, keyProvider *keypool.KeyProvider, keyValidator *keypool.KeyValidator) *KeyService { return &KeyService{ - DB: db, - KeyProvider: keyProvider, + DB: db, + KeyProvider: keyProvider, + KeyValidator: keyValidator, } } @@ -50,6 +58,9 @@ func NewKeyService(db *gorm.DB, keyProvider *keypool.KeyProvider) *KeyService { func (s *KeyService) AddMultipleKeys(groupID uint, keysText string) (*AddKeysResult, error) { // 1. Parse keys from the text block keys := s.ParseKeysFromText(keysText) + if len(keys) > maxRequestKeys { + return nil, fmt.Errorf("batch size exceeds the limit of %d keys, got %d", maxRequestKeys, len(keys)) + } if len(keys) == 0 { return nil, fmt.Errorf("no valid keys found in the input text") } @@ -94,10 +105,16 @@ func (s *KeyService) AddMultipleKeys(groupID uint, keysText string) (*AddKeysRes }, nil } - // 4. Use KeyProvider to add keys - err := s.KeyProvider.AddKeys(groupID, newKeysToCreate) - if err != nil { - return nil, err + // 4. Use KeyProvider to add keys in chunks + for i := 0; i < len(newKeysToCreate); i += chunkSize { + end := i + chunkSize + if end > len(newKeysToCreate) { + end = len(newKeysToCreate) + } + chunk := newKeysToCreate[i:end] + if err := s.KeyProvider.AddKeys(groupID, chunk); err != nil { + return nil, err + } } // 5. Calculate new total count @@ -164,16 +181,28 @@ func (s *KeyService) isValidKeyFormat(key string) bool { // RestoreMultipleKeys handles the business logic of restoring keys from a text block. func (s *KeyService) RestoreMultipleKeys(groupID uint, keysText string) (*RestoreKeysResult, error) { keysToRestore := s.ParseKeysFromText(keysText) + if len(keysToRestore) > maxRequestKeys { + return nil, fmt.Errorf("batch size exceeds the limit of %d keys, got %d", maxRequestKeys, len(keysToRestore)) + } if len(keysToRestore) == 0 { return nil, fmt.Errorf("no valid keys found in the input text") } - restoredCount, err := s.KeyProvider.RestoreMultipleKeys(groupID, keysToRestore) - if err != nil { - return nil, err + var totalRestoredCount int64 + for i := 0; i < len(keysToRestore); i += chunkSize { + end := i + chunkSize + if end > len(keysToRestore) { + end = len(keysToRestore) + } + chunk := keysToRestore[i:end] + restoredCount, err := s.KeyProvider.RestoreMultipleKeys(groupID, chunk) + if err != nil { + return nil, err + } + totalRestoredCount += restoredCount } - ignoredCount := len(keysToRestore) - int(restoredCount) + ignoredCount := len(keysToRestore) - int(totalRestoredCount) var totalInGroup int64 if err := s.DB.Model(&models.APIKey{}).Where("group_id = ?", groupID).Count(&totalInGroup).Error; err != nil { @@ -181,7 +210,7 @@ func (s *KeyService) RestoreMultipleKeys(groupID uint, keysText string) (*Restor } return &RestoreKeysResult{ - RestoredCount: int(restoredCount), + RestoredCount: int(totalRestoredCount), IgnoredCount: ignoredCount, TotalInGroup: totalInGroup, }, nil @@ -200,16 +229,28 @@ func (s *KeyService) ClearAllInvalidKeys(groupID uint) (int64, error) { // DeleteMultipleKeys handles the business logic of deleting keys from a text block. func (s *KeyService) DeleteMultipleKeys(groupID uint, keysText string) (*DeleteKeysResult, error) { keysToDelete := s.ParseKeysFromText(keysText) + if len(keysToDelete) > maxRequestKeys { + return nil, fmt.Errorf("batch size exceeds the limit of %d keys, got %d", maxRequestKeys, len(keysToDelete)) + } if len(keysToDelete) == 0 { return nil, fmt.Errorf("no valid keys found in the input text") } - deletedCount, err := s.KeyProvider.RemoveKeys(groupID, keysToDelete) - if err != nil { - return nil, err + var totalDeletedCount int64 + for i := 0; i < len(keysToDelete); i += chunkSize { + end := i + chunkSize + if end > len(keysToDelete) { + end = len(keysToDelete) + } + chunk := keysToDelete[i:end] + deletedCount, err := s.KeyProvider.RemoveKeys(groupID, chunk) + if err != nil { + return nil, err + } + totalDeletedCount += deletedCount } - ignoredCount := len(keysToDelete) - int(deletedCount) + ignoredCount := len(keysToDelete) - int(totalDeletedCount) var totalInGroup int64 if err := s.DB.Model(&models.APIKey{}).Where("group_id = ?", groupID).Count(&totalInGroup).Error; err != nil { @@ -217,7 +258,7 @@ func (s *KeyService) DeleteMultipleKeys(groupID uint, keysText string) (*DeleteK } return &DeleteKeysResult{ - DeletedCount: int(deletedCount), + DeletedCount: int(totalDeletedCount), IgnoredCount: ignoredCount, TotalInGroup: totalInGroup, }, nil @@ -237,3 +278,32 @@ func (s *KeyService) ListKeysInGroupQuery(groupID uint, statusFilter string, sea return query } + +// TestMultipleKeys handles a one-off validation test for multiple keys. +func (s *KeyService) TestMultipleKeys(ctx context.Context, group *models.Group, keysText string) ([]keypool.KeyTestResult, error) { + keysToTest := s.ParseKeysFromText(keysText) + if len(keysToTest) > maxRequestKeys { + return nil, fmt.Errorf("batch size exceeds the limit of %d keys, got %d", maxRequestKeys, len(keysToTest)) + } + if len(keysToTest) == 0 { + return nil, fmt.Errorf("no valid keys found in the input text") + } + + var allResults []keypool.KeyTestResult + for i := 0; i < len(keysToTest); i += chunkSize { + end := i + chunkSize + if end > len(keysToTest) { + end = len(keysToTest) + } + chunk := keysToTest[i:end] + results, err := s.KeyValidator.TestMultipleKeys(ctx, group, chunk) + if err != nil { + // If one chunk fails, we might want to stop or collect partial results. + // For now, let's stop and return the error. + return nil, err + } + allResults = append(allResults, results...) + } + + return allResults, nil +}