在FastAPI中利用lifespan与依赖注入高效管理Redis连接池

from contextlib import asynccontextmanager
from fastapi import FastAPI

@asynccontextmanager
async def lifespan(app: FastAPI):
    # 应用程序启动时执行的逻辑，例如初始化数据库连接池
    print("Application startup: Initializing resources...")
    app.state.my_resource = "Initialized Resource" # 将资源存储在 app.state
    yield
    # 应用程序关闭时执行的逻辑，例如关闭数据库连接池
    print("Application shutdown: Cleaning up resources...")
    # 清理 app.state.my_resource

from pydantic_settings import BaseSettings

class Settings(BaseSettings):
    REDIS_URL: str = "redis://localhost:6379/0" # 默认Redis连接URL

settings = Settings()

from contextlib import asynccontextmanager
from typing import AsyncGenerator, Annotated

import redis.asyncio as aredis
from fastapi import FastAPI, Depends, APIRouter, Request

from app.core.config import settings # 导入配置

# 定义 lifespan 上下文管理器
@asynccontextmanager
async def lifespan(app: FastAPI):
    """
    FastAPI 应用程序生命周期管理器。
    在应用程序启动时初始化 Redis 连接池，并在关闭时将其断开。
    """
    print("Initializing Redis Connection Pool...")
    # 创建 Redis 连接池并存储在 app.state 中
    # decode_responses=True 使得从 Redis 获取的数据自动解码为字符串
    app.state.redis_pool = aredis.ConnectionPool.from_url(
        settings.REDIS_URL, decode_responses=True
    )
    print("Redis Connection Pool Initialized.")
    yield # 应用程序在此处开始处理请求
    print("Closing Redis Connection Pool...")
    # 应用程序关闭时，断开 Redis 连接池
    await app.state.redis_pool.disconnect()
    print("Redis Connection Pool Closed.")

# 创建 FastAPI 应用实例，并传入 lifespan
app = FastAPI(lifespan=lifespan)
router = APIRouter()

# app/main.py (接着上面的代码)

async def get_redis_client(request: Request) -> AsyncGenerator[aredis.Redis, None]:
    """
    依赖注入函数，为每个请求提供一个 Redis 客户端实例。
    该客户端从应用程序全局连接池中获取连接。
    """
    # 从 app.state 中获取已初始化的 Redis 连接池
    pool = request.app.state.redis_pool
    # 从连接池创建一个 Redis 客户端实例
    # 注意：aredis.Redis 客户端是异步安全的，并且内部会管理从连接池获取和释放连接
    r = aredis.Redis(connection_pool=pool, decode_responses=True)
    try:
        yield r # 将 Redis 客户端提供给路由函数
    finally:
        # 在请求处理完毕后，关闭当前客户端实例，
        # 这会将底层连接释放回连接池。
        await r.aclose()

# 定义一个类型别名，方便在路由函数中使用
RedisDependency = Annotated[aredis.Redis, Depends(get_redis_client)]

# app/main.py (接着上面的代码)

@router.post("/items/")
async def create_item(item_data: dict, r: RedisDependency):
    """
    示例路由：创建一个新的 item 并将其存储到 Redis。
    """
    item_id = item_data.get("id") or "new_item"
    key = f"item:{item_id}"
    await r.set(key, str(item_data))
    value = await r.get(key)
    return {"message": "Item created and stored in Redis", "item_id": item_id, "stored_value": value}

@router.get("/items/{item_id}")
async def read_item(item_id: str, r: RedisDependency):
    """
    示例路由：从 Redis 中读取一个 item。
    """
    key = f"item:{item_id}"
    value = await r.get(key)
    if value:
        return {"item_id": item_id, "data": value}
    return {"message": "Item not found", "item_id": item_id}

app.include_router(router)