hxf/backend/tests/fastapi_test/main.py

from contextvars import ContextVar
from fastapi import FastAPI, Request, Depends, HTTPException
from fastapi.security import OAuth2PasswordBearer
from pydantic import BaseModel
from typing import Optional
import uuid
import uvicorn

app = FastAPI()
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")

# 创建上下文变量存储当前用户和请求ID
current_user_ctx: ContextVar[dict] = ContextVar("current_user", default=None)
request_id_ctx: ContextVar[str] = ContextVar("request_id", default=None)


# 用户模型
class User(BaseModel):
    id: int
    username: str
    email: Optional[str] = None


# 模拟用户服务
class UserService:
    @staticmethod
    def get_current_user_id() -> int:
        """在service中直接获取当前用户ID"""
        user = current_user_ctx.get()
        if not user:
            raise RuntimeError("No current user available")
        return user["id"]

    @staticmethod
    def get_current_user() -> dict:
        """获取完整的当前用户信息"""
        user = current_user_ctx.get()
        if not user:
            raise RuntimeError("No current user available")
        return user


# 业务服务示例
class TaskService:
    def create_task(self, task_data: dict):
        """创建任务时自动添加当前用户ID"""
        current_user_id = UserService.get_current_user_id()

        # 这里模拟数据库操作
        task = {
            **task_data,
            "created_by": current_user_id,
            "created_at": "2023-10-01 12:00:00"
        }

        print(f"Task created by user {current_user_id}: {task}")
        return task

    def get_user_tasks(self):
        """获取当前用户的任务"""
        user = current_user_ctx.get()
        current_user_id = UserService.get_current_user_id()

        # 模拟根据用户ID查询任务
        return [{"id": 1, "title": "Sample task", "user_id": current_user_id}]


# 中间件：设置上下文
@app.middleware("http")
async def set_context_vars(request: Request, call_next):
    # 为每个请求生成唯一ID
    request_id = str(uuid.uuid4())
    request_id_token = request_id_ctx.set(request_id)

    # 尝试提取用户信息
    user_token = None
    try:
        auth_header = request.headers.get("Authorization")
        if auth_header and auth_header.startswith("Bearer "):
            token = auth_header.replace("Bearer ", "")
            user = await decode_token_and_get_user(token)  # 您的认证逻辑
            user_token = current_user_ctx.set(user)

        response = await call_next(request)
        return response
    finally:
        # 清理上下文
        request_id_ctx.reset(request_id_token)
        if user_token:
            current_user_ctx.reset(user_token)


# 模拟认证函数
async def decode_token_and_get_user(token: str) -> dict:
    # 这里应该是您的实际认证逻辑，例如JWT解码或数据库查询
    # 简单模拟：根据token返回用户信息
    if token == "valid_token_123":
        return {"id": 123, "username": "john_doe", "email": "john@example.com"}
    elif token == "valid_token_456":
        return {"id": 456, "username": "jane_doe", "email": "jane@example.com"}
    else:
        return None


# 依赖项：用于路由层认证
async def get_current_user_route(token: str = Depends(oauth2_scheme)) -> dict:
    """路由层的用户认证"""
    user = await decode_token_and_get_user(token)
    if not user:
        raise HTTPException(status_code=401, detail="Invalid credentials")
    return user


# 路由处理函数
@app.post("/tasks")
async def create_task(
        task_data: dict,
        current_user: dict = Depends(get_current_user_route)
):
    """创建任务"""
    # 不需要显式传递user_id到service！
    task_service = TaskService()
    task = task_service.create_task(task_data)
    return {"task": task, "message": "Task created successfully"}


@app.get("/tasks")
async def get_tasks(current_user: dict = Depends(get_current_user_route)):
    """获取当前用户的任务"""
    task_service = TaskService()
    tasks = task_service.get_user_tasks()
    return {"tasks": tasks}


@app.get("/users/me")
async def read_users_me(current_user: dict = Depends(get_current_user_route)):
    """获取当前用户信息"""
    return current_user


# 测试端点 - 直接在路由中获取上下文用户
@app.get("/test-context")
async def test_context():
    """测试直接通过上下文获取用户（不通过依赖注入）"""
    try:
        user = UserService.get_current_user()
        return {"message": "Successfully got user from context", "user": user}
    except RuntimeError as e:
        return {"error": str(e)}


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