Files
gemini-banlancer/internal/repository/repository.go
2025-11-21 19:33:05 +08:00

109 lines
4.1 KiB
Go

// Filename: internal/repository/repository.go
package repository
import (
"gemini-balancer/internal/crypto"
"gemini-balancer/internal/errors"
"gemini-balancer/internal/models"
"gemini-balancer/internal/store"
"io"
"time"
"github.com/sirupsen/logrus"
"gorm.io/gorm"
)
// BasePool 虚拟的临时资源池,用于智能聚合模式。
type BasePool struct {
CandidateGroups []*models.KeyGroup
PollingStrategy models.PollingStrategy
Protocol models.ProtocolType
}
type KeyRepository interface {
// --- 核心选取与调度 --- key_selector
SelectOneActiveKey(group *models.KeyGroup) (*models.APIKey, *models.GroupAPIKeyMapping, error)
SelectOneActiveKeyFromBasePool(pool *BasePool) (*models.APIKey, *models.KeyGroup, error)
// --- 加密与解密 --- key_crud
Decrypt(key *models.APIKey) error
DecryptBatch(keys []models.APIKey) error
// --- 基础增删改查 --- key_crud
AddKeys(keys []models.APIKey) ([]models.APIKey, error)
Update(key *models.APIKey) error
HardDeleteByID(id uint) error
HardDeleteByValues(keyValues []string) (int64, error)
GetKeyByID(id uint) (*models.APIKey, error)
GetKeyByValue(keyValue string) (*models.APIKey, error)
GetKeysByValues(keyValues []string) ([]models.APIKey, error)
GetKeysByIDs(ids []uint) ([]models.APIKey, error) // [新增] 根据一组主键ID批量获取Key
GetKeysByGroup(groupID uint) ([]models.APIKey, error)
CountByGroup(groupID uint) (int64, error)
// --- 多对多关系管理 --- key_mapping
LinkKeysToGroup(groupID uint, keyIDs []uint) error
UnlinkKeysFromGroup(groupID uint, keyIDs []uint) (unlinkedCount int64, err error)
GetGroupsForKey(keyID uint) ([]uint, error)
GetMapping(groupID, keyID uint) (*models.GroupAPIKeyMapping, error)
UpdateMapping(mapping *models.GroupAPIKeyMapping) error
GetPaginatedKeysAndMappingsByGroup(params *models.APIKeyQueryParams) ([]*models.APIKeyDetails, int64, error)
GetKeysByValuesAndGroupID(values []string, groupID uint) ([]models.APIKey, error)
FindKeyValuesByStatus(groupID uint, statuses []string) ([]string, error)
FindKeyIDsByStatus(groupID uint, statuses []string) ([]uint, error)
GetKeyStringsByGroupAndStatus(groupID uint, statuses []string) ([]string, error)
UpdateMappingWithoutCache(mapping *models.GroupAPIKeyMapping) error
// --- 缓存管理 --- key_cache
LoadAllKeysToStore() error
HandleCacheUpdateEventBatch(mappings []*models.GroupAPIKeyMapping) error
// --- 维护与后台任务 --- key_maintenance
StreamKeysToWriter(groupID uint, statusFilter string, writer io.Writer) error
UpdateMasterStatusByValues(keyValues []string, newStatus models.MasterAPIKeyStatus) (int64, error)
UpdateMasterStatusByID(keyID uint, newStatus models.MasterAPIKeyStatus) error
DeleteOrphanKeys() (int64, error)
DeleteOrphanKeysTx(tx *gorm.DB) (int64, error)
GetActiveMasterKeys() ([]*models.APIKey, error)
UpdateAPIKeyStatus(keyID uint, status models.MasterAPIKeyStatus) error
HardDeleteSoftDeletedBefore(date time.Time) (int64, error)
// --- 轮询策略的"写"操作 --- key_writer
UpdateKeyUsageTimestamp(groupID, keyID uint)
// 同步更新缓存,供核心业务使用
SyncKeyStatusInPollingCaches(groupID, keyID uint, newStatus models.APIKeyStatus)
// 异步更新缓存,供事件订阅者使用
HandleCacheUpdateEvent(groupID, keyID uint, newStatus models.APIKeyStatus)
UpdateKeyStatusAfterRequest(group *models.KeyGroup, key *models.APIKey, success bool, apiErr *errors.APIError)
}
type GroupRepository interface {
GetGroupByName(name string) (*models.KeyGroup, error)
GetAllGroups() ([]*models.KeyGroup, error)
UpdateOrderInTransaction(orders map[uint]int) error
}
type gormKeyRepository struct {
db *gorm.DB
store store.Store
logger *logrus.Entry
crypto *crypto.Service
}
type gormGroupRepository struct {
db *gorm.DB
}
func NewKeyRepository(db *gorm.DB, s store.Store, logger *logrus.Logger, crypto *crypto.Service) KeyRepository {
return &gormKeyRepository{
db: db,
store: s,
logger: logger.WithField("component", "repository.key🔗"),
crypto: crypto,
}
}
func NewGroupRepository(db *gorm.DB) GroupRepository {
return &gormGroupRepository{db: db}
}