初始提交

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()