FastAPI 核心功能的完整示例代码,涵盖 WebSocket、后台任务、生命周期事件、中间件配置及会话管理

以下是一个整合 FastAPI 核心功能的完整示例代码,涵盖 WebSocket、后台任务、生命周期事件、中间件配置及会话管理。代码结构参考了多个技术文档的最佳实践:

from contextlib import asynccontextmanager
from fastapi import FastAPI, WebSocket, Request, Response, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from fastapi.middleware.gzip import GZipMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import StreamingResponse
import asyncio

# ===== 生命周期事件 =====
@asynccontextmanager
async def lifespan(app: FastAPI):
    # 启动时初始化资源(数据库连接等)
    print("⚡ Application starting...")
    yield
    # 关闭时释放资源
    print("🛑 Application shutting down...")

app = FastAPI(lifespan=lifespan)

# ===== 中间件配置 =====
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"]
)

app.add_middleware(GZipMiddleware, minimum_size=1024)  # 压缩大于1KB的响应

# ===== 静态文件服务 =====
app.mount("/static", StaticFiles(directory="static"), name="static")

# ===== WebSocket 实现 =====
active_connections = []

@app.websocket("/ws/chat")
async def websocket_chat(websocket: WebSocket):
    await websocket.accept()
    active_connections.append(websocket)
    try:
        while True:
            data = await websocket.receive_text()
            # 广播消息给所有客户端
            for conn in active_connections:
                await conn.send_text(f"User: {data}")
    except Exception:
        active_connections.remove(websocket)

# ===== 后台任务处理 =====
def cpu_intensive_task(message: str):
    """模拟CPU密集型任务"""
    import time
    time.sleep(5)
    return f"Processed: {message}"

async def async_io_task(message: str):
    """模拟异步IO任务"""
    await asyncio.sleep(3)
    return f"Async: {message}"

@app.post("/tasks")
async def create_task(
    bg_tasks: BackgroundTasks,
    message: str = "default"
):
    # 添加两种类型的后台任务
    bg_tasks.add_task(async_io_task, message)
    bg_tasks.add_task(cpu_intensive_task, message)
    return {"status": "Tasks started"}

# ===== 流式响应 =====
async def generate_logs():
    """实时日志生成器"""
    for i in range(10):
        yield f"data: Log entry {i}\n\n"
        await asyncio.sleep(1)

@app.get("/stream-logs")
async def stream_logs():
    return StreamingResponse(generate_logs(), media_type="text/event-stream")

# ===== 会话与Cookie管理 =====
@app.get("/set-cookie")
async def set_cookie(response: Response):
    response.set_cookie(
        key="user_token",
        value="encrypted_value_123",
        max_age=3600,
        httponly=True,
        samesite="Lax"
    )
    return {"message": "Cookie set"}

@app.get("/get-cookie")
async def get_cookie(request: Request):
    token = request.cookies.get("user_token")
    return {"token": token or "No cookie found"}

功能说明与最佳实践:

  1. WebSocket 实现
    • 使用内存列表管理活动连接
    • 支持文本消息的实时广播
    • 异常时自动清理失效连接
    • 前端可通过 new WebSocket('ws://localhost:8000/ws/chat') 连接

  2. 后台任务处理
    BackgroundTasks 支持同步/异步混合任务
    • CPU密集型任务自动使用线程池
    • 异步IO任务直接运行在事件循环中
    • 生产环境建议结合 Celery 处理分布式任务

  3. 生命周期管理
    • 替代已废弃的 startup/shutdown 事件
    • 支持异步资源初始化(数据库连接池等)
    • 确保资源释放的可靠性(try/finally 结构)

  4. 中间件配置
    • CORS 配置支持跨域请求
    • GZip 压缩优化响应速度
    • 静态文件服务通过 /static 路径提供

  5. 会话管理
    • 安全 Cookie 设置(httpOnly, SameSite)
    • 支持服务端主动设置 Cookie
    • 建议生产环境使用 JWT 或 Redis 会话存储

部署建议:

# 安装依赖
pip install uvicorn[standard] websockets

# 开发模式运行
uvicorn main:app --reload --ws-ping-interval 30

# 生产模式
uvicorn main:app --workers 4 --timeout-keep-alive 30

此实现综合了 FastAPI 的核心功能,可根据具体业务需求扩展 ORM 集成(推荐 SQLAlchemy 或 Tortoise-ORM)、身份验证系统(推荐 OAuth2 with JWT)和监控中间件(推荐 Prometheus)。

posted @ 2025-04-04 09:24  iTech  阅读(69)  评论(0)    收藏  举报