development, python

fastapi 中使用ContextVar来访问当前请求信息

By jere on 2021-05-07 0 Comments

在rails 中可以在controller或者service中使用current_user 来获取当前已登录的用户,python也有类似的机制,那便是ContextVar.

这篇文章说的很清楚了.

可以参考如下代码
from uuid import uuid4

from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
from starlette.requests import Request
from starlette.responses import Response

from app.core.context import context, REQUEST_ID


class RequestContextMiddleware(BaseHTTPMiddleware):
    async def dispatch(
            self, request: Request, call_next: RequestResponseEndpoint
    ) -> Response:

        # .....
        context.set(REQUEST_ID, str(uuid4()))

        res = await call_next(request)

        return res
from collections import UserDict

from contextvars import ContextVar, copy_context
from typing import Dict, Any

_context_storage: ContextVar[Dict[Any, Any]] = ContextVar(
    "your key name"
)

REQUEST_ID = 'REQUEST_ID'
CURRENT_USER = 'CURRENT_USER'


# bla bla....


class Context(UserDict):
    def __init__(self, *args: Any, **kwargs: Any):  # noqa
        # not calling super on purpose
        if args or kwargs:
            raise AttributeError("Can't instantiate with attributes")

    @property
    def data(self) -> dict:  # type: ignore
        return _context_storage.get()

    def exists(self) -> bool:  # noqa
        return _context_storage in copy_context()

    def copy(self) -> dict:  # type: ignore
        """Read only context data."""
        import copy

        return copy.copy(self.data)

    @staticmethod
    def init():
        _context_storage.set(dict())

    def set(self, key: Any, value: Any):
        data = self.data
        data[key] = value

    def current_user(self):
        return self.data.get(CURRENT_USER)


context = Context()
打赏

Leave a comment

您的电子邮箱地址不会被公开。