111 lines
4.3 KiB
Go
111 lines
4.3 KiB
Go
// Filename: internal/repository/repository.go
|
|
package repository
|
|
|
|
import (
|
|
"context"
|
|
"gemini-balancer/internal/config"
|
|
"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(ctx context.Context, group *models.KeyGroup) (*models.APIKey, *models.GroupAPIKeyMapping, error)
|
|
SelectOneActiveKeyFromBasePool(ctx context.Context, 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)
|
|
GetKeysByGroup(groupID uint) ([]models.APIKey, error)
|
|
CountByGroup(groupID uint) (int64, error)
|
|
|
|
// --- 多对多关系管理 --- key_mapping
|
|
LinkKeysToGroup(ctx context.Context, groupID uint, keyIDs []uint) error
|
|
UnlinkKeysFromGroup(ctx context.Context, groupID uint, keyIDs []uint) (unlinkedCount int64, err error)
|
|
GetGroupsForKey(ctx context.Context, keyID uint) ([]uint, error)
|
|
GetMapping(groupID, keyID uint) (*models.GroupAPIKeyMapping, error)
|
|
UpdateMapping(ctx context.Context, 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(ctx context.Context) error
|
|
HandleCacheUpdateEventBatch(ctx context.Context, 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(ctx context.Context, keyID uint, status models.MasterAPIKeyStatus) error
|
|
HardDeleteSoftDeletedBefore(date time.Time) (int64, error)
|
|
|
|
// --- 轮询策略的"写"操作 --- key_writer
|
|
UpdateKeyUsageTimestamp(ctx context.Context, groupID, keyID uint)
|
|
SyncKeyStatusInPollingCaches(ctx context.Context, groupID, keyID uint, newStatus models.APIKeyStatus)
|
|
HandleCacheUpdateEvent(ctx context.Context, groupID, keyID uint, newStatus models.APIKeyStatus)
|
|
UpdateKeyStatusAfterRequest(ctx context.Context, 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
|
|
config *config.Config
|
|
}
|
|
|
|
type gormGroupRepository struct {
|
|
db *gorm.DB
|
|
}
|
|
|
|
func NewKeyRepository(db *gorm.DB, s store.Store, logger *logrus.Logger, crypto *crypto.Service, cfg *config.Config) KeyRepository {
|
|
return &gormKeyRepository{
|
|
db: db,
|
|
store: s,
|
|
logger: logger.WithField("component", "repository.key🔗"),
|
|
crypto: crypto,
|
|
config: cfg,
|
|
}
|
|
}
|
|
|
|
func NewGroupRepository(db *gorm.DB) GroupRepository {
|
|
return &gormGroupRepository{db: db}
|
|
}
|