FastAPI请求会话context上下文中间件

一、背景

FastAPI框架当中，我们知道可以通过request对象中获取到前端请求的相关内容。包括HTTP请求头、请求体、请求Query参数等等。

有些情况，我们想在service层或者更加深层次的位置获取本次会话的上下文数据。例如有这么一个场景，sqlmodel的ORM类，有一些参数是每次会跟着本次请求会话携带过来的，需要自动进行填充。要不然开发每次都要从request对象获取jwt token/请求头内容进行填充，这个工作量和代码可维护性太差了.

详细场景例子如下：

1、前端通过请求头传递jwt token至后端

2、后端可以利用FastAPI的middleware中间件进行拦截，校验token合法性，合法则进一步解析为dict或者封装好的class对象

3、将本次会话的token解析后的payload数据，无论是dict还是class实例，挂载设置到request.state这个属性下, 假设就是request.state.jwt_payload jwt_payload类型是一个dict字典

4、FastAPI的router路由函数，就可以自动注入request对象，函数中通过request.state.jwt_payload 获取到这个dict字典数据，进行实际的业务处理

二、JWT中间件代码示例

1、认证中间件拦截代码 jwt_auth.py

python 复制代码

from typing import List

from loguru import logger
from starlette import status
from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
from starlette.requests import Request
from starlette.responses import Response
import jwt
import time

from app.api.common import BaseAPI
from app.core.config import settings


class JwtAuthMiddleware(BaseHTTPMiddleware):
    """
        JWT请求头token拦截，进行鉴权、以及在将解析后的jwt payload放入
        request.state.jwt_payload属性
        请求接口通过: request.state.jwt_payload获取
    """

    def __init__(self, app, public_key: str, algorithms: List[str], exclude_path, *args, **kwargs):
        super().__init__(app)
        self.public_key = public_key
        self.algorithms = algorithms
        self.exclude_path = exclude_path

    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
        jwt_token = request.headers.get("Authorization", "")
        if not jwt_token:
            return BaseAPI.error(code=status.HTTP_400_BAD_REQUEST, message="token缺失",
                                 status_code=status.HTTP_400_BAD_REQUEST)
        try:
            if len(jwt_token.split(".")) != 3:
                return BaseAPI.error(code=status.HTTP_400_BAD_REQUEST, message="token错误",
                                     status_code=status.HTTP_400_BAD_REQUEST)
            decode_jwt_body = jwt.decode(jwt_token, self.public_key, algorithms=self.algorithms)
            exp_ts = decode_jwt_body.get("exp", 0)
            if time.time() > exp_ts:
                BaseAPI.error(code=status.HTTP_401_UNAUTHORIZED, message="token已过期",
                              status_code=status.HTTP_401_UNAUTHORIZED)
            # 解析jwt token拿到的dict字典数据,设置到request.state属性下
            request.state.jwt_payload = decode_jwt_body
        except jwt.InvalidTokenError as e:
            logger.error(e)
            return BaseAPI.error(code=status.HTTP_500_INTERNAL_SERVER_ERROR, message="token解析错误",
                                 status_code=status.HTTP_500_INTERNAL_SERVER_ERROR)
        response = await call_next(request)
        return response

2、main.py中程序中app进行应用

python 复制代码

app = FastAPI(lifespan=core.lifespan, **settings.docs_route_config)

app.add_middleware(JwtAuthMiddleware, public_key=settings.JWT_SECRET_KEY, algorithms=[settings.JWT_ALGORITHM],exclude_path=settings.JWT_AUTH_EXCLUDE_PATH)

3、router路由函数进行使用，获取request.state.jwt_payload

python 复制代码

    async def summary(
                      jwt_user: JwtUser = Depends(deps_get_jwt_user),
                      service: SummarizeService = Depends()):
        pass

deps_get_jwt_user函数定义，从request.state获取解析好的token数据，返回封装的Class:

python 复制代码

class JwtUser(BaseModel):
    """
        jwt payload内容
    """
    username: str
    user_id: str

def deps_get_jwt_user(request: Request) -> JwtUser:
    """
        依赖注入, 从拦截器的request对象获取已经解析好的jwt payload内容
    """
    payload = request.state.jwt_payload
    user = JwtUser(username=payload["username"],user_id=payload["userId"])
    return user

三、通过request对象获取的优缺点

1、优点

直接通过在中间件设置，然后在路由接口函数再通过request对象获取，方便、快捷，也很容易理解

2、缺点

如果我的代码层次很深的地方也想获取到本次HTTP会话请求的一些信息，那么我需要把这个request对象一直往下传递，或者在路由接口处，先把值从request.state取出来，但是还是继续往下传递, 当传递的层次深了以后，代码就不好维护了，开发也不太方便写，每次都把这个request往下传。

