47 lines
1.3 KiB
Python
47 lines
1.3 KiB
Python
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()
|