92 lines
3.1 KiB
Python
92 lines
3.1 KiB
Python
import json
|
||
import uuid
|
||
from contextlib import asynccontextmanager
|
||
from typing import Annotated, AsyncContextManager
|
||
|
||
import redis.asyncio as redis
|
||
from fastapi import Depends
|
||
|
||
from app.config.env import env
|
||
|
||
|
||
class RedisUtils():
|
||
def __init__(self):
|
||
self.redis_pool = None
|
||
|
||
# 检查 Redis 连接，先创建连接池管理对象
|
||
async def check_redis_connection(self):
|
||
self.redis_pool = redis.ConnectionPool(
|
||
host=env.redis_host,
|
||
port=env.redis_port,
|
||
password=env.redis_password,
|
||
db=env.redis_db,
|
||
encoding="utf-8",
|
||
decode_responses=True
|
||
)
|
||
# 立即对redis做一次读写测试
|
||
async with self.get_redis_connection() as redis_client:
|
||
try:
|
||
await redis_client.ping()
|
||
random_id = str(uuid.uuid4())
|
||
await redis_client.set("__init__setup__", random_id)
|
||
cache_random_id = await redis_client.get("__init__setup__")
|
||
if cache_random_id != random_id:
|
||
raise Exception("Redis 连接失败，读写测试不一致")
|
||
print("✅ Redis connection successful：", f"redis://{env.redis_password}@${env.redis_host}:{env.redis_port}/{env.redis_db}")
|
||
except Exception as e:
|
||
print(f"❌ Database connection failed")
|
||
raise e
|
||
|
||
# 获取一个Redis连接客户端
|
||
@asynccontextmanager
|
||
async def get_redis_connection(self) -> AsyncContextManager[redis.Redis]:
|
||
if not self.redis_pool:
|
||
raise Exception("Redis 连接池未初始化")
|
||
redis_client = redis.Redis(connection_pool=self.redis_pool)
|
||
try:
|
||
yield redis_client
|
||
finally:
|
||
await redis_client.close()
|
||
|
||
# 清理Redis连接池
|
||
async def close_redis_connection(self):
|
||
if self.redis_pool:
|
||
await self.redis_pool.disconnect()
|
||
|
||
|
||
redis_utils = RedisUtils()
|
||
|
||
|
||
# 作用：用于在接口中注入得到会话实例对象session，在接口执行完毕之后，自动执行close动作关闭会话
|
||
async def get_redis_client() -> redis.Redis:
|
||
async with redis_utils.get_redis_connection() as redis_client:
|
||
yield redis_client
|
||
|
||
|
||
RedisClientDep = Annotated[redis.Redis, Depends(get_redis_client)]
|
||
|
||
|
||
# 从redis中获取key的缓存，如果没有值则执行默认值获取函数，并保存到redis中
|
||
async def get_redis_cache(key: str, default_value_getter):
|
||
async with redis_utils.get_redis_connection() as redis_client:
|
||
# default_value_getter返回的字段可能嵌套多层对象，这里改成用json字符串缓存
|
||
json_string = await redis_client.get(key)
|
||
exists = False if json_string == "{}" else bool(json_string)
|
||
# print("exists", exists)
|
||
if exists:
|
||
return json.loads(json_string)
|
||
else:
|
||
value = await default_value_getter()
|
||
print("value", value)
|
||
if value is not None:
|
||
new_value = {k: v for k, v in value.items() if v is not None}
|
||
await redis_client.set(key, json.dumps(new_value, ensure_ascii=False))
|
||
return value
|
||
|
||
|
||
# 删除redis缓存
|
||
async def remove_redis_cache(key: str):
|
||
async with redis_utils.get_redis_connection() as redis_client:
|
||
result = await redis_client.delete(key)
|
||
return result > 0 # 返回是否删除成功
|