一文掌握异步web框架FastAPI（四）-- 中间件（添加回包数据、请求计时、用户代理、跨域资源共享CORS、请求限制、请求体大小限制、响应缓存、安全头、数据库事务管理）

from fastapi import FastAPI, Request
import logging
import json

logging.basicConfig(level=logging.INFO)

app = FastAPI()


@app.middleware("http")
async def log_body(request: Request, call_next):
    response = await call_next(request)

    # 捕获响应体内容
    body = b""
    async for chunk in response.body_iterator:
        body += chunk
    # 打印body
    logging.info(f"Response body: {body.decode('utf-8')}")

    # 打印body类型
    logging.info(f"body type: {type(body.decode('utf-8'))}")

    # 解析body为dict
    body_dict = json.loads(body.decode('utf-8'))

    # 新增数据
    body_dict.update({'name': 'jack'})

    # 将修改后的body重新编码为bytes
    body = json.dumps(body_dict).encode('utf-8')

    # 更新响应体长度
    response.headers["Content-Length"] = str(len(body))

    # 创建一个新的异步生成器来返回处理后的内容
    async def async_iterate():
        yield body

    response.body_iterator = async_iterate()
    return response


@app.get("/")
async def read_root():
    return {"Hello": "World"}

请求：

import requests

# 设置请求的URL
url = "http://127.0.0.1:8000/"

# 发送GET请求
response = requests.get(url)

# 输出响应内容
print("状态码:", response.status_code)  # 打印状态码
print('headers: ', response.headers)
print("响应体:", response.text)  # 打印响应的JSON内容

输出：

原接口返回的是{"Hello": "World"}，通过在中间件中新加了body数据，最后返回{"Hello": "World", "name": "jack"}，通过这个过程，可以更直观的看到中间件的运行过程和对中间件回包的处理。

7、请求计时中间件

这个中间件用于记录每个请求的处理时间，这对于性能分析很有帮助。

from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
import time

app = FastAPI()

@app.middleware("http")
async def add_process_time_header(request: Request, call_next):
    start_time = time.time()
    response = await call_next(request)
    process_time = time.time() - start_time
    response.headers["X-Process-Time"] = str(process_time)
    return response

@app.get("/")
async def root():
    time.sleep(1)  # 模拟耗时操作
    return {"message": "Hello World"}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)

请求：

import requests

# 设置请求的URL
url = "http://127.0.0.1:8000/"

# 发送GET请求
response = requests.get(url)

# 输出响应内容
print("状态码:", response.status_code)  # 打印状态码
print('headers: ', response.headers)
print("响应体:", response.text)  # 打印响应的JSON内容

8、用户代理中间件

这个中间件用于检查请求的 User-Agent 字段，确保请求来自合法的客户端。

from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse

app = FastAPI()

@app.middleware("http")
async def check_user_agent(request: Request, call_next):
    user_agent = request.headers.get("User-Agent")
    if user_agent and "Mozilla" in user_agent:
        return await call_next(request)
    return JSONResponse(status_code=403, content={"error": "Invalid User-Agent"})

@app.get("/headers/")
async def root():
    return {"message": "Hello World"}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)

不传User-Agent或传的User-Agent不是Mozilla，返回403

import requests

# 设置请求的URL
url = "http://127.0.0.1:8000/headers/"

# 发送GET请求
response = requests.get(url)

# 输出响应内容
print("状态码:", response.status_code)  # 打印状态码
print('headers: ', response.headers)
print("响应体:", response.json())  # 打印响应的JSON内容

传带Mozilla的User-Agent

import requests

# 设置请求的URL
url = "http://127.0.0.1:8000/headers/"

# 发送GET请求
response = requests.get(url, headers={"User-Agent": "Mozilla/5.0"})

# 输出响应内容
print("状态码:", response.status_code)  # 打印状态码
print('headers: ', response.headers)
print("响应体:", response.json())  # 打印响应的JSON内容