Update request state usage to context variable (#853)

* Update request state usage to context variable

* Update context to custom ctx

* Fix the exception interception in opera log

* restore elapsed

Author: wu-clan

backend/app/admin/service/login_log_service.py

@@ -6,6 +6,7 @@
 
 from backend.app.admin.crud.crud_login_log import login_log_dao
 from backend.app.admin.schema.login_log import CreateLoginLogParam, DeleteLoginLogParam
+from backend.common.context import ctx
 from backend.common.log import log
 from backend.database.db import async_db_session
 
@@ -53,14 +54,14 @@ async def create(
                 user_uuid=user_uuid,
                 username=username,
                 status=status,
-                ip=request.state.ip,
-                country=request.state.country,
-                region=request.state.region,
-                city=request.state.city,
-                user_agent=request.state.user_agent,
-                browser=request.state.browser,
-                os=request.state.os,
-                device=request.state.device,
+                ip=ctx.ip,
+                country=ctx.country,
+                region=ctx.region,
+                city=ctx.city,
+                user_agent=ctx.user_agent,
+                browser=ctx.browser,
+                os=ctx.os,
+                device=ctx.device,
                 msg=msg,
                 login_time=login_time,
             )

backend/app/admin/service/user_service.py

@@ -14,6 +14,7 @@
     ResetPasswordParam,
     UpdateUserParam,
 )
+from backend.common.context import ctx
 from backend.common.enums import UserPermissionType
 from backend.common.exception import errors
 from backend.common.response.response_code import CustomErrorCode
@@ -249,12 +250,12 @@ async def update_email(*, request: Request, captcha: str, email: str) -> int:
             user = await user_dao.get(db, token_payload.id)
             if not user:
                 raise errors.NotFoundError(msg='用户不存在')
-            captcha_code = await redis_client.get(f'{settings.EMAIL_CAPTCHA_REDIS_PREFIX}:{request.state.ip}')
+            captcha_code = await redis_client.get(f'{settings.EMAIL_CAPTCHA_REDIS_PREFIX}:{ctx.ip}')
             if not captcha_code:
                 raise errors.RequestError(msg='验证码已失效，请重新获取')
             if captcha != captcha_code:
                 raise errors.CustomError(error=CustomErrorCode.CAPTCHA_ERROR)
-            await redis_client.delete(f'{settings.EMAIL_CAPTCHA_REDIS_PREFIX}:{request.state.ip}')
+            await redis_client.delete(f'{settings.EMAIL_CAPTCHA_REDIS_PREFIX}:{ctx.ip}')
             count = await user_dao.update_email(db, token_payload.id, email)
             await redis_client.delete(f'{settings.JWT_USER_REDIS_PREFIX}:{user.id}')
             return count

backend/common/context.py

@@ -0,0 +1,32 @@
+from datetime import datetime
+from typing import Any, Protocol
+
+from starlette_context.ctx import _Context, context
+
+
+class TypedContextProtocol(Protocol):
+    perf_time: float
+    start_time: datetime
+
+    ip: str
+    country: str | None
+    region: str | None
+    city: str | None
+
+    user_agent: str
+    os: str | None
+    browser: str | None
+    device: str | None
+
+    permission: str | None
+
+
+class TypedContext(TypedContextProtocol, _Context):
+    def __getattr__(self, name: str) -> Any:
+        return context.get(name)
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        context[name] = value
+
+
+ctx = TypedContext()

backend/common/exception/exception_handler.py

@@ -4,6 +4,7 @@
 from starlette.exceptions import HTTPException
 from uvicorn.protocols.http.h11_impl import STATUS_PHRASES
 
+from backend.common.context import ctx
 from backend.common.exception.errors import BaseExceptionError
 from backend.common.i18n import i18n, t
 from backend.common.response.response_code import CustomResponseCode, StandardResponseCode
@@ -72,8 +73,8 @@ async def _validation_exception_handler(request: Request, exc: RequestValidation
         'msg': msg,
         'data': data,
     }
-    request.state.__request_validation_exception__ = content  # 用于在中间件中获取异常信息
-    content.update(trace_id=get_request_trace_id(request))
+    ctx.__request_validation_exception__ = content  # 用于在中间件中获取异常信息
+    content.update(trace_id=get_request_trace_id())
     return MsgSpecJSONResponse(status_code=StandardResponseCode.HTTP_422, content=content)
 
 
