feat: Optimize massive keys (#371)

tbphp · web-flow · commit e4657cad3a22 · 2026-01-24T09:21:19.000+08:00
* feat: Support txt file upload

* feat: Optimize Batch Cache

* feat: Keys query index
diff --git a/internal/handler/key_handler.go b/internal/handler/key_handler.go
@@ -5,7 +5,9 @@ import (
 	app_errors "gpt-load/internal/errors"
 	"gpt-load/internal/models"
 	"gpt-load/internal/response"
+	"io"
 	"log"
+	"path/filepath"
 	"strconv"
 	"strings"
 	"time"
@@ -107,24 +109,79 @@ func (s *Server) AddMultipleKeys(c *gin.Context) {
 	response.Success(c, result)
 }
 
-// AddMultipleKeysAsync handles creating new keys from a text block within a specific group.
+// AddMultipleKeysAsync handles creating new keys from a text block or file within a specific group.
 func (s *Server) AddMultipleKeysAsync(c *gin.Context) {
-	var req KeyTextRequest
-	if err := c.ShouldBindJSON(&req); err != nil {
-		response.Error(c, app_errors.NewAPIError(app_errors.ErrInvalidJSON, err.Error()))
-		return
+	var groupID uint
+	var keysText string
+
+	// Check content type to determine if it's a file upload or JSON request
+	contentType := c.ContentType()
+
+	if strings.Contains(contentType, "multipart/form-data") {
+		// Handle file upload
+		groupIDStr := c.PostForm("group_id")
+		if groupIDStr == "" {
+			response.ErrorI18nFromAPIError(c, app_errors.ErrBadRequest, "validation.group_id_required")
+			return
+		}
+
+		groupIDInt, err := strconv.Atoi(groupIDStr)
+		if err != nil || groupIDInt <= 0 {
+			response.ErrorI18nFromAPIError(c, app_errors.ErrBadRequest, "validation.invalid_group_id_format")
+			return
+		}
+		groupID = uint(groupIDInt)
+
+		// Get uploaded file
+		file, err := c.FormFile("file")
+		if err != nil {
+			response.ErrorI18nFromAPIError(c, app_errors.ErrBadRequest, "validation.file_required")
+			return
+		}
+
+		// Validate file extension
+		ext := strings.ToLower(filepath.Ext(file.Filename))
+		if ext != ".txt" {
+			response.ErrorI18nFromAPIError(c, app_errors.ErrValidation, "validation.only_txt_supported")
+			return
+		}
+
+		// Read file content
+		fileContent, err := file.Open()
+		if err != nil {
+			response.ErrorI18nFromAPIError(c, app_errors.ErrBadRequest, "validation.failed_to_open_file")
+			return
+		}
+		defer fileContent.Close()
+
+		// Read file content as string using io.ReadAll
+		buf, err := io.ReadAll(fileContent)
+		if err != nil {
+			response.ErrorI18nFromAPIError(c, app_errors.ErrBadRequest, "validation.failed_to_read_file")
+			return
+		}
+		keysText = string(buf)
+	} else {
+		// Handle JSON request (original behavior)
+		var req KeyTextRequest
+		if err := c.ShouldBindJSON(&req); err != nil {
+			response.Error(c, app_errors.NewAPIError(app_errors.ErrInvalidJSON, err.Error()))
+			return
+		}
+		groupID = req.GroupID
+		keysText = req.KeysText
 	}
 
-	group, ok := s.findGroupByID(c, req.GroupID)
+	group, ok := s.findGroupByID(c, groupID)
 	if !ok {
 		return
 	}
 
-	if !validateKeysText(c, req.KeysText) {
+	if !validateKeysText(c, keysText) {
 		return
 	}
 
-	taskStatus, err := s.KeyImportService.StartImportTask(group, req.KeysText)
+	taskStatus, err := s.KeyImportService.StartImportTask(group, keysText)
 	if err != nil {
 		response.Error(c, app_errors.NewAPIError(app_errors.ErrTaskInProgress, err.Error()))
 		return
diff --git a/internal/i18n/locales/en-US.go b/internal/i18n/locales/en-US.go
@@ -60,6 +60,10 @@ var MessagesEnUS = map[string]string{
 	"validation.group_id_required":       "group_id query parameter is required",
 	"validation.invalid_group_id_format": "Invalid group_id format",
 	"validation.keys_text_empty":         "Keys text cannot be empty",
+	"validation.file_required":           "File is required",
+	"validation.only_txt_supported":      "Only .txt files are supported",
+	"validation.failed_to_open_file":     "Failed to open file",
+	"validation.failed_to_read_file":     "Failed to read file content",
 	"validation.invalid_group_type":      "Invalid group type, must be 'standard' or 'aggregate'",
 	"validation.sub_groups_required":     "Aggregate group must contain at least one sub-group",
 	"validation.invalid_sub_group_id":    "Invalid sub-group ID",
diff --git a/internal/i18n/locales/ja-JP.go b/internal/i18n/locales/ja-JP.go
@@ -60,6 +60,10 @@ var MessagesJaJP = map[string]string{
 	"validation.group_id_required":       "group_idクエリパラメータが必要です",
 	"validation.invalid_group_id_format": "無効なgroup_id形式",
 	"validation.keys_text_empty":         "キーテキストは空にできません",
+	"validation.file_required":           "ファイルが必要です",
+	"validation.only_txt_supported":      ".txtファイルのみサポートされています",
+	"validation.failed_to_open_file":     "ファイルを開けませんでした",
+	"validation.failed_to_read_file":     "ファイルの内容を読み取れませんでした",
 	"validation.invalid_group_type":      "無効なグループタイプ、'standard'または'aggregate'である必要があります",
 	"validation.sub_groups_required":     "集約グループには少なくとも1つのサブグループが必要です",
 	"validation.invalid_sub_group_id":    "無効なサブグループID",
diff --git a/internal/i18n/locales/zh-CN.go b/internal/i18n/locales/zh-CN.go
@@ -60,6 +60,10 @@ var MessagesZhCN = map[string]string{
 	"validation.group_id_required":       "需要提供group_id参数",
 	"validation.invalid_group_id_format": "无效的group_id格式",
 	"validation.keys_text_empty":         "密钥文本不能为空",
+	"validation.file_required":           "需要上传文件",
+	"validation.only_txt_supported":      "仅支持.txt文件",
+	"validation.failed_to_open_file":     "无法打开文件",
+	"validation.failed_to_read_file":     "无法读取文件内容",
 	"validation.invalid_group_type":      "无效的分组类型，必须为'standard'或'aggregate'",
 	"validation.sub_groups_required":     "聚合分组必须包含至少一个子分组",
 	"validation.invalid_sub_group_id":    "无效的子分组ID",
diff --git a/internal/keypool/provider.go b/internal/keypool/provider.go
@@ -242,7 +242,7 @@ func (p *KeyProvider) LoadKeysFromDB() error {
 
 	// 1. 分批从数据库加载并使用 Pipeline 写入 Redis
 	allActiveKeyIDs := make(map[uint][]any)
-	batchSize := 1000
+	batchSize := 10000
 	var batchKeys []*models.APIKey
 
 	err := p.db.Model(&models.APIKey{}).FindInBatches(&batchKeys, batchSize, func(tx *gorm.DB, batch int) error {
@@ -308,13 +308,8 @@ func (p *KeyProvider) AddKeys(groupID uint, keys []models.APIKey) error {
 			return err
 		}
 
-		for _, key := range keys {
-			if err := p.addKeyToStore(&key); err != nil {
-				logrus.WithFields(logrus.Fields{"keyID": key.ID, "error": err}).Error("Failed to add key to store after DB creation, rolling back transaction")
-				return err
-			}
-		}
-		return nil
+		// 使用批量方法添加到缓存
+		return p.addKeysToCacheBatch(groupID, keys)
 	})
 
 	return err
@@ -577,6 +572,48 @@ func (p *KeyProvider) addKeyToStore(key *models.APIKey) error {
 	return nil
 }
 
+// addKeysToCacheBatch 批量添加密钥到缓存（用于批量导入场景）
+func (p *KeyProvider) addKeysToCacheBatch(groupID uint, keys []models.APIKey) error {
+	if len(keys) == 0 {
+		return nil
+	}
+
+	// 1. 批量 HSet 密钥详情
+	if pipeliner, ok := p.store.(store.RedisPipeliner); ok {
+		// Redis: 使用 Pipeline 批量操作
+		pipe := pipeliner.Pipeline()
+		for i := range keys {
+			keyHashKey := fmt.Sprintf("key:%d", keys[i].ID)
+			pipe.HSet(keyHashKey, p.apiKeyToMap(&keys[i]))
+		}
+		if err := pipe.Exec(); err != nil {
+			return fmt.Errorf("failed to batch HSet keys: %w", err)
+		}
+	} else {
+		// MemoryStore: 降级为逐个 HSet
+		for i := range keys {
+			keyHashKey := fmt.Sprintf("key:%d", keys[i].ID)
+			if err := p.store.HSet(keyHashKey, p.apiKeyToMap(&keys[i])); err != nil {
+				return fmt.Errorf("failed to HSet key %d: %w", keys[i].ID, err)
+			}
+		}
+	}
+
+	// 2. 收集所有密钥 ID
+	activeKeysListKey := fmt.Sprintf("group:%d:active_keys", groupID)
+	activeKeyIDs := make([]any, len(keys))
+	for i := range keys {
+		activeKeyIDs[i] = keys[i].ID
+	}
+
+	// 3. 批量 LPush 活跃密钥
+	if err := p.store.LPush(activeKeysListKey, activeKeyIDs...); err != nil {
+		return fmt.Errorf("failed to batch LPush keys to group %d: %w", groupID, err)
+	}
+
+	return nil
+}
+
 // removeKeyFromStore is a helper to remove a single key from the cache.
 func (p *KeyProvider) removeKeyFromStore(keyID, groupID uint) error {
 	activeKeysListKey := fmt.Sprintf("group:%d:active_keys", groupID)
diff --git a/internal/models/types.go b/internal/models/types.go
@@ -115,11 +115,10 @@ type APIKey struct {
 	KeyValue     string     `gorm:"type:text;not null" json:"key_value"`
 	KeyHash      string     `gorm:"type:varchar(128);index" json:"key_hash"`
 	GroupID      uint       `gorm:"not null;index" json:"group_id"`
-	Status       string     `gorm:"type:varchar(50);not null;default:'active'" json:"status"`
+	Status       string     `gorm:"type:varchar(50);not null;default:'active';index" json:"status"`
 	Notes        string     `gorm:"type:varchar(255);default:''" json:"notes"`
 	RequestCount int64      `gorm:"not null;default:0" json:"request_count"`
 	FailureCount int64      `gorm:"not null;default:0" json:"failure_count"`
-	LastUsedAt   *time.Time `json:"last_used_at"`
 	CreatedAt    time.Time  `json:"created_at"`
 	UpdatedAt    time.Time  `json:"updated_at"`
 }
diff --git a/internal/services/key_service.go b/internal/services/key_service.go
@@ -301,7 +301,7 @@ func (s *KeyService) ListKeysInGroupQuery(groupID uint, statusFilter string, sea
 		query = query.Where("key_hash = ?", searchHash)
 	}
 
-	query = query.Order("last_used_at desc, updated_at desc")
+	query = query.Order("id desc")
 
 	return query
 }
diff --git a/internal/services/request_log_service.go b/internal/services/request_log_service.go
@@ -227,7 +227,6 @@ func (s *RequestLogService) writeLogsToDB(logs []*models.RequestLog) error {
 			if err := tx.Model(&models.APIKey{}).Where("key_hash IN ?", keyHashes).
 				Updates(map[string]any{
 					"request_count": gorm.Expr(caseStmt.String()),
-					"last_used_at":  time.Now(),
 				}).Error; err != nil {
 				return fmt.Errorf("failed to batch update api_key stats: %w", err)
 			}
diff --git a/web/src/api/keys.ts b/web/src/api/keys.ts
@@ -102,17 +102,25 @@ export const keysApi = {
   },
 
   // 异步批量添加密钥
-  async addKeysAsync(group_id: number, keys_text: string): Promise<TaskInfo> {
-    const res = await http.post(
-      "/keys/add-async",
-      {
-        group_id,
-        keys_text,
-      },
-      {
-        hideMessage: true,
-      }
-    );
+  async addKeysAsync(group_id: number, keys_text?: string, file?: File): Promise<TaskInfo> {
+    let requestData: FormData | { group_id: number; keys_text: string };
+    const config: { hideMessage: boolean; headers?: { "Content-Type": string } } = {
+      hideMessage: true,
+    };
+
+    if (file) {
+      // File upload mode
+      const formData = new FormData();
+      formData.append("group_id", group_id.toString());
+      formData.append("file", file);
+      requestData = formData;
+      config.headers = { "Content-Type": "multipart/form-data" };
+    } else {
+      // Text input mode
+      requestData = { group_id, keys_text: keys_text || "" };
+    }
+
+    const res = await http.post("/keys/add-async", requestData, config);
     return res.data;
   },
 
diff --git a/web/src/components/keys/KeyCreateDialog.vue b/web/src/components/keys/KeyCreateDialog.vue
diff --git a/web/src/locales/en-US.ts b/web/src/locales/en-US.ts
diff --git a/web/src/locales/ja-JP.ts b/web/src/locales/ja-JP.ts
diff --git a/web/src/locales/zh-CN.ts b/web/src/locales/zh-CN.ts

Original file line number	Diff line number	Diff line change
`@@ -301,7 +301,7 @@ func (s *KeyService) ListKeysInGroupQuery(groupID uint, statusFilter string, sea`
`301`	`301`	`query = query.Where("key_hash = ?", searchHash)`
`302`	`302`	`}`
`303`	`303`
`304`		`- query = query.Order("last_used_at desc, updated_at desc")`
	`304`	`+ query = query.Order("id desc")`
`305`	`305`
`306`	`306`	`return query`
`307`	`307`	`}`
Original file line number	Diff line number	Diff line change
`@@ -227,7 +227,6 @@ func (s RequestLogService) writeLogsToDB(logs []models.RequestLog) error {`
`227`	`227`	`if err := tx.Model(&models.APIKey{}).Where("key_hash IN ?", keyHashes).`
`228`	`228`	`Updates(map[string]any{`
`229`	`229`	`"request_count": gorm.Expr(caseStmt.String()),`
`230`		`- "last_used_at": time.Now(),`
`231`	`230`	`}).Error; err != nil {`
`232`	`231`	`return fmt.Errorf("failed to batch update api_key stats: %w", err)`
`233`	`232`	`}`