Browse Source

Merge background tasks from injected BackgroundTasks and Response

When an endpoint injects BackgroundTasks and also returns a Response with
its own background task, the injected tasks were silently discarded. Now
both sets of tasks are merged so all background work executes as expected.

Closes #11215
pull/15010/head
Anandesh Sharma 5 months ago
parent
commit
16ac260aaa
  1. 14
      fastapi/routing.py
  2. 71
      tests/test_background_tasks_merge.py

14
fastapi/routing.py

@ -61,6 +61,7 @@ from fastapi.utils import (
is_body_allowed_for_status_code,
)
from starlette import routing
from starlette.background import BackgroundTasks as StarletteBackgroundTasks
from starlette._exception_handler import wrap_app_handling_exceptions
from starlette._utils import is_async_callable
from starlette.concurrency import run_in_threadpool
@ -436,6 +437,19 @@ def get_request_handler(
if isinstance(raw_response, Response):
if raw_response.background is None:
raw_response.background = solved_result.background_tasks
elif solved_result.background_tasks is not None:
# Merge injected background tasks with the response's
# existing background task(s) so neither set is lost
existing = raw_response.background
merged = StarletteBackgroundTasks()
if isinstance(existing, StarletteBackgroundTasks):
merged.tasks.extend(existing.tasks)
else:
merged.tasks.append(existing)
merged.tasks.extend(
solved_result.background_tasks.tasks
)
raw_response.background = merged
response = raw_response
else:
response_args: dict[str, Any] = {

71
tests/test_background_tasks_merge.py

@ -0,0 +1,71 @@
from fastapi import BackgroundTasks, FastAPI
from fastapi.testclient import TestClient
from starlette.background import BackgroundTask
from starlette.responses import JSONResponse
results: list[str] = []
def task_from_injection():
results.append("injected")
def task_from_response():
results.append("response")
def test_merge_injected_and_response_background_tasks():
"""Both injected BackgroundTasks and Response background tasks should run."""
app = FastAPI()
@app.get("/")
async def endpoint(background_tasks: BackgroundTasks):
background_tasks.add_task(task_from_injection)
return JSONResponse(
{"status": "ok"},
background=BackgroundTask(task_from_response),
)
results.clear()
client = TestClient(app)
response = client.get("/")
assert response.status_code == 200
assert "injected" in results
assert "response" in results
def test_only_injected_background_tasks():
"""When Response has no background, injected tasks still run."""
app = FastAPI()
@app.get("/")
async def endpoint(background_tasks: BackgroundTasks):
background_tasks.add_task(task_from_injection)
return JSONResponse({"status": "ok"})
results.clear()
client = TestClient(app)
response = client.get("/")
assert response.status_code == 200
assert "injected" in results
assert "response" not in results
def test_only_response_background_task():
"""When no BackgroundTasks are injected, response background task runs."""
app = FastAPI()
@app.get("/")
async def endpoint():
return JSONResponse(
{"status": "ok"},
background=BackgroundTask(task_from_response),
)
results.clear()
client = TestClient(app)
response = client.get("/")
assert response.status_code == 200
assert "response" in results
assert "injected" not in results
Loading…
Cancel
Save