tiangolo.fastapi/fastapi/lifespan.py

"""Unified lifespan management for FastAPI."""

import warnings
from collections.abc import AsyncGenerator, Callable, Sequence
from contextlib import asynccontextmanager
from typing import Any, Optional, TypeVar

from starlette.types import ASGIApp, Lifespan

from fastapi.exceptions import FastAPIDeprecationWarning

AppType = TypeVar("AppType", bound="ASGIApp")


class LifespanError(Exception):
    """Exception raised when lifespan handling fails."""

    def __init__(
        self,
        message: str,
        *,
        phase: str,
        handler: Optional[Callable[..., Any]] = None,
        original_error: Optional[BaseException] = None,
    ):
        super().__init__(message)
        self.phase = phase
        self.handler = handler
        self.original_error = original_error

    def __str__(self) -> str:
        base = f"Lifespan {self.phase} failed"
        if self.handler:
            handler_name = getattr(self.handler, "__name__", str(self.handler))
            base += f" in handler '{handler_name}'"
        if self.original_error:
            base += f": {self.original_error}"
        return base


def _emit_on_event_deprecation_warning(event_type: str) -> None:
    warnings.warn(
        f"on_event('{event_type}') is deprecated. Use the lifespan context manager.",
        FastAPIDeprecationWarning,
        stacklevel=4,
    )


class UnifiedLifespanManager:
    """Manages unified lifespan handling with rollback semantics."""

    def __init__(
        self,
        *,
        on_startup: Optional[Sequence[Callable[..., Any]]] = None,
        on_shutdown: Optional[Sequence[Callable[..., Any]]] = None,
        lifespan: Optional[Lifespan[AppType]] = None,
        emit_deprecation_warnings: bool = True,
    ):
        self._on_startup = list(on_startup or [])
        self._on_shutdown = list(on_shutdown or [])
        self._lifespan = lifespan
        self._emit_warnings = emit_deprecation_warnings
        self._completed_startup_handlers: list[Callable[..., Any]] = []

        if lifespan and (on_startup or on_shutdown):
            warnings.warn(
                "Using both 'lifespan' and 'on_startup'/'on_shutdown' is not recommended.",
                FastAPIDeprecationWarning,
                stacklevel=2,
            )

    @asynccontextmanager
    async def lifespan(self, app: AppType) -> AsyncGenerator[dict[str, Any], None]:
        """Unified lifespan context manager with rollback semantics."""
        state: dict[str, Any] = {}
        self._completed_startup_handlers = []
        lifespan_cm = None

        try:
            await self._execute_startup_handlers()

            if self._lifespan:
                lifespan_cm = self._lifespan(app)
                lifespan_state = await lifespan_cm.__aenter__()
                if lifespan_state:
                    state.update(lifespan_state)

            try:
                yield state
            finally:
                shutdown_errors: list[tuple[str, BaseException]] = []

                if lifespan_cm is not None:
                    try:
                        await lifespan_cm.__aexit__(None, None, None)
                    except Exception as e:
                        shutdown_errors.append(("lifespan context", e))

                handler_errors = await self._execute_shutdown_handlers()
                shutdown_errors.extend(handler_errors)

                if shutdown_errors:
                    error_msgs = [f"{name}: {err}" for name, err in shutdown_errors]
                    warnings.warn(
                        f"Errors during shutdown: {'; '.join(error_msgs)}",
                        RuntimeWarning,
                        stacklevel=2,
                    )

        except BaseException as startup_error:
            await self._rollback_startup(startup_error)
            raise

    async def _execute_startup_handlers(self) -> None:
        for handler in self._on_startup:
            if self._emit_warnings:
                _emit_on_event_deprecation_warning("startup")

            try:
                result = handler()
                if hasattr(result, "__await__"):
                    await result
                self._completed_startup_handlers.append(handler)
            except Exception as e:
                raise LifespanError(
                    f"Startup handler failed",
                    phase="startup",
                    handler=handler,
                    original_error=e,
                ) from e

    async def _execute_shutdown_handlers(self) -> list[tuple[str, BaseException]]:
        errors: list[tuple[str, BaseException]] = []
        for handler in reversed(self._on_shutdown):
            if self._emit_warnings:
                _emit_on_event_deprecation_warning("shutdown")

            try:
                result = handler()
                if hasattr(result, "__await__"):
                    await result
            except Exception as e:
                handler_name = getattr(handler, "__name__", str(handler))
                errors.append((handler_name, e))
        return errors

    async def _rollback_startup(self, original_error: BaseException) -> None:
        if not self._completed_startup_handlers:
            return

        completed_count = len(self._completed_startup_handlers)
        handlers_to_cleanup = list(reversed(self._on_shutdown[:completed_count]))

        for handler in handlers_to_cleanup:
            try:
                result = handler()
                if hasattr(result, "__await__"):
                    await result
            except Exception as cleanup_error:
                handler_name = getattr(handler, "__name__", str(handler))
                warnings.warn(
                    f"Error during rollback in '{handler_name}': {cleanup_error}",
                    RuntimeWarning,
                    stacklevel=2,
                )

    def add_startup_handler(self, handler: Callable[..., Any]) -> None:
        self._on_startup.append(handler)

    def add_shutdown_handler(self, handler: Callable[..., Any]) -> None:
        self._on_shutdown.append(handler)


def create_unified_lifespan(
    *,
    on_startup: Optional[Sequence[Callable[..., Any]]] = None,
    on_shutdown: Optional[Sequence[Callable[..., Any]]] = None,
    lifespan: Optional[Lifespan[AppType]] = None,
) -> Lifespan[AppType]:
    """Create a unified lifespan from mixed inputs."""
    manager = UnifiedLifespanManager(
        on_startup=on_startup,
        on_shutdown=on_shutdown,
        lifespan=lifespan,
    )
    return manager.lifespan  # type: ignore


def merge_lifespan_state(
    parent_state: Optional[dict[str, Any]],
    child_state: Optional[dict[str, Any]],
    *,
    child_overrides_parent: bool = True,
) -> dict[str, Any]:
    """Merge lifespan state dictionaries."""
    parent = parent_state or {}
    child = child_state or {}

    if child_overrides_parent:
        return {**parent, **child}
    else:
        return {**child, **parent}