asyncexitstack.py 1.0 KB

12345678910111213141516171819202122232425
  1. from typing import Optional
  2. from fastapi.concurrency import AsyncExitStack
  3. from starlette.types import ASGIApp, Receive, Scope, Send
  4. class AsyncExitStackMiddleware:
  5. def __init__(self, app: ASGIApp, context_name: str = "fastapi_astack") -> None:
  6. self.app = app
  7. self.context_name = context_name
  8. async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
  9. dependency_exception: Optional[Exception] = None
  10. async with AsyncExitStack() as stack:
  11. scope[self.context_name] = stack
  12. try:
  13. await self.app(scope, receive, send)
  14. except Exception as e:
  15. dependency_exception = e
  16. raise e
  17. if dependency_exception:
  18. # This exception was possibly handled by the dependency but it should
  19. # still bubble up so that the ServerErrorMiddleware can return a 500
  20. # or the ExceptionMiddleware can catch and handle any other exceptions
  21. raise dependency_exception