初始提交

shierfeng-fastapi-backend/middlewares/trace_middleware/__init__.py

@@ -0,0 +1,18 @@
+from fastapi import FastAPI
+
+from .ctx import TraceCtx
+from .middle import TraceASGIMiddleware
+
+__all__ = ('TraceASGIMiddleware', 'TraceCtx')
+
+__version__ = '0.1.0'
+
+
+def add_trace_middleware(app: FastAPI) -> None:
+    """
+    添加trace中间件
+
+    :param app: FastAPI对象
+    :return:
+    """
+    app.add_middleware(TraceASGIMiddleware)

shierfeng-fastapi-backend/middlewares/trace_middleware/ctx.py

@@ -0,0 +1,64 @@
+import contextvars
+from uuid import uuid4
+
+CTX_TRACE_ID: contextvars.ContextVar[str] = contextvars.ContextVar('trace-id', default='')
+CTX_REQUEST_ID: contextvars.ContextVar[str] = contextvars.ContextVar('request-id', default='')
+CTX_SPAN_ID: contextvars.ContextVar[str] = contextvars.ContextVar('span-id', default='')
+CTX_REQUEST_PATH: contextvars.ContextVar[str] = contextvars.ContextVar('request-path', default='')
+CTX_REQUEST_METHOD: contextvars.ContextVar[str] = contextvars.ContextVar('request-method', default='')
+
+
+class TraceCtx:
+    @staticmethod
+    def set_trace_id() -> str:
+        _id = uuid4().hex
+        CTX_TRACE_ID.set(_id)
+        return _id
+
+    @staticmethod
+    def get_trace_id() -> str:
+        return CTX_TRACE_ID.get()
+
+    @staticmethod
+    def set_request_id() -> str:
+        _id = uuid4().hex
+        CTX_REQUEST_ID.set(_id)
+        return _id
+
+    @staticmethod
+    def get_request_id() -> str:
+        return CTX_REQUEST_ID.get()
+
+    @staticmethod
+    def set_span_id() -> str:
+        _id = uuid4().hex
+        CTX_SPAN_ID.set(_id)
+        return _id
+
+    @staticmethod
+    def get_span_id() -> str:
+        return CTX_SPAN_ID.get()
+
+    @staticmethod
+    def set_request_path(path: str) -> None:
+        CTX_REQUEST_PATH.set(path)
+
+    @staticmethod
+    def get_request_path() -> str:
+        return CTX_REQUEST_PATH.get()
+
+    @staticmethod
+    def set_request_method(method: str) -> None:
+        CTX_REQUEST_METHOD.set(method)
+
+    @staticmethod
+    def get_request_method() -> str:
+        return CTX_REQUEST_METHOD.get()
+
+    @staticmethod
+    def clear() -> None:
+        CTX_TRACE_ID.set('')
+        CTX_REQUEST_ID.set('')
+        CTX_SPAN_ID.set('')
+        CTX_REQUEST_PATH.set('')
+        CTX_REQUEST_METHOD.set('')

shierfeng-fastapi-backend/middlewares/trace_middleware/middle.py

@@ -0,0 +1,46 @@
+from functools import wraps
+
+from starlette.types import ASGIApp, Message, Receive, Scope, Send
+
+from .ctx import TraceCtx
+from .span import Span, get_current_span
+
+
+class TraceASGIMiddleware:
+    """
+    fastapi-example:
+        app = FastAPI()
+        app.add_middleware(TraceASGIMiddleware)
+    """
+
+    def __init__(self, app: ASGIApp) -> None:
+        self.app = app
+
+    @staticmethod
+    async def my_receive(receive: Receive, span: Span) -> Receive:
+        await span.request_before()
+
+        @wraps(receive)
+        async def my_receive() -> Message:
+            message = await receive()
+            await span.request_after(message)
+            return message
+
+        return my_receive
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        if scope['type'] != 'http':
+            await self.app(scope, receive, send)
+            return
+
+        async with get_current_span(scope) as span:
+            handle_outgoing_receive = await self.my_receive(receive, span)
+
+            async def handle_outgoing_request(message: 'Message') -> None:
+                await span.response(message)
+                await send(message)
+
+            try:
+                await self.app(scope, handle_outgoing_receive, handle_outgoing_request)
+            finally:
+                TraceCtx.clear()

shierfeng-fastapi-backend/middlewares/trace_middleware/span.py

@@ -0,0 +1,54 @@
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
+
+from starlette.types import Message, Scope
+
+from .ctx import TraceCtx
+
+
+class Span:
+    """
+    整个http生命周期：
+        request(before) --> request(after) --> response(before) --> response(after)
+    """
+
+    def __init__(self, scope: Scope) -> None:
+        self.scope = scope
+
+    async def request_before(self) -> None:
+        """
+        request_before: 处理header信息等, 如记录请求体信息
+        """
+        TraceCtx.set_trace_id()
+        TraceCtx.set_request_id()
+        TraceCtx.set_span_id()
+        TraceCtx.set_request_path(self.scope.get('path', ''))
+        TraceCtx.set_request_method(self.scope.get('method', ''))
+
+    async def request_after(self, message: Message) -> Message:
+        """
+        request_after: 处理请求bytes， 如记录请求参数
+
+        example:
+            message: {'type': 'http.request', 'body': b'{\r\n    "name": "\xe8\x8b\x8f\xe8\x8b\x8f\xe8\x8b\x8f"\r\n}', 'more_body': False}
+        """
+        return message
+
+    async def response(self, message: Message) -> Message:
+        """
+        if message['type'] == "http.response.start":   -----> request-before
+            pass
+        if message['type'] == "http.response.body":    -----> request-after
+            message.get('body', b'')
+            pass
+        """
+        if message['type'] == 'http.response.start':
+            message['headers'].append((b'request-id', TraceCtx.get_request_id().encode()))
+            message['headers'].append((b'trace-id', TraceCtx.get_trace_id().encode()))
+            message['headers'].append((b'span-id', TraceCtx.get_span_id().encode()))
+        return message
+
+
+@asynccontextmanager
+async def get_current_span(scope: Scope) -> AsyncGenerator[Span, None]:
+    yield Span(scope)