FastAPI -- 第三弹(自定义响应、中间件、代理、WebSockets)

from fastapi import FastAPI, Response
from starlette.requests import Request
from starlette.responses import JSONResponse
from starlette.middleware.cors import CORSMiddleware
from starlette.exceptions import ExceptionMiddleware
 
app = FastAPI()
 
# 自定义响应类
class MyJSONResponse(JSONResponse):
    media_type = "application/vnd.mysite+json"
 
# 自定义错误处理
@app.exception_handler(Exception)
async def custom_exception_handler(request: Request, exc: Exception):
    return MyJSONResponse(
        content={"message": "An error occurred"},
        status_code=500,
    )
 
# 应用全局中间件
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)
 
# 自定义异常中间件
app.add_middleware(ExceptionMiddleware)
 
# 路由代理示例
@app.api_route("/proxy/{item_id}")
async def proxy_route(item_id: int):
    # 这里可以使用requests或者其他库来代理请求
    # 假设有一个后端服务运行在 http://backend-service:8000
    response = await some_async_request_function(f"http://backend-service:8000/items/{item_id}")
    return Response(media_type=response.media_type, content=response.content)
 
# WebSocket代理示例
@app.websocket_route("/ws/{item_id}")
async def websocket_route(websocket: WebSocket, item_id: int):
    await websocket.accept()
    # 可以在这里建立WebSocket连接并代理数据
    # 假设有一个WebSocket服务运行在 ws://websocket-service
    # async with some_websocket_client_library(f"ws://websocket-service/{item_id}") as websocket:
    #     await websocket.send_text("Hello, WebSocket!")
    #     receive_text = await websocket.receive_text()
    #     await websocket.send_text(f"Received: {receive_text}")

在这个示例中，我们演示了如何自定义响应类、处理异常并应用跨源资源共享(CORS)以及如何创建路由代理和WebSocket代理。这些技术可以应用于构建高度可扩展和灵活的后端服务。