FastAPI 依赖注入之使用 yield 依赖：生成器依赖和资源管理

FastAPI 依赖注入之使用 yield 依赖：生成器依赖和资源管理

在 FastAPI 中，使用 yield 创建的依赖项能够管理资源的生命周期并处理异常。通过将 yield 与 try、except、finally 配合使用，依赖项可以在路由处理函数中提供必要的初始化，并在请求完成后清理资源。FastAPI 支持使用上下文管理器来简化这种依赖管理，确保依赖项的退出代码按正确顺序执行。这种方法非常适合用于数据库连接、文件处理、网络请求等需要清理资源的场景。本文介绍了如何通过 yield 和上下文管理器来创建高效且可管理的依赖项。

以下示例中使用的 Python 版本为 Python 3.10.15，FastAPI 版本为 0.115.4。

一 示例代码

from fastapi import Depends, FastAPI, HTTPException

app = FastAPI()

data = {
    "plumbus": {"description": "Freshly pickled plumbus", "owner": "Morty"},
    "portal-gun": {"description": "Gun to create portals", "owner": "Rick"},
}

class OwnerError(Exception):
    pass

def get_username():
    try:
        yield "Rick"
        print("哈哈")
    except OwnerError as e:
        raise HTTPException(status_code=400, detail=f"Owner error: {e}")

@app.get("/items01/{item_id}")
def get_item(item_id: str, username: str = Depends(get_username)):
    if item_id not in data:
        raise HTTPException(status_code=404, detail="Item not found")
    item = data[item_id]
    if item["owner"] != username:
        raise OwnerError(username)
    return item

FastAPI支持使用 yield 执行额外步骤，每个依赖只能用一次 yield ，运行代码文件 di06.py 来启动应用：

$ uvicorn di06:app --reload

在 SwaggerUI 中可以查看在线文档：http://127.0.0.1:8000/docs。
**注：**可以将使用 @contextlib.contextmanager 或 @contextlib.asynccontextmanager 的函数作为依赖项，FastAPI 内部也使用这两个装饰器。

二 使用 yield 的数据库依赖项

async def get_db():
    db = DBSession()
    try:
        yield db
    finally:
        db.close()

使用yield创建数据库会话并在完成后关闭，yield语句及其之前的代码在响应前执行，之后的代码在响应创建后发送前执行。无论是async函数还是普通函数，FastAPI都会正确处理依赖。

三 包含 yield 和 try 的依赖项

在包含yield的依赖中使用try代码块，可以捕获依赖或路由函数中的异常。可通过except捕获特定异常，并使用finally确保退出步骤执行。

四 使用 yield 的子依赖项

# 生成依赖A的异步函数
async def generate_dep_a():
    pass  # 这里可以进行一些资源的创建或初始化操作

# 生成依赖B的异步函数
async def generate_dep_b():
    pass  # 这里可以进行一些资源的创建或初始化操作

# 生成依赖C的异步函数
async def generate_dep_c():
    pass  # 这里可以进行一些资源的创建或初始化操作

# 依赖A的生成器函数
async def dependency_a():
    dep_a = generate_dep_a()  # 创建依赖A的实例
    try:
        yield dep_a  # 将依赖A传递给调用者（例如路由函数或其他依赖）
    finally:
        dep_a.close()  # 在退出时关闭依赖A

# 依赖B的生成器函数，依赖于依赖A
async def dependency_b(dep_a=Depends(dependency_a)):  
    dep_b = generate_dep_b()  # 创建依赖B的实例
    try:
        yield dep_b  # 将依赖B传递给调用者
    finally:
        dep_b.close(dep_a)  # 在退出时关闭依赖B，并确保依赖A的资源得到正确释放

# 依赖C的生成器函数，依赖于依赖B
async def dependency_c(dep_b=Depends(dependency_b)):  
    dep_c = generate_dep_c()  # 创建依赖C的实例
    try:
        yield dep_c  # 将依赖C传递给调用者
    finally:
        dep_c.close(dep_b)  # 在退出时关闭依赖C，并确保依赖B的资源得到正确释放

可以声明任意层级的树状依赖，任一依赖都可以使用yield。FastAPI会按正确顺序执行所有依赖，确保每个带yield的依赖的退出代码按预期运行。例如，dependency_c依赖于dependency_b，而dependency_b又依赖dependency_a。也可以混合使用带yield或return的依赖，FastAPI会确保它们按正确顺序执行。该过程由Python的上下文管理器实现。

五 包含 yield 和 HTTPException 的依赖项

def get_username():
    try:
        yield "Rick"
    except OwnerError as e:
        raise HTTPException(status_code=400, detail=f"Owner error: {e}")


@app.get("/items/{item_id}")
def get_item(item_id: str, username: Annotated[str, Depends(get_username)]):
    if item_id not in data:
        raise HTTPException(status_code=404, detail="Item not found")
    item = data[item_id]
    if item["owner"] != username:
        raise OwnerError(username)
    return item

