import asyncio import logging import uuid import json import redis try: from Config.Config import ( REDIS_HOST, REDIS_PORT, REDIS_DB, REDIS_PASSWORD, REDIS_DECODE_RESPONSES, REDIS_MAX_CONNECTIONS ) except ModuleNotFoundError: import importlib.util import os _root_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) _config_path = os.path.join(_root_dir, "Config", "Config.py") _spec = importlib.util.spec_from_file_location("project_config_fallback", _config_path) _cfg = importlib.util.module_from_spec(_spec) assert _spec.loader is not None _spec.loader.exec_module(_cfg) REDIS_HOST = _cfg.REDIS_HOST REDIS_PORT = _cfg.REDIS_PORT REDIS_DB = _cfg.REDIS_DB REDIS_PASSWORD = getattr(_cfg, "REDIS_PASSWORD", None) REDIS_DECODE_RESPONSES = _cfg.REDIS_DECODE_RESPONSES REDIS_MAX_CONNECTIONS = _cfg.REDIS_MAX_CONNECTIONS # 创建logger实例 logger = logging.getLogger(__name__) class RedisKit: """ 异步Redis工具类，提供所有Redis操作的异步实现 单例模式确保全局只有一个Redis连接池 """ _instance = None _lock = asyncio.Lock() _redis_pool = None def __new__(cls): if cls._instance is None: cls._instance = super().__new__(cls) return cls._instance @classmethod async def _ensure_pool(cls): """ 确保Redis连接池已创建 """ if cls._redis_pool is None: async with cls._lock: if cls._redis_pool is None: try: # 创建同步Redis连接池 sync_pool = redis.ConnectionPool( host=REDIS_HOST, port=REDIS_PORT, db=REDIS_DB, password=REDIS_PASSWORD, decode_responses=REDIS_DECODE_RESPONSES, max_connections=REDIS_MAX_CONNECTIONS ) # 创建Redis连接实例 cls._redis_pool = redis.Redis( connection_pool=sync_pool, encoding='utf-8', decode_responses=REDIS_DECODE_RESPONSES ) # 测试连接 await asyncio.to_thread(cls._redis_pool.ping) logger.info("Redis连接池创建成功") except Exception as e: logger.error(f"Redis连接池创建失败: {e}") raise @classmethod async def get_connection(cls): """ 获取Redis连接实例 Returns: redis.Redis: Redis连接实例 """ await cls._ensure_pool() return cls._redis_pool async def get_data(self, key): """ 异步从Redis中获取数据 Args: key (str): Redis键名 Returns: any: Redis中存储的值，如果键不存在则返回None """ try: await self._ensure_pool() return await asyncio.to_thread(RedisKit._redis_pool.get, key) except Exception as e: logger.error(f"从Redis获取数据失败(key={key}): {e}") return None async def exists(self, key): """ 异步检查Redis键是否存在 Args: key (str): Redis键名 Returns: bool: 键存在返回True，不存在或失败返回False """ try: await self._ensure_pool() result = await asyncio.to_thread(RedisKit._redis_pool.exists, key) return bool(result) except Exception as e: logger.error(f"检查Redis键是否存在失败(key={key}): {e}") return False async def set_data(self, key, value, expire=None): """ 异步向Redis中写入数据 Args: key (str): Redis键名 value (str): 要存储的值 expire (int, optional): 过期时间(秒)，默认不过期 Returns: bool: 写入成功返回True，失败返回False """ try: await self._ensure_pool() if expire: result = await asyncio.to_thread(RedisKit._redis_pool.setex, key, expire, value) else: result = await asyncio.to_thread(RedisKit._redis_pool.set, key, value) return result except Exception as e: logger.error(f"向Redis写入数据失败(key={key}): {e}") return False async def delete_data(self, key): """ 异步从Redis中删除数据 Args: key (str): Redis键名 Returns: bool: 删除成功返回True，失败返回False """ try: await self._ensure_pool() result = await asyncio.to_thread(RedisKit._redis_pool.delete, key) return result > 0 # delete返回删除的键数量，大于0表示删除成功 except Exception as e: logger.error(f"从Redis删除数据失败(key={key}): {e}") return False async def keys(self, pattern='*'): """ 异步获取匹配模式的所有键 Args: pattern (str): 匹配模式 Returns: list: 匹配的键列表 """ try: await self._ensure_pool() return await asyncio.to_thread(RedisKit._redis_pool.keys, pattern) except Exception as e: logger.error(f"获取Redis键列表失败(pattern={pattern}): {e}") return [] async def delete(self, *names): """ 异步删除一个或多个键 Args: *names: 要删除的键名 Returns: int: 删除的键数量 """ try: await self._ensure_pool() if not names: return 0 return await asyncio.to_thread(RedisKit._redis_pool.delete, *names) except Exception as e: logger.error(f"删除Redis键失败: {e}") return 0 async def publish(self, channel, message): """ 异步发布消息到Redis频道 Args: channel (str): Redis频道名称 message (str): 要发布的消息 Returns: int: 接收消息的订阅者数量 """ try: await self._ensure_pool() result = await asyncio.to_thread(RedisKit._redis_pool.publish, channel, message) return result except Exception as e: logger.error(f"Redis发布消息失败(channel={channel}): {e}") return 0 async def subscribe(self, channel): """ 异步订阅Redis频道 Args: channel (str): Redis频道名称 Returns: redis.client.PubSub: PubSub对象，用于接收消息 """ try: await self._ensure_pool() pubsub = RedisKit._redis_pool.pubsub() await asyncio.to_thread(pubsub.subscribe, channel) return pubsub except Exception as e: logger.error(f"Redis订阅频道失败(channel={channel}): {e}") return None async def save_user_token(self, person_id, platform, token, expire=86400): """ 异步保存用户的token到Redis Args: person_id (str): 用户ID platform (str): 平台 token (str): 生成的token expire (int): 过期时间（秒），默认24小时 Returns: bool: 保存成功返回True，失败返回False """ try: # 使用固定格式的key：user_token:{person_id}:{platform} key = f"user_token:{person_id}:{platform}" return await self.set_data(key, token, expire) except Exception as e: logger.error(f"保存用户token失败(person_id={person_id}, platform={platform}): {e}") return False async def get_user_token(self, person_id, platform): """ 异步获取用户在Redis中存储的token Args: person_id (str): 用户ID platform (str): 平台 Returns: str: 用户的token，如果不存在或过期则返回None """ try: key = f"user_token:{person_id}:{platform}" return await self.get_data(key) except Exception as e: logger.error(f"获取用户token失败(person_id={person_id}, platform={platform}): {e}") return None async def validate_user_token(self, person_id, platform, token): """ 异步验证用户的token是否与Redis中存储的一致 Args: person_id (str): 用户ID platform (str): 平台、终端 token (str): 待验证的token Returns: bool: token有效返回True，无效返回False """ try: stored_token = await self.get_user_token(person_id, platform) # 如果没有存储的token，或者存储的token与当前token不一致，则token无效 return stored_token is not None and stored_token == token except Exception as e: logger.error(f"验证用户token失败(person_id={person_id}, platform={platform}): {e}") return False async def delete_user_token(self, person_id, platform=None): """ 异步删除用户在Redis中存储的token Args: person_id (str): 用户ID platform (str, optional): 平台，如果为None则删除所有平台的token Returns: bool: 删除成功返回True，失败返回False """ try: await self._ensure_pool() if platform: # 删除特定平台的token key = f"user_token:{person_id}:{platform}" return await self.delete_data(key) else: # 删除所有平台的token（使用模式匹配） pattern = f"user_token:{person_id}:*" cursor = b'0' deleted_count = 0 # 使用scan命令查找所有匹配的键 while cursor: cursor, keys = await asyncio.to_thread( RedisKit._redis_pool.scan, cursor=cursor, match=pattern, count=100 ) if keys: deleted = await asyncio.to_thread( RedisKit._redis_pool.delete, *keys ) deleted_count += deleted return deleted_count > 0 except Exception as e: logger.error(f"删除用户token失败(person_id={person_id}): {e}") return False async def list_push(self, key, value, expire=None): """ 异步向Redis List中推入数据（LPUSH） Args: key (str): Redis键名 value (str): 要推入的值 expire (int, optional): 过期时间(秒)，如果设置则同时为key设置过期时间 Returns: int: 推入后list的长度，失败返回-1 """ try: await self._ensure_pool() result = await asyncio.to_thread(RedisKit._redis_pool.lpush, key, value) # 如果设置了过期时间，为key设置TTL if expire and result > 0: await asyncio.to_thread(RedisKit._redis_pool.expire, key, expire) return result except Exception as e: logger.error(f"向Redis List推入数据失败(key={key}): {e}") return -1 async def list_pop(self, key): """ 异步从Redis List中弹出数据（RPOP） Args: key (str): Redis键名 Returns: str: 弹出的值，如果list为空或失败返回None """ try: await self._ensure_pool() result = await asyncio.to_thread(RedisKit._redis_pool.rpop, key) return result except Exception as e: logger.error(f"从Redis List弹出数据失败(key={key}): {e}") return None async def list_length(self, key): """ 异步获取Redis List的长度 Args: key (str): Redis键名 Returns: int: list的长度，如果失败返回-1 """ try: await self._ensure_pool() result = await asyncio.to_thread(RedisKit._redis_pool.llen, key) return result except Exception as e: logger.error(f"获取Redis List长度失败(key={key}): {e}") return -1 async def list_get_all(self, key): """ 异步获取Redis List中的所有数据（LRANGE 0 -1） Args: key (str): Redis键名 Returns: list: list中的所有数据，如果失败返回空列表 """ try: await self._ensure_pool() result = await asyncio.to_thread(RedisKit._redis_pool.lrange, key, 0, -1) return result if result else [] except Exception as e: logger.error(f"获取Redis List所有数据失败(key={key}): {e}") return [] async def list_clear(self, key): """ 异步清空Redis List（删除key） Args: key (str): Redis键名 Returns: bool: 清空成功返回True，失败返回False """ try: return await self.delete_data(key) except Exception as e: logger.error(f"清空Redis List失败(key={key}): {e}") return False async def set_is_member(self, key, value): """ 异步检查值是否在Redis集合中 Args: key (str): Redis键名 value: 要检查的值 Returns: bool: 值在集合中返回True，不在或失败返回False """ try: await self._ensure_pool() result = await asyncio.to_thread(RedisKit._redis_pool.sismember, key, value) return bool(result) except Exception as e: logger.error(f"检查值是否在Redis集合中失败(key={key}): {e}") return False async def set_add(self, key, value): """ 异步向Redis集合中添加值 Args: key (str): Redis键名 value: 要添加的值 Returns: bool: 添加成功返回True，失败返回False """ try: await self._ensure_pool() result = await asyncio.to_thread(RedisKit._redis_pool.sadd, key, value) return result > 0 # sadd返回添加的元素数量，大于0表示添加成功 except Exception as e: logger.error(f"向Redis集合添加值失败(key={key}): {e}") return False async def set_remove(self, key, value): """ 异步从Redis集合中移除值 Args: key (str): Redis键名 value: 要移除的值 Returns: bool: 移除成功返回True，失败返回False """ try: await self._ensure_pool() result = await asyncio.to_thread(RedisKit._redis_pool.srem, key, value) return result > 0 # srem返回移除的元素数量，大于0表示移除成功 except Exception as e: logger.error(f"从Redis集合移除值失败(key={key}): {e}") return False async def set_cardinality(self, key): """ 异步获取Redis集合的成员数量 Args: key (str): Redis键名 Returns: int: 集合的成员数量，如果失败返回0 """ try: await self._ensure_pool() result = await asyncio.to_thread(RedisKit._redis_pool.scard, key) return result except Exception as e: logger.error(f"获取Redis集合成员数量失败(key={key}): {e}") return 0 # 分布式锁相关方法 async def acquire_lock(self, lock_name, timeout=300): """ 获取分布式锁 Args: lock_name (str): 锁的名称 timeout (int): 锁的超时时间(秒)，默认300秒 Returns: str: 锁ID，如果获取失败返回None """ try: await self._ensure_pool() # 生成唯一锁ID lock_id = str(uuid.uuid4()) key = f"lock:{lock_name}" # 使用Redis原子操作设置锁 result = await asyncio.to_thread( RedisKit._redis_pool.setnx, key, lock_id ) if result: # 设置锁的过期时间 await asyncio.to_thread( RedisKit._redis_pool.expire, key, timeout ) logger.info(f"获取分布式锁成功: {lock_name}, lock_id: {lock_id}") return lock_id else: logger.info(f"获取分布式锁失败: {lock_name} (锁已存在)") return None except Exception as e: logger.error(f"获取分布式锁失败: {lock_name}, 错误: {e}") return None async def release_lock(self, lock_name, lock_id): """ 释放分布式锁 Args: lock_name (str): 锁的名称 lock_id (str): 锁ID Returns: bool: 是否成功释放锁 """ try: await self._ensure_pool() key = f"lock:{lock_name}" # 使用Lua脚本原子性检查并释放锁 lua_script = """ if redis.call('get', KEYS[1]) == ARGV[1] then return redis.call('del', KEYS[1]) else return 0 end """ result = await asyncio.to_thread( RedisKit._redis_pool.eval, lua_script, 1, key, lock_id ) if result == 1: logger.info(f"释放分布式锁成功: {lock_name}, lock_id: {lock_id}") return True else: logger.info(f"释放分布式锁失败: {lock_name}, lock_id: {lock_id} (锁不存在或已被其他进程持有)") return False except Exception as e: logger.error(f"释放分布式锁失败: {lock_name}, 错误: {e}") return False async def renew_lock(self, lock_name, lock_id, timeout=300): """ 续约分布式锁 Args: lock_name (str): 锁的名称 lock_id (str): 锁ID timeout (int): 续约时间(秒)，默认300秒 Returns: bool: 是否成功续约 """ try: await self._ensure_pool() key = f"lock:{lock_name}" # 使用Lua脚本原子性检查并续约锁 lua_script = """ if redis.call('get', KEYS[1]) == ARGV[1] then return redis.call('expire', KEYS[1], ARGV[2]) else return 0 end """ result = await asyncio.to_thread( RedisKit._redis_pool.eval, lua_script, 1, key, lock_id, timeout ) if result == 1: logger.info(f"续约分布式锁成功: {lock_name}, lock_id: {lock_id}") return True else: logger.info(f"续约分布式锁失败: {lock_name}, lock_id: {lock_id} (锁不存在或已被其他进程持有)") return False except Exception as e: logger.error(f"续约分布式锁失败: {lock_name}, 错误: {e}") return False # 任务进度监控相关方法 async def set_task_progress(self, task_id, progress, status="processing", message=""): """ 设置任务进度 Args: task_id (str): 任务ID progress (int): 进度百分比(0-100) status (str): 任务状态(waiting, processing, completed, failed) message (str): 任务消息 Returns: bool: 设置成功返回True，失败返回False """ try: await self._ensure_pool() key = f"task:progress:{task_id}" # 构造进度数据 progress_data = { "task_id": task_id, "progress": min(100, max(0, progress)), # 确保进度在0-100之间 "status": status, "message": message, "update_time": asyncio.get_event_loop().time() } # 序列化进度数据 progress_json = json.dumps(progress_data, ensure_ascii=False) # 保存到Redis await asyncio.to_thread( RedisKit._redis_pool.set, key, progress_json ) # 设置过期时间为24小时 await asyncio.to_thread( RedisKit._redis_pool.expire, key, 86400 ) logger.info(f"设置任务进度成功: {task_id}, 进度: {progress}%, 状态: {status}") return True except Exception as e: logger.error(f"设置任务进度失败: {task_id}, 错误: {e}") return False async def get_task_progress(self, task_id): """ 获取任务进度 Args: task_id (str): 任务ID Returns: dict: 任务进度数据，如果不存在返回None """ try: await self._ensure_pool() key = f"task:progress:{task_id}" # 从Redis获取进度数据 progress_json = await asyncio.to_thread( RedisKit._redis_pool.get, key ) if progress_json: return json.loads(progress_json) else: return None except Exception as e: logger.error(f"获取任务进度失败: {task_id}, 错误: {e}") return None async def complete_task(self, task_id, success=True, message=""): """ 完成任务 Args: task_id (str): 任务ID success (bool): 任务是否成功完成 message (str): 完成消息 Returns: bool: 设置成功返回True，失败返回False """ status = "completed" if success else "failed" progress = 100 if success else 0 return await self.set_task_progress(task_id, progress, status, message) async def delete_task_progress(self, task_id): """ 删除任务进度记录 Args: task_id (str): 任务ID Returns: bool: 删除成功返回True，失败返回False """ try: await self._ensure_pool() key = f"task:progress:{task_id}" # 从Redis删除进度记录 result = await asyncio.to_thread( RedisKit._redis_pool.delete, key ) logger.info(f"删除任务进度记录成功: {task_id}") return result > 0 except Exception as e: logger.error(f"删除任务进度记录失败: {task_id}, 错误: {e}") return False @classmethod async def close(cls): """ 关闭Redis连接池 """ if cls._redis_pool: # 同步Redis连接池不需要等待关闭 cls._redis_pool.connection_pool.disconnect() cls._redis_pool = None logger.info("Redis连接池已关闭") # 创建全局实例 redisKit = RedisKit() async def main(): """ 测试异步Redis操作的main函数 """ print("开始测试异步Redis工具类...") try: # 测试基本数据操作 print("\n=== 测试基本数据操作 ===") # 测试设置和获取数据 test_key = "test_key" test_value = "test_value_123" result = await redisKit.set_data(test_key, test_value, expire=60) print(f"设置数据结果: {result}") get_value = await redisKit.get_data(test_key) print(f"获取数据结果: {get_value}") # 测试用户token操作 print("\n=== 测试用户token操作 ===") # 保存用户token person_id = "user123" platform = "web" token = "jwt_token_xyz789" save_result = await redisKit.save_user_token(person_id, platform, token, expire=120) print(f"保存token结果: {save_result}") # 获取用户token stored_token = await redisKit.get_user_token(person_id, platform) print(f"获取token结果: {stored_token}") # 验证token validate_result = await redisKit.validate_user_token(person_id, platform, token) print(f"验证token结果: {validate_result}") # 验证错误的token wrong_token = "wrong_token" wrong_validate_result = await redisKit.validate_user_token(person_id, platform, wrong_token) print(f"验证错误token结果: {wrong_validate_result}") # 测试删除操作 print("\n=== 测试删除操作 ===") # 删除测试数据 delete_result = await redisKit.delete_data(test_key) print(f"删除数据结果: {delete_result}") # 删除用户token delete_token_result = await redisKit.delete_user_token(person_id, platform) print(f"删除token结果: {delete_token_result}") # 测试批量token删除功能 print("\n=== 测试批量token删除功能 ===") # 保存多个平台的token platforms = ["android", "ios", "web"] for i, platform in enumerate(platforms): await redisKit.save_user_token(f"user{i + 1}", platform, f"token{i + 1}", expire=300) print(f"保存用户{i+1}的{platform}平台token") # 删除特定用户的所有token user_to_delete = "user2" batch_delete_result = await redisKit.delete_user_token(user_to_delete) print(f"批量删除{user_to_delete}的所有token结果: {batch_delete_result}") # 验证删除结果 for platform in platforms: token_check = await redisKit.get_user_token(user_to_delete, platform) print(f"删除后{user_to_delete}-{platform}平台的token状态: {'已删除' if token_check is None else '存在'}") print("\n=== 所有测试完成 ===") except Exception as e: print(f"测试过程中发生错误: {e}") import traceback traceback.print_exc() finally: # 清理测试数据并关闭连接池 print("\n清理测试数据并关闭连接池...") # 清理测试用的数据 test_keys_to_cleanup = ["test_key"] for key in test_keys_to_cleanup: await redisKit.delete_data(key) # 清理测试用户token cleanup_users = ["user123", "user1", "user2", "user3"] for user in cleanup_users: try: await redisKit.delete_user_token(user) except: pass # 忽略清理过程中的错误 # 关闭Redis连接池 await redisKit.close() print("Redis连接池已关闭") if __name__ == "__main__": asyncio.run(main())