From c214b8f66d7acb70e27807fd7bc0bbc33dce242b Mon Sep 17 00:00:00 2001
From: tbphp <admin@tbphp.net>
Date: Wed, 9 Jul 2025 16:31:04 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E6=8B=89=E5=8F=96key=E6=95=B0=E6=8D=AE?=
 =?UTF-8?q?=E5=A2=9E=E5=8A=A0=E5=88=9D=E5=A7=8B=E6=A0=87=E8=AE=B0?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .gitignore                   | 13 -------------
 internal/keypool/provider.go | 19 ++++++++++++++++++-
 2 files changed, 18 insertions(+), 14 deletions(-)

diff --git a/.gitignore b/.gitignore
index 8854d7e..655e33d 100644
--- a/.gitignore
+++ b/.gitignore
@@ -84,21 +84,10 @@ node_modules/
 .cache
 .parcel-cache
 
-# Next.js build output
-.next
-
-# Nuxt.js build / generate output
-.nuxt
-dist
-
 # Gatsby files
 .cache/
 public
 
-# Storybook build outputs
-.out
-.storybook-out
-
 # Temporary folders
 tmp/
 temp/
@@ -226,5 +215,3 @@ charts/*.tgz
 run-local.sh
 start-local.sh
 dev-setup.sh
-
-cmd/gpt-load/dist
diff --git a/internal/keypool/provider.go b/internal/keypool/provider.go
index ff2381c..85f605a 100644
--- a/internal/keypool/provider.go
+++ b/internal/keypool/provider.go
@@ -187,13 +187,26 @@ func (p *KeyProvider) handleFailure(keyID uint, keyHashKey, activeKeysListKey st
 
 // LoadKeysFromDB 从数据库加载所有分组和密钥，并填充到 Store 中。
 func (p *KeyProvider) LoadKeysFromDB() error {
+	initFlagKey := "initialization:db_keys_loaded"
+
+	exists, err := p.store.Exists(initFlagKey)
+	if err != nil {
+		return fmt.Errorf("failed to check initialization flag: %w", err)
+	}
+
+	if exists {
+		logrus.Debug("Keys have already been loaded into the store. Skipping.")
+		return nil
+	}
+
+	logrus.Debug("First time startup, loading keys from DB...")
 
 	// 1. 分批从数据库加载并使用 Pipeline 写入 Redis
 	allActiveKeyIDs := make(map[uint][]any)
 	batchSize := 1000
 	var batchKeys []*models.APIKey
 
-	err := p.db.Model(&models.APIKey{}).FindInBatches(&batchKeys, batchSize, func(tx *gorm.DB, batch int) error {
+	err = p.db.Model(&models.APIKey{}).FindInBatches(&batchKeys, batchSize, func(tx *gorm.DB, batch int) error {
 		logrus.Debugf("Processing batch %d with %d keys...", batch, len(batchKeys))
 
 		var pipeline store.Pipeliner
@@ -242,6 +255,10 @@ func (p *KeyProvider) LoadKeysFromDB() error {
 		}
 	}
 
+	if err := p.store.Set(initFlagKey, []byte("1"), 0); err != nil {
+		logrus.WithField("flagKey", initFlagKey).Error("Failed to set initialization flag after loading keys")
+	}
+
 	return nil
 }