diff --git a/lite_bootstrap/bootstrappers/fastapi_bootstrapper.py b/lite_bootstrap/bootstrappers/fastapi_bootstrapper.py index 24e352f..8f706ba 100644 --- a/lite_bootstrap/bootstrappers/fastapi_bootstrapper.py +++ b/lite_bootstrap/bootstrappers/fastapi_bootstrapper.py @@ -35,7 +35,6 @@ if import_checker.is_fastapi_installed: import fastapi from fastapi.middleware.cors import CORSMiddleware - from fastapi.routing import _merge_lifespan_context if import_checker.is_fastapi_installed and import_checker.is_sentry_installed: from sentry_sdk.integrations.fastapi import FastApiIntegration @@ -299,11 +298,14 @@ def __init__(self, bootstrap_config: FastAPIConfig) -> None: self._attach_teardown_once(application, lambda: self._wrap_lifespan(application)) def _wrap_lifespan(self, application: "fastapi.FastAPI") -> None: - old_lifespan_manager = application.router.lifespan_context - application.router.lifespan_context = _merge_lifespan_context( - old_lifespan_manager, - self.lifespan_manager, - ) + user_lifespan = application.router.lifespan_context + + @contextlib.asynccontextmanager + async def merged_lifespan(app: "fastapi.FastAPI") -> AsyncGenerator[Mapping[str, typing.Any]]: + async with self.lifespan_manager(app) as bootstrap_state, user_lifespan(app) as user_state: + yield {**bootstrap_state, **(user_state or {})} + + application.router.lifespan_context = merged_lifespan def is_ready(self) -> bool: return import_checker.is_fastapi_installed diff --git a/pyproject.toml b/pyproject.toml index 59fe217..d2f42f7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -82,8 +82,6 @@ free-all = [ "lite-bootstrap[sentry,otl,logging,pyroscope]", ] fastapi = [ - # >=0.112.2 for fastapi.routing._merge_lifespan_context, which FastAPIBootstrapper._wrap_lifespan - # uses to chain teardown onto the user's lifespan. Private, so it can move without a deprecation. "fastapi>=0.112.2", # fastapi 0.112.2 requires typing-extensions>=4.8. "typing-extensions>=4.8", diff --git a/tests/test_fastapi_bootstrap.py b/tests/test_fastapi_bootstrap.py index 5aa3b55..c9e647a 100644 --- a/tests/test_fastapi_bootstrap.py +++ b/tests/test_fastapi_bootstrap.py @@ -5,7 +5,7 @@ import typing import uuid import warnings -from collections.abc import Generator +from collections.abc import AsyncGenerator, Generator from unittest.mock import patch import fastapi @@ -64,6 +64,33 @@ def fastapi_config() -> FastAPIConfig: ) +def test_fastapi_exports_user_shutdown_span_and_preserves_lifespan_state() -> None: + @contextlib.asynccontextmanager + async def lifespan(_app: fastapi.FastAPI) -> AsyncGenerator[dict[str, str]]: + yield {"service": "ready"} + with get_tracer_provider().get_tracer(__name__).start_as_current_span("user-shutdown"): + pass + + app = fastapi.FastAPI(lifespan=lifespan) + + @app.get("/state") + async def read_state(request: fastapi.Request) -> str: + return request.state.service + + bootstrapper = FastAPIBootstrapper(bootstrap_config=FastAPIConfig(application=app, opentelemetry_log_traces=True)) + application = bootstrapper.bootstrap() + provider = get_tracer_provider() + assert isinstance(provider, SDKTracerProvider) + exporter = InMemorySpanExporter() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + + with TestClient(application) as client: + assert client.get("/state").json() == "ready" + + assert "user-shutdown" in [span.name for span in exporter.get_finished_spans()] + assert not bootstrapper.is_bootstrapped + + def _instrument_app_kwargs(config: FastAPIConfig) -> dict[str, typing.Any]: """Bootstrap with FastAPIInstrumentor stubbed, and return what it was handed.""" with patch.object(fastapi_bootstrapper, "FastAPIInstrumentor") as mock_instrumentor: