【FastAPI】中间件

大神神奇奇奇

已于 2024-12-25 19:57:16 修改

阅读量1.1k

点赞数 14

文章标签： fastapi 中间件 python

于 2024-12-25 19:29:10 首次发布

本文链接：https://blog.youkuaiyun.com/xiaodoudouer/article/details/144725836

版权

【FastAPI】中间件

一、概述
二、作用
三、总结
四、相关链接

一、概述

FastAPI的中间件提供了一种强大的机制，允许开发者在请求被处理之前以及响应返回给客户端之前对它们进行拦截和操作。
这为实现诸如日志记录、身份验证、性能监控、错误处理等功能提供了极大的灵活性。

二、作用

2.1 日志记录与监控

中间件可以用来记录所有进入系统的HTTP请求及其对应的响应。
例如，你可以创建一个中间件来测量每个请求的处理时间，并将这个信息作为自定义头部X-Process-Time添加到响应中。
这对于性能分析非常有用，因为它可以帮助你了解哪些请求消耗了较多的时间，从而优化系统性能。

import json
import logging
import os
import time
from fastapi import FastAPI, Request, Response
from starlette.middleware.base import BaseHTTPMiddleware

# 设置日志格式和级别
LOG_LEVEL = os.getenv('LOG_LEVEL', 'INFO').upper()
logging.basicConfig(
    level=LOG_LEVEL,
    format='%(asctime)s - %(levelname)s - %(message)s'
)

class LoggingMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next) -> Response:
        start_time = time.perf_counter()
        response = await call_next(request)
        process_time = time.perf_counter() - start_time
        
        log_data = {
            "client_ip": request.client.host,
            "method": request.method,
            "path": request.url.path,
            "status_code": response.status_code,
            "process_time": f"{process_time:.4f}s",
            "timestamp": datetime.datetime.fromtimestamp(time.time()).strftime('%Y-%m-%d %H:%M:%S')
        }
        
        # 根据环境变量决定是否打印详细的日志内容
        if LOG_LEVEL == 'DEBUG':
            log_data["headers"] = dict(request.headers)
            log_data["query_params"] = dict(request.query_params)
        
        logging.info(json.dumps(log_data))
        
        return response

app = FastAPI()
app.add_middleware(LoggingMiddleware)

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

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

测试方法：

启动应用：python main.py
使用浏览器或Postman访问http://localhost:8000/
查看命令行输出的日志信息，确认日志已正确记录。

2024-12-25 19:21:56,568 - INFO - {"client_ip": "127.0.0.1", "method": "GET", "path": "/", "status_code": 200, "process_time": "0.0003s", "timestamp": "2024-12-25 19:21:56"}
INFO:     127.0.0.1:53956 - "GET / HTTP/1.1" 200 OK

设置日志配置
这里设置了日志记录器的基本配置。通过环境变量 LOG_LEVEL 来控制日志级别，默认为 INFO。
这意味着除非设置了更低的日志级别（如 DEBUG），否则不会输出调试信息。
此外，还指定了日志格式，包括时间戳、日志级别和消息内容。

LOG_LEVEL = os.getenv('LOG_LEVEL', 'INFO').upper()
logging.basicConfig(
    level=LOG_LEVEL,
    format='%(asctime)s - %(levelname)s - %(message)s'
)

自定义中间件 LoggingMiddleware
中间件是用来处理请求和响应之间逻辑的理想场所。
在这个例子中，LoggingMiddleware 被用来捕获每次请求的时间点，然后在请求完成后计算处理所需的时间，并将相关信息以 JSON 格式记录下来。这有助于监控 API 性能并诊断问题。

开始计时：使用 time.perf_counter() 获取当前时间。
调用下一个处理器：通过 await call_next(request) 将请求传递给下一个中间件或最终的路由处理器。
结束计时：再次调用 time.perf_counter() 计算总耗时。
构建日志数据：创建一个字典 log_data 包含客户端 IP、HTTP 方法、路径、状态码及处理时间等信息。
条件性添加更多信息：如果日志级别为 DEBUG，还会额外添加请求头和查询参数到日志中，以便于更详细的调试。
记录日志：最后，将构造好的日志数据转换成 JSON 字符串并通过 logging.info() 输出。

定义 FastAPI 应用程序
这里创建了一个 FastAPI 实例 app 并注册了之前定义的 LoggingMiddleware。随后定义了一个 GET 方法对应的路由 /，当访问根路径时会返回一个 JSON 响应，其中包含了一条简单的欢迎信息 “Hello World”。

app = FastAPI()
app.add_middleware(LoggingMiddleware)

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

启动 Uvicorn 服务器
这段代码确保只有当脚本作为主模块运行时才会启动 Uvicorn 服务器。
Uvicorn 是一个用于运行 ASGI 应用的异步服务器，它将托管 FastAPI 应用程序，并监听指定的主机地址 (0.0.0.0) 和端口 (8000)。
同时，也将日志级别传递给 Uvicorn，以保证一致性。

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