@@ -96,8 +97,8 @@ async def http_exception_handler(request: Request, exc: HTTPException):
         else:
             res = response_base.fail(res=CustomResponseCode.HTTP_400)
             content = res.model_dump()
-        request.state.__request_http_exception__ = content
-        content.update(trace_id=get_request_trace_id(request))
+        ctx.__request_http_exception__ = content
+        content.update(trace_id=get_request_trace_id())
         return MsgSpecJSONResponse(
             status_code=_get_exception_code(exc.status_code),
             content=content,
@@ -144,8 +145,8 @@ async def assertion_error_handler(request: Request, exc: AssertionError):
         else:
             res = response_base.fail(res=CustomResponseCode.HTTP_500)
             content = res.model_dump()
-        request.state.__request_assertion_error__ = content
-        content.update(trace_id=get_request_trace_id(request))
+        ctx.__request_assertion_error__ = content
+        content.update(trace_id=get_request_trace_id())
         return MsgSpecJSONResponse(
             status_code=StandardResponseCode.HTTP_500,
             content=content,
@@ -165,8 +166,8 @@ async def custom_exception_handler(request: Request, exc: BaseExceptionError):
             'msg': str(exc.msg),
             'data': exc.data or None,
         }
-        request.state.__request_custom_exception__ = content
-        content.update(trace_id=get_request_trace_id(request))
+        ctx.__request_custom_exception__ = content
+        content.update(trace_id=get_request_trace_id())
         return MsgSpecJSONResponse(
             status_code=_get_exception_code(exc.code),
             content=content,
@@ -191,7 +192,7 @@ async def all_unknown_exception_handler(request: Request, exc: Exception):
         else:
             res = response_base.fail(res=CustomResponseCode.HTTP_500)
             content = res.model_dump()
-        content.update(trace_id=get_request_trace_id(request))
+        content.update(trace_id=get_request_trace_id())
         return MsgSpecJSONResponse(
             status_code=StandardResponseCode.HTTP_500,
             content=content,

backend/common/log.py

@@ -4,12 +4,13 @@
 import re
 import sys
 
-from asgi_correlation_id import correlation_id
 from loguru import logger
 
+from backend.common.context import ctx
 from backend.core.conf import settings
 from backend.core.path_conf import LOG_DIR
 from backend.utils.timezone import timezone
+from backend.utils.trace_id import get_request_trace_id
 
 
 class InterceptHandler(logging.Handler):
@@ -75,11 +76,14 @@ def setup_logging() -> None:
     # 移除 loguru 默认处理器
     logger.remove()
 
-    # correlation_id 过滤器
-    # https://github.com/snok/asgi-correlation-id/issues/7
-    def correlation_id_filter(record: logging.LogRecord) -> logging.LogRecord:
-        cid = correlation_id.get(settings.TRACE_ID_LOG_DEFAULT_VALUE)
-        record['correlation_id'] = cid[: settings.TRACE_ID_LOG_LENGTH]
+    # request_id 过滤器
+    def request_id_filter(record: logging.LogRecord) -> logging.LogRecord:
+        if ctx.exists():
+            rid = get_request_trace_id()
+            record['request_id'] = rid[: settings.TRACE_ID_LOG_LENGTH]
+        else:
+            record['request_id'] = settings.TRACE_ID_LOG_DEFAULT_VALUE
+
         return record
 
     # 配置 loguru 处理器
@@ -89,7 +93,7 @@ def correlation_id_filter(record: logging.LogRecord) -> logging.LogRecord:
                 'sink': sys.stdout,
                 'level': settings.LOG_STD_LEVEL,
                 'format': default_formatter,
-                'filter': lambda record: correlation_id_filter(record),
+                'filter': lambda record: request_id_filter(record),
             },
         ],
     )

backend/common/security/permission.py

@@ -3,6 +3,7 @@
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from backend.app.admin.crud.crud_data_scope import data_scope_dao
+from backend.common.context import ctx
 from backend.common.enums import RoleDataRuleExpressionType, RoleDataRuleOperatorType
 from backend.common.exception import errors
 from backend.core.conf import settings
