FastAPI启动事件中AsyncGenerator依赖注入的正确实践

import uvicorn
from fastapi import FastAPI, Depends
import redis.asyncio as redis
from redis.asyncio import Redis
from typing import AsyncGenerator
from rq import Queue # 假设rq是任务队列库

# 配置Redis连接
REDIS_HOST = "localhost"
REDIS_PORT = 6379

redis_pool = redis.ConnectionPool.from_url(f"redis://{REDIS_HOST}:{REDIS_PORT}")

async def get_async_redis_client() -> AsyncGenerator[Redis, None]:
    """
    异步生成器，用于提供Redis客户端连接。
    """
    async with Redis.from_pool(redis_pool) as client:
        yield client

def process_data(data: str):
    """
    模拟一个处理数据的函数。
    """
    print(f"Processing data: {data}")

def create_app():
    app = FastAPI(docs_url='/')
    task_queue: Queue = None # 声明为None，稍后初始化

    @app.on_event("startup")
    async def startup_event(redis_conn: redis.asyncio.Redis = Depends(get_async_redis_client)):
        """
        尝试在startup事件中使用Depends()注入Redis连接。
        """
        nonlocal task_queue
        task_queue = Queue("task_queue", connection=redis_conn)
        print("Redis connection initialized in startup event.")

    @app.post("/add_data")
    async def add_data(data: str):
        """
        添加数据到任务队列。
        """
        if task_queue:
            task_queue.enqueue(process_data, data)
            return {"message": "Book in processing"}
        return {"message": "Task queue not initialized", "status": "error"}

    @app.get("/get_data")
    async def get_data():
        """
        示例接口。
        """
        return {"data": "kek"}

    return app

def main():
    uvicorn.run(
        f"{__name__}:create_app",
        host='0.0.0.0', port=8888,
        reload=True
    )

if __name__ == '__main__':
    main()

import uvicorn
from fastapi import FastAPI
import redis.asyncio as redis
from redis.asyncio import Redis
from typing import AsyncGenerator
from rq import Queue # 假设rq是任务队列库
from contextlib import asynccontextmanager

# 配置Redis连接
REDIS_HOST = "localhost"
REDIS_PORT = 6379

redis_pool = redis.ConnectionPool.from_url(f"redis://{REDIS_HOST}:{REDIS_PORT}")

async def get_async_redis_client() -> AsyncGenerator[Redis, None]:
    """
    异步生成器，用于提供Redis客户端连接。
    """
    print("Opening Redis connection...")
    async with Redis.from_pool(redis_pool) as client:
        yield client
    print("Closing Redis connection...") # 应用关闭时执行

def process_data(data: str):
    """
    模拟一个处理数据的函数。
    """
    print(f"Processing data: {data}")

# 定义一个全局变量来存储任务队列
task_queue: Queue = None

@asynccontextmanager
async def lifespan(app: FastAPI):
    """
    FastAPI应用生命周期管理器。
    在应用启动时初始化资源，在应用关闭时清理资源。
    """
    global task_queue # 声明使用全局变量

    # 手动调用异步生成器以获取Redis连接
    # 注意：这里直接调用get_async_redis_client()，并迭代它
    # app.dependency_overrides.get(get_async_redis_client, get_async_redis_client)
    # 这一步是为了兼容可能存在的依赖覆盖，确保获取到的是最终的依赖函数
    redis_generator_func = app.dependency_overrides.get(get_async_redis_client, get_async_redis_client)
    async for redis_conn in redis_generator_func():
        # 在这里，redis_conn已经是实际的Redis客户端对象
        task_queue = Queue("task_queue", connection=redis_conn)
        print("Redis connection and Task Queue initialized via lifespan.")
        yield # 应用在此处启动并处理请求
    # 应用关闭时，生成器会继续执行，清理Redis连接
    print("Application shutdown: Resources released.")

def create_app():
    app = FastAPI(
        docs_url='/',
        lifespan=lifespan # 将lifespan上下文管理器传递给FastAPI
    )

    @app.post("/add_data")
    async def add_data(data: str):
        """
        添加数据到任务队列。
        """
        if task_queue:
            task_queue.enqueue(process_data, data)
            return {"message": "Book in processing"}
        return {"message": "Task queue not initialized", "status": "error"}

    @app.get("/get_data")
    async def get_data():
        """
        示例接口。
        """
        return {"data": "kek"}

    return app

def main():
    uvicorn.run(
        f"{__name__}:create_app",
        host='0.0.0.0', port=8888,
        reload=True
    )

if __name__ == '__main__':
    main()