🐛 Fix support for StreamingResponses with dependencies with yield or UploadFiles, close after the response is done (#14099)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: tiangolo

docs/en/docs/advanced/advanced-dependencies.md

@@ -63,3 +63,91 @@ In the chapters about security, there are utility functions that are implemented
 If you understood all this, you already know how those utility tools for security work underneath.
 
 ///
+
+## Dependencies with `yield`, `HTTPException`, `except` and Background Tasks { #dependencies-with-yield-httpexception-except-and-background-tasks }
+
+/// warning
+
+You most probably don't need these technical details.
+
+These details are useful mainly if you had a FastAPI application older than 0.118.0 and you are facing issues with dependencies with `yield`.
+
+///
+
+Dependencies with `yield` have evolved over time to account for the different use cases and to fix some issues, here's a summary of what has changed.
+
+### Dependencies with `yield` and `StreamingResponse`, Technical Details { #dependencies-with-yield-and-streamingresponse-technical-details }
+
+Before FastAPI 0.118.0, if you used a dependency with `yield`, it would run the exit code after the *path operation function* returned but right before sending the response.
+
+The intention was to avoid holding resources for longer than necessary, waiting for the response to travel through the network.
+
+This change also meant that if you returned a `StreamingResponse`, the exit code of the dependency with `yield` would have been already run.
+
+For example, if you had a database session in a dependency with `yield`, the `StreamingResponse` would not be able to use that session while streaming data because the session would have already been closed in the exit code after `yield`.
+
+This behavior was reverted in 0.118.0, to make the exit code after `yield` be executed after the response is sent.
+
+/// info
+
+As you will see below, this is very similar to the behavior before version 0.106.0, but with several improvements and bug fixes for corner cases.
+
+///
+
+#### Use Cases with Early Exit Code { #use-cases-with-early-exit-code }
+
+There are some use cases with specific conditions that could benefit from the old behavior of running the exit code of dependencies with `yield` before sending the response.
+
+For example, imagine you have code that uses a database session in a dependency with `yield` only to verify a user, but the database session is never used again in the *path operation function*, only in the dependency, **and** the response takes a long time to be sent, like a `StreamingResponse` that sends data slowly, but for some reason doesn't use the database.
+
+In this case, the database session would be held until the response is finished being sent, but if you don't use it, then it wouldn't be necessary to hold it.
+
+Here's how it could look like:
+
+{* ../../docs_src/dependencies/tutorial013_an_py310.py *}
+
+The exit code, the automatic closing of the `Session` in:
+
+{* ../../docs_src/dependencies/tutorial013_an_py310.py ln[19:21] *}
+
+...would be run after the the response finishes sending the slow data:
+
+{* ../../docs_src/dependencies/tutorial013_an_py310.py ln[30:38] hl[31:33] *}
+
+But as `generate_stream()` doesn't use the database session, it is not really necessary to keep the session open while sending the response.
+
+If you have this specific use case using SQLModel (or SQLAlchemy), you could explicitly close the session after you don't need it anymore:
+
+{* ../../docs_src/dependencies/tutorial014_an_py310.py ln[24:28] hl[28] *}
+
+That way the session would release the database connection, so other requests could use it.
+
+If you have a different use case that needs to exit early from a dependency with `yield`, please create a <a href="https://github.com/fastapi/fastapi/discussions/new?category=questions" class="external-link" target="_blank">GitHub Discussion Question</a> with your specific use case and why you would benefit from having early closing for dependencies with `yield`.
+
+If there are compelling use cases for early closing in dependencies with `yield`, I would consider adding a new way to opt in to early closing.
+
+### Dependencies with `yield` and `except`, Technical Details { #dependencies-with-yield-and-except-technical-details }
+
+Before FastAPI 0.110.0, if you used a dependency with `yield`, and then you captured an exception with `except` in that dependency, and you didn't raise the exception again, the exception would be automatically raised/forwarded to any exception handlers or the internal server error handler.
+
+This was changed in version 0.110.0 to fix unhandled memory consumption from forwarded exceptions without a handler (internal server errors), and to make it consistent with the behavior of regular Python code.
+
+### Background Tasks and Dependencies with `yield`, Technical Details { #background-tasks-and-dependencies-with-yield-technical-details }
+
+Before FastAPI 0.106.0, raising exceptions after `yield` was not possible, the exit code in dependencies with `yield` was executed *after* the response was sent, so [Exception Handlers](../handling-errors.md#install-custom-exception-handlers){.internal-link target=_blank} would have already run.
+
+This was designed this way mainly to allow using the same objects "yielded" by dependencies inside of background tasks, because the exit code would be executed after the background tasks were finished.
+
+This was changed in FastAPI 0.106.0 with the intention to not hold resources while waiting for the response to travel through the network.
+
+/// tip
+
+Additionally, a background task is normally an independent set of logic that should be handled separately, with its own resources (e.g. its own database connection).
+
+So, this way you will probably have cleaner code.
+
+///
+
+If you used to rely on this behavior, now you should create the resources for background tasks inside the background task itself, and use internally only data that doesn't depend on the resources of dependencies with `yield`.
+
+For example, instead of using the same database session, you would create a new database session inside of the background task, and you would obtain the objects from the database using this new session. And then instead of passing the object from the database as a parameter to the background task function, you would pass the ID of that object and then obtain the object again inside the background task function.

docs/en/docs/tutorial/dependencies/dependencies-with-yield.md

@@ -35,7 +35,7 @@ The yielded value is what is injected into *path operations* and other dependenc
 
 {* ../../docs_src/dependencies/tutorial007.py hl[4] *}
 
-The code following the `yield` statement is executed after creating the response but before sending it:
+The code following the `yield` statement is executed after the response:
 
 {* ../../docs_src/dependencies/tutorial007.py hl[5:6] *}
 
@@ -51,7 +51,7 @@ You can use `async` or regular functions.
 
 If you use a `try` block in a dependency with `yield`, you'll receive any exception that was thrown when using the dependency.
 
-For example, if some code at some point in the middle, in another dependency or in a *path operation*, made a database transaction "rollback" or create any other error, you will receive the exception in your dependency.
+For example, if some code at some point in the middle, in another dependency or in a *path operation*, made a database transaction "rollback" or created any other exception, you would receive the exception in your dependency.
 
 So, you can look for that specific exception inside the dependency with `except SomeException`.
 
@@ -95,9 +95,11 @@ This works thanks to Python's <a href="https://docs.python.org/3/library/context
 
 ## Dependencies with `yield` and `HTTPException` { #dependencies-with-yield-and-httpexception }
 
-You saw that you can use dependencies with `yield` and have `try` blocks that catch exceptions.
+You saw that you can use dependencies with `yield` and have `try` blocks that try to execute some code and then run some exit code after `finally`.
 
-The same way, you could raise an `HTTPException` or similar in the exit code, after the `yield`.
+You can also use `except` to catch the exception that was raised and do something with it.
+
+For example, you can raise a different exception, like `HTTPException`.
 
 /// tip
 
@@ -109,7 +111,7 @@ But it's there for you if you need it. 🤓
 
 {* ../../docs_src/dependencies/tutorial008b_an_py39.py hl[18:22,31] *}
 
-An alternative you could use to catch exceptions (and possibly also raise another `HTTPException`) is to create a [Custom Exception Handler](../handling-errors.md#install-custom-exception-handlers){.internal-link target=_blank}.
+If you want to catch exceptions and create a custom response based on that, create a [Custom Exception Handler](../handling-errors.md#install-custom-exception-handlers){.internal-link target=_blank}.
 
 ## Dependencies with `yield` and `except` { #dependencies-with-yield-and-except }
 
@@ -121,7 +123,7 @@ In this case, the client will see an *HTTP 500 Internal Server Error* response a
 
 ### Always `raise` in Dependencies with `yield` and `except` { #always-raise-in-dependencies-with-yield-and-except }
 
-If you catch an exception in a dependency with `yield`, unless you are raising another `HTTPException` or similar, you should re-raise the original exception.
+If you catch an exception in a dependency with `yield`, unless you are raising another `HTTPException` or similar, **you should re-raise the original exception**.
 
 You can re-raise the same exception using `raise`:
 
@@ -178,48 +180,15 @@ After one of those responses is sent, no other response can be sent.
 
 /// tip
 
-This diagram shows `HTTPException`, but you could also raise any other exception that you catch in a dependency with `yield` or with a [Custom Exception Handler](../handling-errors.md#install-custom-exception-handlers){.internal-link target=_blank}.
-
-If you raise any exception, it will be passed to the dependencies with yield, including `HTTPException`. In most cases you will want to re-raise that same exception or a new one from the dependency with `yield` to make sure it's properly handled.
+If you raise any exception in the code from the *path operation function*, it will be passed to the dependencies with yield, including `HTTPException`. In most cases you will want to re-raise that same exception or a new one from the dependency with `yield` to make sure it's properly handled.
 
 ///
 
 ## Dependencies with `yield`, `HTTPException`, `except` and Background Tasks { #dependencies-with-yield-httpexception-except-and-background-tasks }
 
-/// warning
-
-You most probably don't need these technical details, you can skip this section and continue below.
-
-These details are useful mainly if you were using a version of FastAPI prior to 0.106.0 and used resources from dependencies with `yield` in background tasks.
-
-///
-
-### Dependencies with `yield` and `except`, Technical Details { #dependencies-with-yield-and-except-technical-details }
-
-Before FastAPI 0.110.0, if you used a dependency with `yield`, and then you captured an exception with `except` in that dependency, and you didn't raise the exception again, the exception would be automatically raised/forwarded to any exception handlers or the internal server error handler.
-
-This was changed in version 0.110.0 to fix unhandled memory consumption from forwarded exceptions without a handler (internal server errors), and to make it consistent with the behavior of regular Python code.
-
-### Background Tasks and Dependencies with `yield`, Technical Details { #background-tasks-and-dependencies-with-yield-technical-details }
-
-Before FastAPI 0.106.0, raising exceptions after `yield` was not possible, the exit code in dependencies with `yield` was executed *after* the response was sent, so [Exception Handlers](../handling-errors.md#install-custom-exception-handlers){.internal-link target=_blank} would have already run.
-
-This was designed this way mainly to allow using the same objects "yielded" by dependencies inside of background tasks, because the exit code would be executed after the background tasks were finished.
-
-Nevertheless, as this would mean waiting for the response to travel through the network while unnecessarily holding a resource in a dependency with yield (for example a database connection), this was changed in FastAPI 0.106.0.
-
-/// tip
-
-Additionally, a background task is normally an independent set of logic that should be handled separately, with its own resources (e.g. its own database connection).
-
-So, this way you will probably have cleaner code.
-
-///
-
-If you used to rely on this behavior, now you should create the resources for background tasks inside the background task itself, and use internally only data that doesn't depend on the resources of dependencies with `yield`.
-
-For example, instead of using the same database session, you would create a new database session inside of the background task, and you would obtain the objects from the database using this new session. And then instead of passing the object from the database as a parameter to the background task function, you would pass the ID of that object and then obtain the object again inside the background task function.
+Dependencies with `yield` have evolved over time to cover different use cases and fix some issues.
 
+If you want to see what has changed in different versions of FastAPI, you can read more about it in the advanced guide, in [Advanced Dependencies - Dependencies with `yield`, `HTTPException`, `except` and Background Tasks](../../advanced/advanced-dependencies.md#dependencies-with-yield-httpexception-except-and-background-tasks){.internal-link target=_blank}.
 ## Context Managers { #context-managers }
 
 ### What are "Context Managers" { #what-are-context-managers }

docs_src/dependencies/tutorial013_an_py310.py

@@ -0,0 +1,38 @@
+import time
+from typing import Annotated
+
+from fastapi import Depends, FastAPI, HTTPException
+from fastapi.responses import StreamingResponse
+from sqlmodel import Field, Session, SQLModel, create_engine
+
+engine = create_engine("postgresql+psycopg://postgres:postgres@localhost/db")
+
+
+class User(SQLModel, table=True):
+    id: int | None = Field(default=None, primary_key=True)
+    name: str
+
+
+app = FastAPI()
+
+
+def get_session():
+    with Session(engine) as session:
+        yield session
+
+
+def get_user(user_id: int, session: Annotated[Session, Depends(get_session)]):
+    user = session.get(User, user_id)
+    if not user:
+        raise HTTPException(status_code=403, detail="Not authorized")
+
+
+def generate_stream(query: str):
+    for ch in query:
+        yield ch
+        time.sleep(0.1)
+
+
+@app.get("/generate", dependencies=[Depends(get_user)])
+def generate(query: str):
+    return StreamingResponse(content=generate_stream(query))

docs_src/dependencies/tutorial014_an_py310.py

@@ -0,0 +1,39 @@
+import time
+from typing import Annotated
+
+from fastapi import Depends, FastAPI, HTTPException
+from fastapi.responses import StreamingResponse
+from sqlmodel import Field, Session, SQLModel, create_engine
+
+engine = create_engine("postgresql+psycopg://postgres:postgres@localhost/db")
+
+
+class User(SQLModel, table=True):
+    id: int | None = Field(default=None, primary_key=True)
+    name: str
+
+
+app = FastAPI()
+
+
+def get_session():
+    with Session(engine) as session:
+        yield session
+
+
+def get_user(user_id: int, session: Annotated[Session, Depends(get_session)]):
+    user = session.get(User, user_id)
+    if not user:
+        raise HTTPException(status_code=403, detail="Not authorized")
+    session.close()
+
+
+def generate_stream(query: str):
+    for ch in query:
+        yield ch
+        time.sleep(0.1)
+
+
+@app.get("/generate", dependencies=[Depends(get_user)])
+def generate(query: str):
+    return StreamingResponse(content=generate_stream(query))

fastapi/applications.py

@@ -22,6 +22,7 @@
 )
 from fastapi.exceptions import RequestValidationError, WebSocketRequestValidationError
 from fastapi.logger import logger
+from fastapi.middleware.asyncexitstack import AsyncExitStackMiddleware
 from fastapi.openapi.docs import (
     get_redoc_html,
     get_swagger_ui_html,
@@ -36,10 +37,12 @@
 from starlette.exceptions import HTTPException
 from starlette.middleware import Middleware
 from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.middleware.errors import ServerErrorMiddleware
+from starlette.middleware.exceptions import ExceptionMiddleware
 from starlette.requests import Request
 from starlette.responses import HTMLResponse, JSONResponse, Response
 from starlette.routing import BaseRoute
-from starlette.types import ASGIApp, Lifespan, Receive, Scope, Send
+from starlette.types import ASGIApp, ExceptionHandler, Lifespan, Receive, Scope, Send
 from typing_extensions import Annotated, Doc, deprecated
 
 AppType = TypeVar("AppType", bound="FastAPI")
@@ -990,6 +993,54 @@ class Item(BaseModel):
         self.middleware_stack: Union[ASGIApp, None] = None
         self.setup()
 
+    def build_middleware_stack(self) -> ASGIApp:
+        # Duplicate/override from Starlette to add AsyncExitStackMiddleware
+        # inside of ExceptionMiddleware, inside of custom user middlewares
+        debug = self.debug
+        error_handler = None
+        exception_handlers: dict[Any, ExceptionHandler] = {}
+
+        for key, value in self.exception_handlers.items():
+            if key in (500, Exception):
+                error_handler = value
+            else:
+                exception_handlers[key] = value
+
+        middleware = (
+            [Middleware(ServerErrorMiddleware, handler=error_handler, debug=debug)]
+            + self.user_middleware
+            + [
+                Middleware(
+                    ExceptionMiddleware, handlers=exception_handlers, debug=debug
+                ),
+                # Add FastAPI-specific AsyncExitStackMiddleware for closing files.
+                # Before this was also used for closing dependencies with yield but
+                # those now have their own AsyncExitStack, to properly support
+                # streaming responses while keeping compatibility with the previous
+                # versions (as of writing 0.117.1) that allowed doing
+                # except HTTPException inside a dependency with yield.
+                # This needs to happen after user middlewares because those create a
+                # new contextvars context copy by using a new AnyIO task group.
+                # This AsyncExitStack preserves the context for contextvars, not
+                # strictly necessary for closing files but it was one of the original
+                # intentions.
+                # If the AsyncExitStack lived outside of the custom middlewares and
+                # contextvars were set, for example in a dependency with 'yield'
+                # in that internal contextvars context, the values would not be
+                # available in the outer context of the AsyncExitStack.
+                # By placing the middleware and the AsyncExitStack here, inside all
+                # user middlewares, the same context is used.
+                # This is currently not needed, only for closing files, but used to be
+                # important when dependencies with yield were closed here.
+                Middleware(AsyncExitStackMiddleware),
+            ]
+        )
+
+        app = self.router
+        for cls, args, kwargs in reversed(middleware):
+            app = cls(app, *args, **kwargs)
+        return app
+
     def openapi(self) -> Dict[str, Any]:
         """
         Generate the OpenAPI schema of the application. This is called by FastAPI

fastapi/middleware/asyncexitstack.py

@@ -0,0 +1,18 @@
+from contextlib import AsyncExitStack
+
+from starlette.types import ASGIApp, Receive, Scope, Send
+
+
+# Used mainly to close files after the request is done, dependencies are closed
+# in their own AsyncExitStack
+class AsyncExitStackMiddleware:
+    def __init__(
+        self, app: ASGIApp, context_name: str = "fastapi_middleware_astack"
+    ) -> None:
+        self.app = app
+        self.context_name = context_name
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        async with AsyncExitStack() as stack:
+            scope[self.context_name] = stack
+            await self.app(scope, receive, send)