@@ -38,7 +39,7 @@ async def __call__(self, request: Request) -> None:
             if not isinstance(self.value, str):
                 raise errors.ServerError
             # 附加权限标识到请求状态
-            request.state.permission = self.value
+            ctx.permission = self.value
 
 
 async def filter_data_permission(db: AsyncSession, request: Request) -> ColumnElement[bool]:  # noqa: C901

backend/common/security/rbac.py

@@ -1,5 +1,6 @@
 from fastapi import Depends, Request
 
+from backend.common.context import ctx
 from backend.common.enums import MethodType, StatusType
 from backend.common.exception import errors
 from backend.common.log import log
@@ -49,7 +50,7 @@ async def rbac_verify(request: Request, _token: str = DependsJwtAuth) -> None:
 
     # RBAC 鉴权
     if settings.RBAC_ROLE_MENU_MODE:
-        path_auth_perm = getattr(request.state, 'permission', None)
+        path_auth_perm = ctx.permission
 
         # 没有菜单操作权限标识不校验
         if not path_auth_perm:

backend/core/conf.py

@@ -150,7 +150,7 @@ class Settings(BaseSettings):
 
     # 日志
     LOG_FORMAT: str = (
-        '<green>{time:YYYY-MM-DD HH:mm:ss.SSS}</> | <lvl>{level: <8}</> | <cyan>{correlation_id}</> | <lvl>{message}</>'
+        '<green>{time:YYYY-MM-DD HH:mm:ss.SSS}</> | <lvl>{level: <8}</> | <cyan>{request_id}</> | <lvl>{message}</>'
     )
 
     # 日志（控制台）

backend/core/registrar.py

@@ -6,18 +6,20 @@
 
 import socketio
 
-from asgi_correlation_id import CorrelationIdMiddleware
 from fastapi import Depends, FastAPI
 from fastapi_limiter import FastAPILimiter
 from fastapi_pagination import add_pagination
 from starlette.middleware.authentication import AuthenticationMiddleware
 from starlette.middleware.cors import CORSMiddleware
 from starlette.staticfiles import StaticFiles
 from starlette.types import ASGIApp
+from starlette_context.middleware import ContextMiddleware
+from starlette_context.plugins import RequestIdPlugin
 
 from backend import __version__
 from backend.common.exception.exception_handler import register_exception
 from backend.common.log import set_custom_logfile, setup_logging
+from backend.common.response.response_code import StandardResponseCode
 from backend.core.conf import settings
 from backend.core.path_conf import STATIC_DIR, UPLOAD_DIR
 from backend.database.db import create_tables
@@ -154,8 +156,15 @@ def register_middleware(app: FastAPI) -> None:
     # Access log
     app.add_middleware(AccessMiddleware)
 
-    # Trace ID
-    app.add_middleware(CorrelationIdMiddleware, validator=False)
+    # ContextVar
+    app.add_middleware(
+        ContextMiddleware,
+        plugins=[RequestIdPlugin(validate=True)],
+        default_error_response=MsgSpecJSONResponse(
+            content={'code': StandardResponseCode.HTTP_400, 'msg': 'BAD_REQUEST', 'data': None},
+            status_code=StandardResponseCode.HTTP_400,
+        ),
+    )
 
 
 def register_router(app: FastAPI) -> None:

backend/middleware/access_middleware.py

@@ -3,6 +3,7 @@
 from fastapi import Request, Response
 from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
 
+from backend.common.context import ctx
 from backend.common.log import log
 from backend.utils.timezone import timezone
 
@@ -24,21 +25,19 @@ async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -
             log.debug(f'--> 请求开始[{path}]')
 
         perf_time = time.perf_counter()
-        request.state.perf_time = perf_time
+        ctx.perf_time = perf_time
 
         start_time = timezone.now()
-        request.state.start_time = start_time
+        ctx.start_time = start_time
 
         response = await call_next(request)
 
-        elapsed = (time.perf_counter() - perf_time) * 1000
-
         if request.method != 'OPTIONS':
             log.debug('<-- 请求结束')
 
             log.info(
                 f'{request.client.host: <15} | {request.method: <8} | {response.status_code: <6} | '
-                f'{path} | {elapsed:.3f}ms',
+                f'{path} | {(time.perf_counter() - perf_time) * 1000:.3f}ms',
             )
 
         return response