435 lines
12 KiB
Go
435 lines
12 KiB
Go
package database
|
||
|
||
import (
|
||
"context"
|
||
"crypto/sha256"
|
||
"encoding/json"
|
||
"fmt"
|
||
"time"
|
||
|
||
"github.com/redis/go-redis/v9"
|
||
)
|
||
|
||
const (
|
||
BlacklistKeyPrefix = "blacklist:token:"
|
||
InspirationFlowKeyPrefix = "inspiration_flow:"
|
||
AssetLikersKeyPrefix = "asset_likers:"
|
||
ExpiringAssetsKey = "expiring_assets"
|
||
ExpiringAssetsPerStarKey = "expiring_assets:star:"
|
||
)
|
||
|
||
// AssetLikersCache 缓存数据结构
|
||
type AssetLikersCache struct {
|
||
Users []AssetLikerWithTotal `json:"users"` // 完整用户列表（按 liked_at DESC）
|
||
Total int64 `json:"total"` // 总数
|
||
UpdatedAt int64 `json:"updated_at"` // 缓存更新时间
|
||
}
|
||
|
||
// AssetLikerWithTotal 用户+点赞时间
|
||
type AssetLikerWithTotal struct {
|
||
UserID int64 `json:"user_id"`
|
||
Nickname string `json:"nickname"`
|
||
Avatar string `json:"avatar"`
|
||
FanLevel int32 `json:"fan_level"`
|
||
LikedAt int64 `json:"liked_at"`
|
||
StarID int64 `json:"star_id"` // 用于 JOIN fan_profiles，缓存时保留
|
||
}
|
||
|
||
// RedisClient Redis 客户端单例
|
||
var RedisClient *redis.Client
|
||
|
||
// Config Redis 配置
|
||
type RedisConfig struct {
|
||
Host string
|
||
Port int
|
||
Password string
|
||
DB int
|
||
}
|
||
|
||
// InitRedis 初始化 Redis 连接
|
||
func InitRedis(cfg RedisConfig) error {
|
||
RedisClient = redis.NewClient(&redis.Options{
|
||
Addr: fmt.Sprintf("%s:%d", cfg.Host, cfg.Port),
|
||
Password: cfg.Password,
|
||
DB: cfg.DB,
|
||
})
|
||
|
||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||
defer cancel()
|
||
|
||
if err := RedisClient.Ping(ctx).Err(); err != nil {
|
||
return fmt.Errorf("failed to connect to redis: %w", err)
|
||
}
|
||
|
||
return nil
|
||
}
|
||
|
||
// CloseRedis 关闭 Redis 连接
|
||
func CloseRedis() error {
|
||
if RedisClient != nil {
|
||
return RedisClient.Close()
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// GetRedis 获取 Redis 客户端实例
|
||
func GetRedis() *redis.Client {
|
||
return RedisClient
|
||
}
|
||
|
||
// RedisHealthCheck 健康检查
|
||
func RedisHealthCheck() error {
|
||
if RedisClient == nil {
|
||
return fmt.Errorf("redis client is not initialized")
|
||
}
|
||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||
defer cancel()
|
||
return RedisClient.Ping(ctx).Err()
|
||
}
|
||
|
||
// BlacklistEntry 黑名单条目
|
||
type BlacklistEntry struct {
|
||
UserID int64 `json:"user_id"`
|
||
Reason string `json:"reason"`
|
||
}
|
||
|
||
// tokenToHash 将 Token 转换为 SHA256 哈希作为 Key
|
||
func tokenToHash(token string) string {
|
||
hash := sha256.Sum256([]byte(token))
|
||
return fmt.Sprintf("%x", hash)
|
||
}
|
||
|
||
// AddToBlacklist 添加 Token 到黑名单
|
||
func AddToBlacklist(ctx context.Context, token string, userID int64, banReason string, ttl time.Duration) error {
|
||
if token == "" {
|
||
return fmt.Errorf("token is empty")
|
||
}
|
||
if RedisClient == nil {
|
||
return fmt.Errorf("redis client is not initialized")
|
||
}
|
||
|
||
key := BlacklistKeyPrefix + tokenToHash(token)
|
||
entry := BlacklistEntry{UserID: userID, Reason: banReason}
|
||
value, err := json.Marshal(entry)
|
||
if err != nil {
|
||
return fmt.Errorf("failed to marshal blacklist entry: %w", err)
|
||
}
|
||
|
||
return RedisClient.Set(ctx, key, value, ttl).Err()
|
||
}
|
||
|
||
// IsBlacklisted 检查 Token 是否在黑名单
|
||
func IsBlacklisted(ctx context.Context, token string) (bool, int64, string, error) {
|
||
if token == "" {
|
||
return false, 0, "", nil
|
||
}
|
||
if RedisClient == nil {
|
||
return false, 0, "", fmt.Errorf("redis client is not initialized")
|
||
}
|
||
|
||
key := BlacklistKeyPrefix + tokenToHash(token)
|
||
value, err := RedisClient.Get(ctx, key).Result()
|
||
if err == redis.Nil {
|
||
return false, 0, "", nil
|
||
}
|
||
if err != nil {
|
||
return false, 0, "", err
|
||
}
|
||
|
||
var entry BlacklistEntry
|
||
if err := json.Unmarshal([]byte(value), &entry); err != nil {
|
||
return false, 0, "", fmt.Errorf("failed to unmarshal blacklist entry: %w", err)
|
||
}
|
||
|
||
return true, entry.UserID, entry.Reason, nil
|
||
}
|
||
|
||
// RemoveFromBlacklist 从黑名单移除 Token（用于解封）
|
||
func RemoveFromBlacklist(ctx context.Context, token string) error {
|
||
if token == "" {
|
||
return nil
|
||
}
|
||
if RedisClient == nil {
|
||
return fmt.Errorf("redis client is not initialized")
|
||
}
|
||
|
||
key := BlacklistKeyPrefix + tokenToHash(token)
|
||
return RedisClient.Del(ctx, key).Err()
|
||
}
|
||
|
||
// InspirationFlowCacheEntry 单个展品缓存数据
|
||
type InspirationFlowCacheEntry struct {
|
||
AssetID int64 `json:"asset_id"`
|
||
Name string `json:"name"`
|
||
CoverURL string `json:"cover_url"`
|
||
LikeCount int32 `json:"like_count"`
|
||
OwnerNickname string `json:"owner_nickname"`
|
||
OwnerAvatar string `json:"owner_avatar"`
|
||
Span int32 `json:"span"`
|
||
MaterialType string `json:"material_type"`
|
||
}
|
||
|
||
// InspirationFlowCache 会话缓存结构
|
||
type InspirationFlowCache struct {
|
||
DisplayedIDs []int64 `json:"displayed_ids"` // 已展示ID列表
|
||
History map[int64]InspirationFlowCacheEntry `json:"history"` // 历史数据详情
|
||
}
|
||
|
||
// InspirationFlowKey 生成灵感瀑布流缓存 Key
|
||
func InspirationFlowKey(starID int64, sessionID string) string {
|
||
return fmt.Sprintf("%s%d:%s", InspirationFlowKeyPrefix, starID, sessionID)
|
||
}
|
||
|
||
// GetInspirationFlowCache 获取灵感瀑布流会话缓存
|
||
func GetInspirationFlowCache(ctx context.Context, starID int64, sessionID string) (*InspirationFlowCache, error) {
|
||
if RedisClient == nil {
|
||
return nil, fmt.Errorf("redis client is not initialized")
|
||
}
|
||
|
||
key := InspirationFlowKey(starID, sessionID)
|
||
data, err := RedisClient.Get(ctx, key).Result()
|
||
if err == redis.Nil {
|
||
return &InspirationFlowCache{
|
||
DisplayedIDs: []int64{},
|
||
History: make(map[int64]InspirationFlowCacheEntry),
|
||
}, nil
|
||
}
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
var cache InspirationFlowCache
|
||
if err := json.Unmarshal([]byte(data), &cache); err != nil {
|
||
return nil, err
|
||
}
|
||
return &cache, nil
|
||
}
|
||
|
||
// SaveInspirationFlowCache 保存灵感瀑布流会话缓存
|
||
func SaveInspirationFlowCache(ctx context.Context, starID int64, sessionID string, cache *InspirationFlowCache, ttl time.Duration) error {
|
||
if RedisClient == nil {
|
||
return fmt.Errorf("redis client is not initialized")
|
||
}
|
||
|
||
key := InspirationFlowKey(starID, sessionID)
|
||
data, err := json.Marshal(cache)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
|
||
return RedisClient.Set(ctx, key, data, ttl).Err()
|
||
}
|
||
|
||
// AddToInspirationFlowCache 添加展品到会话缓存
|
||
func AddToInspirationFlowCache(ctx context.Context, starID int64, sessionID string, entry InspirationFlowCacheEntry, ttl time.Duration) error {
|
||
cache, err := GetInspirationFlowCache(ctx, starID, sessionID)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
|
||
// 检查是否已存在
|
||
for _, id := range cache.DisplayedIDs {
|
||
if id == entry.AssetID {
|
||
return nil // 已存在，跳过
|
||
}
|
||
}
|
||
|
||
// 添加到已展示列表
|
||
cache.DisplayedIDs = append(cache.DisplayedIDs, entry.AssetID)
|
||
// 添加到历史详情
|
||
if cache.History == nil {
|
||
cache.History = make(map[int64]InspirationFlowCacheEntry)
|
||
}
|
||
cache.History[entry.AssetID] = entry
|
||
|
||
return SaveInspirationFlowCache(ctx, starID, sessionID, cache, ttl)
|
||
}
|
||
|
||
// GetHistoryPage 获取历史数据的某一页
|
||
func GetHistoryPage(cache *InspirationFlowCache, offset, limit int) []InspirationFlowCacheEntry {
|
||
if cache == nil || cache.History == nil {
|
||
return []InspirationFlowCacheEntry{}
|
||
}
|
||
|
||
// 按展示顺序反向遍历（最新展示的在前面）
|
||
items := make([]InspirationFlowCacheEntry, 0)
|
||
for i := len(cache.DisplayedIDs) - 1; i >= 0; i-- {
|
||
if entry, ok := cache.History[cache.DisplayedIDs[i]]; ok {
|
||
items = append(items, entry)
|
||
}
|
||
}
|
||
|
||
// 分页
|
||
start := offset
|
||
end := offset + limit
|
||
if start >= len(items) {
|
||
return []InspirationFlowCacheEntry{}
|
||
}
|
||
if end > len(items) {
|
||
end = len(items)
|
||
}
|
||
return items[start:end]
|
||
}
|
||
|
||
// AssetLikersKey 生成藏品点赞用户列表缓存 Key
|
||
func AssetLikersKey(assetID int64) string {
|
||
return fmt.Sprintf("%s%d", AssetLikersKeyPrefix, assetID)
|
||
}
|
||
|
||
// GetAssetLikersCache 获取藏品点赞用户列表缓存
|
||
func GetAssetLikersCache(ctx context.Context, assetID int64) (*AssetLikersCache, error) {
|
||
if RedisClient == nil {
|
||
return nil, fmt.Errorf("redis client is not initialized")
|
||
}
|
||
|
||
key := AssetLikersKey(assetID)
|
||
data, err := RedisClient.Get(ctx, key).Result()
|
||
if err == redis.Nil {
|
||
return nil, nil // 缓存不存在
|
||
}
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
var cache AssetLikersCache
|
||
if err := json.Unmarshal([]byte(data), &cache); err != nil {
|
||
return nil, err
|
||
}
|
||
return &cache, nil
|
||
}
|
||
|
||
// SetAssetLikersCache 设置藏品点赞用户列表缓存
|
||
func SetAssetLikersCache(ctx context.Context, assetID int64, cache *AssetLikersCache, ttl time.Duration) error {
|
||
if RedisClient == nil {
|
||
return fmt.Errorf("redis client is not initialized")
|
||
}
|
||
|
||
key := AssetLikersKey(assetID)
|
||
data, err := json.Marshal(cache)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
|
||
return RedisClient.Set(ctx, key, data, ttl).Err()
|
||
}
|
||
|
||
// InvalidateAssetLikersCache 删除藏品点赞用户列表缓存
|
||
func InvalidateAssetLikersCache(ctx context.Context, assetID int64) error {
|
||
if RedisClient == nil {
|
||
return nil // Redis 未初始化时跳过
|
||
}
|
||
|
||
key := AssetLikersKey(assetID)
|
||
return RedisClient.Del(ctx, key).Err()
|
||
}
|
||
|
||
// AddExpiringAsset 添加到过期资产 ZSET (score = expire_at)
|
||
func AddExpiringAsset(ctx context.Context, assetID int64, expireAt int64) error {
|
||
if RedisClient == nil {
|
||
return nil
|
||
}
|
||
// 使用 ZADD，score 为过期时间戳
|
||
return RedisClient.ZAdd(ctx, ExpiringAssetsKey, redis.Z{
|
||
Score: float64(expireAt),
|
||
Member: assetID,
|
||
}).Err()
|
||
}
|
||
|
||
// AddExpiringAssetToStar 添加到指定 star 的过期资产 ZSET
|
||
func AddExpiringAssetToStar(ctx context.Context, starID int64, assetID int64, expireAt int64) error {
|
||
if RedisClient == nil {
|
||
return nil
|
||
}
|
||
key := fmt.Sprintf("%s%d", ExpiringAssetsPerStarKey, starID)
|
||
return RedisClient.ZAdd(ctx, key, redis.Z{
|
||
Score: float64(expireAt),
|
||
Member: assetID,
|
||
}).Err()
|
||
}
|
||
|
||
// RemoveExpiringAsset 从过期资产 ZSET 移除
|
||
func RemoveExpiringAsset(ctx context.Context, assetID int64) error {
|
||
if RedisClient == nil {
|
||
return nil
|
||
}
|
||
return RedisClient.ZRem(ctx, ExpiringAssetsKey, assetID).Err()
|
||
}
|
||
|
||
// RemoveExpiringAssetFromStar 从指定 star 的过期资产 ZSET 移除
|
||
func RemoveExpiringAssetFromStar(ctx context.Context, starID int64, assetID int64) error {
|
||
if RedisClient == nil {
|
||
return nil
|
||
}
|
||
key := fmt.Sprintf("%s%d", ExpiringAssetsPerStarKey, starID)
|
||
return RedisClient.ZRem(ctx, key, assetID).Err()
|
||
}
|
||
|
||
// GetExpiredAssets 获取已过期的资产ID列表 (score <= now)
|
||
func GetExpiredAssets(ctx context.Context, now int64) ([]int64, error) {
|
||
if RedisClient == nil {
|
||
return nil, fmt.Errorf("redis client not initialized")
|
||
}
|
||
// ZRANGEBYSCORE 获取 score 在 [0, now] 范围内的成员
|
||
results, err := RedisClient.ZRangeByScore(ctx, ExpiringAssetsKey, &redis.ZRangeBy{
|
||
Min: "0",
|
||
Max: fmt.Sprintf("%d", now),
|
||
}).Result()
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
assetIDs := make([]int64, 0, len(results))
|
||
for _, r := range results {
|
||
var id int64
|
||
if _, err := fmt.Sscanf(r, "%d", &id); err == nil {
|
||
assetIDs = append(assetIDs, id)
|
||
}
|
||
}
|
||
return assetIDs, nil
|
||
}
|
||
|
||
// GetExpiredAssetsFromStar 获取指定 star 已过期的资产ID列表
|
||
func GetExpiredAssetsFromStar(ctx context.Context, starID int64, now int64) ([]int64, error) {
|
||
if RedisClient == nil {
|
||
return nil, fmt.Errorf("redis client not initialized")
|
||
}
|
||
key := fmt.Sprintf("%s%d", ExpiringAssetsPerStarKey, starID)
|
||
results, err := RedisClient.ZRangeByScore(ctx, key, &redis.ZRangeBy{
|
||
Min: "0",
|
||
Max: fmt.Sprintf("%d", now),
|
||
}).Result()
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
assetIDs := make([]int64, 0, len(results))
|
||
for _, r := range results {
|
||
var id int64
|
||
if _, err := fmt.Sscanf(r, "%d", &id); err == nil {
|
||
assetIDs = append(assetIDs, id)
|
||
}
|
||
}
|
||
return assetIDs, nil
|
||
}
|
||
|
||
// GetAllExpiringAssets 获取所有待处理的资产ID（未过期的也返回，用于校验）
|
||
func GetAllExpiringAssets(ctx context.Context) ([]int64, error) {
|
||
if RedisClient == nil {
|
||
return nil, fmt.Errorf("redis client not initialized")
|
||
}
|
||
results, err := RedisClient.ZRange(ctx, ExpiringAssetsKey, 0, -1).Result()
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
assetIDs := make([]int64, 0, len(results))
|
||
for _, r := range results {
|
||
var id int64
|
||
if _, err := fmt.Sscanf(r, "%d", &id); err == nil {
|
||
assetIDs = append(assetIDs, id)
|
||
}
|
||
}
|
||
return assetIDs, nil
|
||
}
|
||
|