2.2 身份验证与授权

通过中间件，可以在请求到达路由处理器之前检查用户的身份验证状态。

如果请求未通过身份验证，则可以直接返回401 Unauthorized响应；
如果通过了身份验证但没有足够的权限访问特定资源，则可以返回403 Forbidden响应。

这种方式确保了只有经过适当授权的用户才能访问受保护的API端点。

from datetime import datetime, timedelta
from fastapi import Depends, FastAPI, HTTPException, status
from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm
from jose import JWTError, jwt
from pydantic import BaseModel
from typing import Optional
from fastapi.middleware.cors import CORSMiddleware
from fastapi.middleware.base import BaseHTTPMiddleware

SECRET_KEY = "your_secret_key"
ALGORITHM = "HS256"
ACCESS_TOKEN_EXPIRE_MINUTES = 30

oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")

class Token(BaseModel):
    access_token: str
    token_type: str

class TokenData(BaseModel):
    username: Optional[str] = None

fake_users_db = {
    "johndoe": {
        "username": "johndoe",
        "full_name": "John Doe",
        "email": "johndoe@example.com",
        "hashed_password": "fakehashedsecret",
        "disabled": False,
    }
}

def verify_password(plain_password, hashed_password):
    return plain_password == hashed_password

def get_user(db, username: str):
    if username in db:
        user_dict = db[username]
        return UserInDB(**user_dict)

def authenticate_user(fake_db, username: str, password: str):
    user = get_user(fake_db, username)
    if not user:
        return False
    if not verify_password(password, user.hashed_password):
        return False
    return user

def create_access_token(data: dict, expires_delta: Optional[timedelta] = None):
    to_encode = data.copy()
    if expires_delta:
        expire = datetime.utcnow() + expires_delta
    else:
        expire = datetime.utcnow() + timedelta(minutes=15)
    to_encode.update({"exp": expire})
    encoded_jwt = jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
    return encoded_jwt

async def get_current_user(token: str = Depends(oauth2_scheme)):
    credentials_exception = HTTPException(
        status_code=status.HTTP_401_UNAUTHORIZED,
        detail="Could not validate credentials",
        headers={"WWW-Authenticate": "Bearer"},
    )
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        username: str = payload.get("sub")
        exp: int = payload.get("exp")
        if username is None or exp is None:
            raise credentials_exception
        if datetime.utcfromtimestamp(exp) < datetime.utcnow():
            raise credentials_exception
        token_data = TokenData(username=username)
    except JWTError:
        raise credentials_exception
    user = get_user(fake_users_db, username=token_data.username)
    if user is None:
        raise credentials_exception
    return user

class AuthMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next) -> Response:
        token = request.headers.get('Authorization')
        if token:
            try:
                await get_current_user(token.split()[1])  # Assuming 'Bearer <token>'
            except HTTPException as e:
                return JSONResponse(status_code=e.status_code, content={"detail": e.detail})
        else:
            return JSONResponse(status_code=401, content={"detail": "Unauthorized"})
        return await call_next(request)

app = FastAPI()
app.add_middleware(AuthMiddleware)

@app.post("/token", response_model=Token)
async def login_for_access_token(form_data: OAuth2PasswordRequestForm = Depends()):
    user = authenticate_user(fake_users_db, form_data.username, form_data.password)
    if not user:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Incorrect username or password",
            headers={"WWW-Authenticate": "Bearer"},
        )
    access_token_expires = timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
    access_token = create_access_token(
        data={"sub": user.username}, expires_delta=access_token_expires
    )
    return {"access_token": access_token, "token_type": "bearer"}

@app.get("/users/me", response_model=User)
async def read_users_me(current_user: User = Depends(get_current_user)):
    return current_user

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

2.3 CORS（跨域资源共享）

为了使Web应用程序能够从不同的域名发出AJAX请求，你需要配置CORS策略。
FastAPI内置了CORSMiddleware，它允许你轻松地设置允许的来源、方法、头等参数，从而实现安全且灵活的跨域通信2。
下面是一个支持CORS配置的应用程序示例，允许来自特定源的请求，并处理预检请求（Preflight Requests）。

from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware

app = FastAPI()

origins = [
    "http://localhost:3000",
]

app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
    expose_headers=["X-Total-Count"],
    max_age=86400
)

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

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

2.4 Gzip压缩

对于需要提高传输效率的情况，可以使用Gzip中间件来压缩响应内容。这样不仅可以减少网络带宽占用，还能加快页面加载速度，特别是在移动设备上效果尤为明显2。

from fastapi import FastAPI
from fastapi.middleware.gzip import GZipMiddleware

app = FastAPI()

app.add_middleware(
    GZipMiddleware,
    minimum_size=1000,
    compresslevel=9
)

@app.get("/")
async def root():
    large_text = "a" * 1024  # 创建一个较大的字符串用于测试压缩效果
    return {"message": large_text}

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