FastAPI后台开发基础（17）： Middleware 的使用

from __future__ import annotations

import datetime
import json
import uuid

import uvicorn
from fastapi import FastAPI, HTTPException, Request, Response

app = FastAPI()


@app.get('/')
async def handler_func():
    """
    curl -X 'GET' \
    'http://127.0.0.1:18081/' \
    -H 'accept: application/json'

    {
       "message": "Hello World",
       "unique_id": "40261b47aa4648fb9b1bec6456cf08b9",
       "extra": "This is some extra data"
    }
    """
    print('handler begin')
    return {'message': 'Hello World', 'unique_id': uuid.uuid4().hex}


"""
使用 @app.middleware('http') 装饰器是一种添加中间件的方法，这种中间件会在每个 HTTP 请求和响应的处理过程中被调用
这里的 'http' 参数指定了中间件的类型
在 FastAPI（以及其底层的 Starlette 框架）中，'http' 是目前唯一支持的中间件类型，它用于处理 HTTP 层面的请求和响应
"""


@app.middleware('http')
async def add_process_time_header(request: Request, call_next):
    """
    中间件
    """
    print('add_process_time_header middleware begin')
    response: Response = await call_next(request)
    # Custom proprietary headers have historically been used with an X- prefix
    response.headers['X-Process-Time'] = datetime.datetime.now().isoformat()
    # print('add_process_time_header got response:', response.headers)
    print('add_process_time_header middleware end')
    return response


# 中间件二：添加自定义响应头
@app.middleware("http")
async def add_custom_header(request: Request, call_next):
    print('add_custom_header middleware begin')
    response: Response = await call_next(request)
    response.headers['X-Custom-Header'] = uuid.uuid4().hex
    # print('add_custom_header got response:', response.headers)
    print('add_custom_header middleware end')
    return response


@app.middleware("http")
async def modify_response_body(request: Request, call_next):
    print('modify_response_body middleware begin')
    response = await call_next(request)

    if request.url.path == '/docs':  # 过滤部分路径
        return response

    # 读取原始响应体
    original_body = b''
    async for chunk in response.body_iterator:
        original_body += chunk

    print('modify_response_body middleware, original response body size:', len(original_body))

    try:
        # 解码原始响应体（假设它是JSON）
        data = json.loads(original_body.decode())
    except json.JSONDecodeError:
        raise HTTPException(status_code = 500, detail = "Failed to decode response body")
    else:
        # 修改响应数据
        data['extra'] = "This is some extra data"

        # 编码修改后的响应体
        modified_body = json.dumps(data).encode()
        print('modify_response_body middleware, modified response body size:', len(modified_body))

        # 创建新响应对象
        new_response = Response(content = modified_body,
                                status_code = response.status_code,
                                headers = dict(response.headers),
                                media_type = "application/json")

        # 显式删除 Content-Length, 让框架在发送响应时自动计算并添加正确的 Content-Length 头部
        if 'content-length' in new_response.headers:
            del new_response.headers['content-length']

        # 复制除 Content-Length 外的其他头部
        # for key, value in response.headers.items():
        #     if key.lower() != 'content-length':
        #         new_response.headers[key] = value

        return new_response


if __name__ == '__main__':
    uvicorn.run(app, host = '127.0.0.1', port = 18081)