可以在带有yield的依赖中使用try捕获异常，并在yield后的退出代码中抛出HTTPException等异常。虽然这是一种高级技巧，大多数情况下可以在路由函数中抛出异常，但也可以在依赖中处理，还可以创建自定义异常处理器捕获异常并抛出其他Exception。

六 包含 yield 和 except 的依赖项 要 raise

class InternalError(Exception):
    pass


def get_username01():
    try:
        yield "Rick"
    except InternalError:
        print("Oops, we didn't raise again, Britney ")
        raise

如果在包含yield的依赖中使用except捕获异常但不重新抛出，FastAPI不会察觉异常，客户端会收到HTTP 500 Internal Server Error，且服务器日志仅记录FastAPIError，没有具体错误信息。若使用raise重新抛出异常，服务器日志将记录自定义的InternalError。因此，在yield依赖中捕获异常后，应重新抛出，除非抛出如HTTPException之类的异常。

七 包含 yield, HTTPException, except 的依赖项和后台任务

后台任务应独立处理其逻辑并使用自己的资源（如数据库连接），使代码更加简洁。避免在后台任务中使用属于yield依赖项的资源。例如，后台任务应创建独立的数据库会话查询数据，传递对象ID作为参数，重新获取该对象而不是直接传递数据库对象。

八 在使用 yield 的依赖项中使用上下文管理器

class MySuperContextManager:
    def __init__(self):
        self.db = DBSession()

    def __enter__(self):
        return self.db

    def __exit__(self, exc_type, exc_value, traceback):
        self.db.close()


async def get_db():
    with MySuperContextManager() as db:
        yield db

在Python中，通过创建带有 __enter__() 和 __exit__() 方法的类定义上下文管理器，并在 FastAPI 的 yield 依赖项中使用 with 或 async with。FastAPI 会自动使用 @contextlib.contextmanager 或 @contextlib.asynccontextmanager 装饰器将单个 yield 函数转换为上下文管理器。

什么是上下文管理器

上下文管理器是可以在 with 语句中使用的Python对象。例如，使用 with 读取文件：

with open("./somefile.txt") as f:
    contents = f.read()
    print(contents)

这里，open("./somefile.txt") 创建了一个上下文管理器对象。当 with 代码块结束时，它会确保关闭文件，即使发生异常。

在FastAPI中，使用 yield 创建依赖项时，FastAPI 会将其转换为上下文管理器，并与其他相关工具结合使用。

九 完整代码示例

from fastapi import Depends, FastAPI, HTTPException

app = FastAPI()

data = {
    "plumbus": {"description": "Freshly pickled plumbus", "owner": "Morty"},
    "portal-gun": {"description": "Gun to create portals", "owner": "Rick"},
}


class OwnerError(Exception):
    pass


class InternalError(Exception):
    pass


def get_username():
    try:
        yield "Rick"
        print("哈哈")
    except OwnerError as e:
        raise HTTPException(status_code=400, detail=f"Owner error: {e}")


def get_username01():
    try:
        yield "Rick"
    except InternalError:
        print("Oops, we didn't raise again, Britney ")
        raise


@app.get("/items01/{item_id}")
def get_item(item_id: str, username: str = Depends(get_username)):
    if item_id not in data:
        raise HTTPException(status_code=404, detail="Item not found")
    item = data[item_id]
    if item["owner"] != username:
        raise OwnerError(username)
    return item


@app.get("/items02/{item_id}")
def get_item(item_id: str, username: str = Depends(get_username01)):
    if item_id == "portal-gun":
        raise InternalError(
            f"The portal gun is too dangerous to be owned by {username}"
        )
    if item_id != "plumbus":
        raise HTTPException(
            status_code=404, detail="Item not found, there's only a plumbus here"
        )
    return item_id


async def DBSession():
    return None


async def get_db():
    db = DBSession()
    try:
        yield db
    finally:
        db.close()


async def generate_dep_a():
    pass


async def generate_dep_b():
    pass


async def generate_dep_c():
    pass


async def dependency_a():
    dep_a = generate_dep_a()
    try:
        yield dep_a
    finally:
        dep_a.close()


async def dependency_b(dep_a=Depends(dependency_a)):
    dep_b = generate_dep_b()
    try:
        yield dep_b
    finally:
        dep_b.close(dep_a)


async def dependency_c(dep_b=Depends(dependency_b)):
    dep_c = generate_dep_c()
    try:
        yield dep_c
    finally:
        dep_c.close(dep_b)


class MySuperContextManager:
    def __init__(self):
        self.db = DBSession()

    def __enter__(self):
        return self.db

    def __exit__(self, exc_type, exc_value, traceback):
        self.db.close()


async def get_db():
    with MySuperContextManager() as db:
        yield db

十 源码地址

详情见：GitHub FastApiProj

十一 参考

[1] FastAPI 文档