fork-it/gpt-load
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: 批量分配处理数据

Browse Source
This commit is contained in:
tbphp
2025-07-08 10:00:14 +08:00
parent 22c9b16f5c
commit 356b620ffb
3 changed files with 106 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
internal/handler/handler.go
View File

@@ -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,

27
internal/handler/key_handler.go
View File

@@ -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 {
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
}

90
internal/services/key_service.go
View File

@@ -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"`
@@ -36,13 +42,15 @@ type RestoreKeysResult struct {
type KeyService struct {
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,
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,11 +105,17 @@ 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 {
// 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
totalInGroup := int64(len(existingKeys) + len(newKeysToCreate))
@@ -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)
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)
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
}