四、解决深层次传递request请求会话上下文context

1、问题分析

我们在想，可不可以这样，有没有这种机制，就是在本次请求会话的范围内，存在一个全局的容器，可以存储本次会话的相关信息。但是这个全局的容器只在本次请求有效，隔离，不影响其他请求或者整个app应用。

这样不管代码层次多深，我都可以在中间件先拦截、设置，最后在代码里面任意位置，都可以获取到，这样就不需要从FastAPI的路由函数一直往下传递了。

这个会话全局数据，一般我们称为会话context上下文

2、starlette-context会话上下文管理中间件

github地址: https://github.com/tomwojcik/starlette-context

FastAPI本身就是基于starlette构建的. 所以使用starlette-context这个中间件来管理请求会话context上下文，底层使用 contextvars 实现.

contextvars可以在本次协程建立一个全局的存储容器，作为本次协程的上下文. 都可以从里面设置数据和获取数据.

contextvars的知识大家可以下来自行学习，这里不展开讨论。

简单理解就是，在FastAPI的中间件拦截器里面，你使用这个starlette-context中间件去拦截和设置你的数据，那么后续本次会话，你可以在任何位置直接、安全地、隔离地获取上下文数据，而不需要一直传递request对象到深层次的位置才能获取数据。

并且starlette-context的上下文是安全的、隔离的，只针对本次的HTTP请求，不会造成全局数据污染或者干扰别的HTTP请求。

3、使用方式

1、编写一个Plugin插件

python 复制代码

class SessionAuthContext:
    """
        HTTP会话认证信息上下文
    """
    jwt_user: JwtUser

    def __init__(self, jwt_user: JwtUser):
        self.jwt_user = jwt_user

class JwtPlugin(Plugin):
    
    # key就是后期你通过context对象获取数据设置的key
    key = "auth"

    async def extract_value_from_header_by_key(
            self, request: Union[Request, HTTPConnection]
    ) -> Optional[Any]:
        """
            从请求state中获取已经经过jwt_auth中间件处理的数据,设置进入上下文组件

            from starlette_context import context
            data: SessionAuthContext  = context.get("auth")
        """
        if hasattr(request.state, "jwt_payload"):
            jwt_user = deps_get_jwt_user(request)
            data = SessionAuthContext(token=request.state.jwt_token, jwt_user=jwt_user)
            return data
        else:
            return None

2、main.py引用中间件,加载plugin

python 复制代码

# FILO  先进后出, 中间件执行顺序从下往上
# 设置请求上下文context中间件
app.add_middleware(ContextMiddleware, plugins=[context_plugins.JwtPlugin()])

# JWT 中间件
app.add_middleware(JwtAuthMiddleware, public_key=settings.JWT_SECRET_KEY, algorithms=[settings.JWT_ALGORITHM],
                   exclude_path=settings.JWT_AUTH_EXCLUDE_PATH)

3、sqlmodel 模型类，自动填充会话拿到的数据

python 复制代码

def fill_create_by() -> str:
    """"
        填充sqlmodel模型 默认字段值, 从本次请求的jwt token payload中自动获取且设置值
        从starlette_context导入context  不要导错路径

        from starlette_context import context
    """
    auth_ctx: SessionAuthContext = context.get("auth")
    return auth_ctx.jwt_user.user_id


class WorkspaceMembersModel(SQLModel, table=True):

    __tablename__ = 'workspace_members'
    __table_args__ = {'comment': '工作空间成员表'}

    id: Optional[int] = Field(default_factory=gen_snowflake_id, primary_key=True, nullable=False,
                              description='工作空间ID')
    user_id: Optional[int] = Field(nullable=False, description='用户ID')
    nickname: str = Field(default='', max_length=30, description='成员在团队中的昵称')
     
    # default_factory 设置为 fill_create_by函数
    create_by: Optional[int] = Field(nullable=False, default_factory=fill_create_by, description='创建用户id')


if __main__ == "__main__":
    # 创建实例没传递值，则自动触发fill_create_by函数的执行
    # fill_create_by  从 context中获取auth属性，再获取前面中间件拦截设置的数据
    # 最终实现，深层次代码不需要一直传递request对象，也能优雅地拿到请求上下文数据
    model = WorkspaceMembersModel()
    print(model.create_by)

五、总结

类比Java我们可能也会把数据放在servlet阶段性，只存在本次会话的对象上。这样深层次安全地获取本次HTTP请求会话上下文数据就变得很方便，也安全。不需要一直把request对象往下传递，层次一深，代码就变得非常臃肿、难以维护。

但是这里一定要注意，这个context只针对本次会话有效，不污染全局变量、不污染线程或者进程变量，这是首先要考虑的点。否则用不好，导致，A用户的请求上下文被B用户使用到了，那就严重事故了!