From ea54897e311f72b2e534c80c4dff33556c42e7bf Mon Sep 17 00:00:00 2001 From: tbphp Date: Mon, 7 Jul 2025 20:19:35 +0800 Subject: [PATCH] =?UTF-8?q?fix:=20=E4=BC=98=E5=8C=96=E9=80=BB=E8=BE=91?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/keypool/provider.go | 26 +++++++++++--------------- internal/services/key_service.go | 24 +++++++----------------- internal/store/memory.go | 7 ++++--- internal/store/redis.go | 4 ++-- internal/store/store.go | 2 +- 5 files changed, 25 insertions(+), 38 deletions(-) diff --git a/internal/keypool/provider.go b/internal/keypool/provider.go index dbdcdbe..c5bd7bf 100644 --- a/internal/keypool/provider.go +++ b/internal/keypool/provider.go @@ -103,7 +103,7 @@ func (p *KeyProvider) handleSuccess(keyID uint, keyHashKey, activeKeysListKey st return } - if err := p.store.HSet(keyHashKey, "failure_count", 0); err != nil { + if err := p.store.HSet(keyHashKey, map[string]any{"failure_count": 0}); err != nil { logrus.WithFields(logrus.Fields{"keyID": keyID, "error": err}).Error("Failed to reset failure count in store, aborting DB update.") return } @@ -112,7 +112,7 @@ func (p *KeyProvider) handleSuccess(keyID uint, keyHashKey, activeKeysListKey st if isInvalid { logrus.WithField("keyID", keyID).Info("Key has recovered and is being restored to active pool.") - if err := p.store.HSet(keyHashKey, "status", models.KeyStatusActive); err != nil { + if err := p.store.HSet(keyHashKey, map[string]any{"status": models.KeyStatusActive}); err != nil { logrus.WithFields(logrus.Fields{"keyID": keyID, "error": err}).Error("Failed to update key status to active in store, aborting DB update.") return } @@ -160,7 +160,7 @@ func (p *KeyProvider) handleFailure(keyID uint, keyHashKey, activeKeysListKey st logrus.WithFields(logrus.Fields{"keyID": keyID, "error": err}).Error("Failed to LRem key from active list, aborting DB update.") return } - if err := p.store.HSet(keyHashKey, "status", models.KeyStatusInvalid); err != nil { + if err := p.store.HSet(keyHashKey, map[string]any{"status": models.KeyStatusInvalid}); err != nil { logrus.WithFields(logrus.Fields{"keyID": keyID, "error": err}).Error("Failed to update key status to invalid in store, aborting DB update.") return } @@ -222,10 +222,8 @@ func (p *KeyProvider) LoadKeysFromDB() error { if pipeline != nil { pipeline.HSet(keyHashKey, keyDetails) } else { - for field, value := range keyDetails { - if err := p.store.HSet(keyHashKey, field, value); err != nil { - logrus.WithFields(logrus.Fields{"keyID": key.ID, "error": err}).Error("Failed to HSet key details") - } + if err := p.store.HSet(keyHashKey, keyDetails); err != nil { + logrus.WithFields(logrus.Fields{"keyID": key.ID, "error": err}).Error("Failed to HSet key details") } } @@ -307,7 +305,9 @@ func (p *KeyProvider) RemoveKeys(groupID uint, keyValues []string) (int64, error return nil } - result := tx.Where("group_id = ? AND key_value IN ?", groupID, keyValues).Delete(&models.APIKey{}) + keyIDsToDelete := pluckIDs(keysToDelete) + + result := tx.Where("id IN ?", keyIDsToDelete).Delete(&models.APIKey{}) if result.Error != nil { return result.Error } @@ -394,17 +394,13 @@ func (p *KeyProvider) addKeyToStore(key *models.APIKey) error { // 1. Store key details in HASH keyHashKey := fmt.Sprintf("key:%d", key.ID) keyDetails := p.apiKeyToMap(key) - for field, value := range keyDetails { - if err := p.store.HSet(keyHashKey, field, value); err != nil { - return fmt.Errorf("failed to HSet key details for key %d: %w", key.ID, err) - } + if err := p.store.HSet(keyHashKey, keyDetails); err != nil { + return fmt.Errorf("failed to HSet key details for key %d: %w", key.ID, err) } // 2. If active, add to the active LIST if key.Status == models.KeyStatusActive { activeKeysListKey := fmt.Sprintf("group:%d:active_keys", key.GroupID) - // To prevent duplicates, first remove any existing instance of the key from the list. - // This makes the add operation idempotent regarding the list. if err := p.store.LRem(activeKeysListKey, 0, key.ID); err != nil { return fmt.Errorf("failed to LRem key %d before LPush for group %d: %w", key.ID, key.GroupID, err) } @@ -432,7 +428,7 @@ func (p *KeyProvider) removeKeyFromStore(keyID, groupID uint) error { // apiKeyToMap converts an APIKey model to a map for HSET. func (p *KeyProvider) apiKeyToMap(key *models.APIKey) map[string]any { return map[string]any{ - "id": fmt.Sprint(key.ID), // Use fmt.Sprint for consistency in pipeline + "id": fmt.Sprint(key.ID), "key_string": key.KeyValue, "status": key.Status, "failure_count": key.FailureCount, diff --git a/internal/services/key_service.go b/internal/services/key_service.go index 5613e38..f737fea 100644 --- a/internal/services/key_service.go +++ b/internal/services/key_service.go @@ -20,9 +20,9 @@ type AddKeysResult struct { // DeleteKeysResult holds the result of deleting multiple keys. type DeleteKeysResult struct { - DeletedCount int `json:"deleted_count"` - IgnoredCount int `json:"ignored_count"` - TotalInGroup int64 `json:"total_in_group"` + DeletedCount int `json:"deleted_count"` + IgnoredCount int `json:"ignored_count"` + TotalInGroup int64 `json:"total_in_group"` } // KeyService provides services related to API keys. @@ -80,26 +80,21 @@ func (s *KeyService) AddMultipleKeys(groupID uint, keysText string) (*AddKeysRes } if len(newKeysToCreate) == 0 { - var totalInGroup int64 - s.DB.Model(&models.APIKey{}).Where("group_id = ?", groupID).Count(&totalInGroup) return &AddKeysResult{ AddedCount: 0, IgnoredCount: len(keys), - TotalInGroup: totalInGroup, + TotalInGroup: int64(len(existingKeys)), }, nil } - // 4. Use KeyProvider to add keys, which handles DB and cache + // 4. Use KeyProvider to add keys err := s.KeyProvider.AddKeys(groupID, newKeysToCreate) if err != nil { return nil, err } - // 5. Get the new total count - var totalInGroup int64 - if err := s.DB.Model(&models.APIKey{}).Where("group_id = ?", groupID).Count(&totalInGroup).Error; err != nil { - return nil, err - } + // 5. Calculate new total count + totalInGroup := int64(len(existingKeys) + len(newKeysToCreate)) return &AddKeysResult{ AddedCount: len(newKeysToCreate), @@ -171,13 +166,11 @@ 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) { - // 1. Parse keys from the text block keysToDelete := s.ParseKeysFromText(keysText) if len(keysToDelete) == 0 { return nil, fmt.Errorf("no valid keys found in the input text") } - // 2. Use KeyProvider to delete keys, which handles DB and cache deletedCount, err := s.KeyProvider.RemoveKeys(groupID, keysToDelete) if err != nil { return nil, err @@ -185,7 +178,6 @@ func (s *KeyService) DeleteMultipleKeys(groupID uint, keysText string) (*DeleteK ignoredCount := len(keysToDelete) - int(deletedCount) - // 3. Get the new total count var totalInGroup int64 if err := s.DB.Model(&models.APIKey{}).Where("group_id = ?", groupID).Count(&totalInGroup).Error; err != nil { return nil, err @@ -199,7 +191,6 @@ func (s *KeyService) DeleteMultipleKeys(groupID uint, keysText string) (*DeleteK } // ListKeysInGroupQuery builds a query to list all keys within a specific group, filtered by status. -// It returns a GORM query builder, allowing the handler to apply pagination. func (s *KeyService) ListKeysInGroupQuery(groupID uint, statusFilter string, searchKeyword string) *gorm.DB { query := s.DB.Model(&models.APIKey{}).Where("group_id = ?", groupID) @@ -208,7 +199,6 @@ func (s *KeyService) ListKeysInGroupQuery(groupID uint, statusFilter string, sea } if searchKeyword != "" { - // Use LIKE for fuzzy search on the key_value query = query.Where("key_value LIKE ?", "%"+searchKeyword+"%") } diff --git a/internal/store/memory.go b/internal/store/memory.go index d966cd4..9c6f240 100644 --- a/internal/store/memory.go +++ b/internal/store/memory.go @@ -37,7 +37,6 @@ func (s *MemoryStore) Close() error { return nil } - // Set stores a key-value pair. func (s *MemoryStore) Set(key string, value []byte, ttl time.Duration) error { s.mu.Lock() @@ -147,7 +146,7 @@ func (s *MemoryStore) SetNX(key string, value []byte, ttl time.Duration) (bool, // --- HASH operations --- -func (s *MemoryStore) HSet(key, field string, value any) error { +func (s *MemoryStore) HSet(key string, values map[string]any) error { s.mu.Lock() defer s.mu.Unlock() @@ -164,7 +163,9 @@ func (s *MemoryStore) HSet(key, field string, value any) error { } } - hash[field] = fmt.Sprint(value) + for field, value := range values { + hash[field] = fmt.Sprint(value) + } return nil } diff --git a/internal/store/redis.go b/internal/store/redis.go index 10b024c..d941874 100644 --- a/internal/store/redis.go +++ b/internal/store/redis.go @@ -61,8 +61,8 @@ func (s *RedisStore) Close() error { // --- HASH operations --- -func (s *RedisStore) HSet(key, field string, value any) error { - return s.client.HSet(context.Background(), key, field, value).Err() +func (s *RedisStore) HSet(key string, values map[string]any) error { + return s.client.HSet(context.Background(), key, values).Err() } func (s *RedisStore) HGetAll(key string) (map[string]string, error) { diff --git a/internal/store/store.go b/internal/store/store.go index a5ee81b..57b42f3 100644 --- a/internal/store/store.go +++ b/internal/store/store.go @@ -33,7 +33,7 @@ type Store interface { SetNX(key string, value []byte, ttl time.Duration) (bool, error) // HASH operations - HSet(key, field string, value any) error + HSet(key string, values map[string]any) error HGetAll(key string) (map[string]string, error) HIncrBy(key, field string, incr int64) (int64, error)