118 lines
3.5 KiB
Python
118 lines
3.5 KiB
Python
"""Redis Pub/Sub 消息通道模块
|
|
|
|
提供基于 Redis 的发布/订阅机制，用于实时推送提醒消息到 WebSocket 客户端。
|
|
|
|
主要功能：
|
|
- publish: 发布提醒消息到指定用户的频道
|
|
- subscribe: 订阅用户的提醒频道
|
|
- unsubscribe: 取消订阅
|
|
|
|
依赖：
|
|
- redis.asyncio: Redis 异步客户端
|
|
"""
|
|
|
|
import json
|
|
import logging
|
|
from typing import AsyncGenerator, Optional
|
|
|
|
import redis.asyncio as redis
|
|
|
|
from backend.app.core.config import get_settings
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def _get_redis_url() -> str:
|
|
"""获取 Redis 连接 URL"""
|
|
settings = get_settings()
|
|
password_part = f":{settings.redis_password}@" if settings.redis_password else ""
|
|
return f"redis://{password_part}{settings.redis_host}:{settings.redis_port}/{settings.redis_db}"
|
|
|
|
|
|
class ReminderPubSub:
|
|
"""提醒消息的发布/订阅管理器"""
|
|
|
|
CHANNEL_PREFIX = "reminder"
|
|
|
|
def __init__(self):
|
|
self._redis: Optional[redis.Redis] = None
|
|
|
|
async def get_redis(self) -> redis.Redis:
|
|
"""获取 Redis 连接（懒加载）"""
|
|
if self._redis is None:
|
|
try:
|
|
redis_url = _get_redis_url()
|
|
self._redis = redis.from_url(redis_url, decode_responses=True)
|
|
await self._redis.ping()
|
|
logger.info("Redis Pub/Sub 连接成功")
|
|
except Exception as e:
|
|
logger.error(f"Redis Pub/Sub 连接失败: {e}")
|
|
raise
|
|
return self._redis
|
|
|
|
def _get_channel(self, user_id: int) -> str:
|
|
"""获取用户专属的频道名称"""
|
|
return f"{self.CHANNEL_PREFIX}:{user_id}"
|
|
|
|
async def publish(self, user_id: int, reminder: dict) -> bool:
|
|
"""发布提醒消息到指定用户的频道
|
|
|
|
Args:
|
|
user_id: 目标用户 ID
|
|
reminder: 提醒数据字典
|
|
|
|
Returns:
|
|
bool: 是否发布成功
|
|
"""
|
|
try:
|
|
r = await self.get_redis()
|
|
channel = self._get_channel(user_id)
|
|
message = json.dumps(reminder, ensure_ascii=False)
|
|
await r.publish(channel, message)
|
|
logger.info(f"发布提醒到用户 {user_id}: {reminder.get('title', '')}")
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"发布提醒失败: {e}")
|
|
return False
|
|
|
|
async def subscribe(self, user_id: int) -> redis.client.PubSub:
|
|
"""订阅用户的提醒频道
|
|
|
|
Args:
|
|
user_id: 用户 ID
|
|
|
|
Returns:
|
|
PubSub 对象，可用于接收消息
|
|
"""
|
|
r = await self.get_redis()
|
|
pubsub = r.pubsub()
|
|
channel = self._get_channel(user_id)
|
|
await pubsub.subscribe(channel)
|
|
logger.info(f"用户 {user_id} 订阅频道 {channel}")
|
|
return pubsub
|
|
|
|
async def unsubscribe(self, pubsub: redis.client.PubSub, user_id: int) -> None:
|
|
"""取消订阅
|
|
|
|
Args:
|
|
pubsub: PubSub 对象
|
|
user_id: 用户 ID
|
|
"""
|
|
try:
|
|
channel = self._get_channel(user_id)
|
|
await pubsub.unsubscribe(channel)
|
|
await pubsub.close()
|
|
logger.info(f"用户 {user_id} 取消订阅频道 {channel}")
|
|
except Exception as e:
|
|
logger.error(f"取消订阅失败: {e}")
|
|
|
|
async def close(self) -> None:
|
|
"""关闭 Redis 连接"""
|
|
if self._redis:
|
|
await self._redis.close()
|
|
self._redis = None
|
|
|
|
|
|
# 全局单例
|
|
reminder_pubsub = ReminderPubSub()
|