fix(DI): Ensure dependencies are cleaned up when exception occurrs during cleanup (#4148)

Author: provinzkraut

docs/usage/dependency-injection.rst

@@ -153,9 +153,9 @@ and committing otherwise.
 .. admonition:: Best Practice
     :class: tip
 
-    You should always wrap `yield` in a `try`/`finally` block, regardless of whether you
-    want to handle exceptions, to ensure that the cleanup code is run even when exceptions
-    occurred:
+    You should always wrap ``yield`` in a ``try``/``finally`` block, regardless of
+    whether you want to handle exceptions, to ensure that the cleanup code is run even
+    when exceptions occurred:
 
     .. code-block:: python
 
@@ -168,9 +168,19 @@ and committing otherwise.
 
 .. attention::
 
-   Do not re-raise exceptions within the dependency. Exceptions caught within these
-   dependencies will still be handled by the regular mechanisms without an explicit
-   re-raise
+    Do not re-raise exceptions within the dependency. Exceptions caught within these
+    dependencies will still be handled by the regular mechanisms without an explicit
+    re-raise
+
+
+.. important::
+
+    Exceptions raised during the cleanup step of a dependency will be re-raised in an
+    :exc:`ExceptionGroup` (for Python versions < 3.11, the
+    `exceptiongroup <https://github.com/agronholm/exceptiongroup>`_ will be used). This
+    happens after all dependencies have been cleaned up, so exceptions raised during
+    cleanup of one dependencies do not affect the cleanup of other dependencies.
+
 
 
 Dependency keyword arguments

litestar/_kwargs/cleanup.py

@@ -1,8 +1,13 @@
 from __future__ import annotations
 
-from inspect import Traceback, isasyncgen
+import sys
+from contextlib import AbstractAsyncContextManager
+from inspect import isasyncgen
 from typing import TYPE_CHECKING, Any, AsyncGenerator, Awaitable, Callable, Generator
 
+if sys.version_info < (3, 11):
+    from exceptiongroup import ExceptionGroup
+
 from anyio import create_task_group
 
 from litestar.utils import ensure_async_callable
@@ -12,10 +17,12 @@
 
 
 if TYPE_CHECKING:
+    from types import TracebackType
+
     from litestar.types import AnyGenerator
 
 
-class DependencyCleanupGroup:
+class DependencyCleanupGroup(AbstractAsyncContextManager):
     """Wrapper for generator based dependencies.
 
     Simplify cleanup by wrapping :func:`next` / :func:`anext` calls and providing facilities to
@@ -24,8 +31,6 @@ class DependencyCleanupGroup:
     exceptions caught in this manner will be re-raised after they have been thrown in the generators.
     """
 
-    __slots__ = ("_closed", "_generators")
-
     def __init__(self, generators: list[AnyGenerator] | None = None) -> None:
         """Initialize ``DependencyCleanupGroup``.
 
@@ -45,7 +50,7 @@ def add(self, generator: Generator[Any, None, None] | AsyncGenerator[Any, None])
             None
         """
         if self._closed:
-            raise RuntimeError("Cannot call cleanup on a closed DependencyCleanupGroup")
+            raise RuntimeError("Cannot call .add on a closed DependencyCleanupGroup")
         self._generators.append(generator)
 
     @staticmethod
@@ -62,18 +67,25 @@ def wrapped() -> None:
 
         return ensure_async_callable(wrapped)
 
-    async def cleanup(self) -> None:
+    async def close(self, exc: BaseException | None = None) -> None:
+        if self._closed:
+            raise RuntimeError("Cannot call cleanup on a closed DependencyCleanupGroup")
+
+        self._closed = True
+
+        if exc is None:
+            await self._cleanup()
+        else:
+            await self._throw(exc)
+
+    async def _cleanup(self) -> None:
         """Execute cleanup by calling :func:`next` / :func:`anext` on all generators.
 
         If there are multiple generators to be called, they will be executed in a :class:`anyio.TaskGroup`.
 
         Returns:
             None
         """
-        if self._closed:
-            raise RuntimeError("Cannot call cleanup on a closed DependencyCleanupGroup")
-
-        self._closed = True
 
         if not self._generators:
             return
@@ -95,18 +107,18 @@ async def __aexit__(
         self,
         exc_type: type[BaseException] | None,
         exc_val: BaseException | None,
-        exc_tb: Traceback | None,
+        exc_tb: TracebackType | None,
     ) -> None:
         """If an exception was raised within the contextmanager block, throw it into all generators."""
-        if exc_val:
-            await self.throw(exc_val)
+        await self.close(exc_val)
 
-    async def throw(self, exc: BaseException) -> None:
+    async def _throw(self, exc: BaseException) -> None:
         """Throw an exception in all generators sequentially.
 
         Args:
             exc: Exception to throw
         """
+        exceptions = []
         for gen in self._generators:
             try:
                 if isasyncgen(gen):
@@ -115,3 +127,12 @@ async def throw(self, exc: BaseException) -> None:
                     gen.throw(exc)  # type: ignore[union-attr]
             except (StopIteration, StopAsyncIteration):
                 continue
+            except Exception as cleanup_exc:  # noqa: BLE001
+                if cleanup_exc is not exc:
+                    exceptions.append(cleanup_exc)
+
+        if exceptions:
+            raise ExceptionGroup(
+                "Exceptions occurred during cleanup of dependencies",
+                exceptions,
+            ) from exc

litestar/routes/http.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import contextlib
 from itertools import chain
 from typing import TYPE_CHECKING, Any
 
@@ -17,7 +18,6 @@
 
 if TYPE_CHECKING:
     from litestar._kwargs import KwargsModel
-    from litestar._kwargs.cleanup import DependencyCleanupGroup
     from litestar.connection import Request
     from litestar.types import ASGIApp, HTTPScope, Method, Receive, Scope, Send
 
@@ -134,72 +134,54 @@ async def _get_response_for_request(
             scope=scope, request=request, parameter_model=parameter_model, route_handler=route_handler
         )
 
-    async def _call_handler_function(
+    async def _call_handler_function(  # type: ignore[return]
         self, scope: Scope, request: Request, parameter_model: KwargsModel, route_handler: HTTPRouteHandler
-    ) -> ASGIApp:
+    ) -> ASGIApp:  # pyright: ignore[reportGeneralTypeIssues]
         """Call the before request handlers, retrieve any data required for the route handler, and call the route
         handler's ``to_response`` method.
 
         This is wrapped in a try except block - and if an exception is raised,
         it tries to pass it to an appropriate exception handler - if defined.
         """
         response_data: Any = None
-        cleanup_group: DependencyCleanupGroup | None = None
 
         if before_request_handler := route_handler.resolve_before_request():
             response_data = await before_request_handler(request)
 
-        if not response_data:
-            response_data, cleanup_group = await self._get_response_data(
-                route_handler=route_handler, parameter_model=parameter_model, request=request
-            )
+        # create and enter an AsyncExit stack as we may or may not have a
+        # 'DependencyCleanupGroup' to enter and exit
+        stack = contextlib.AsyncExitStack()
 
-        response: ASGIApp = await route_handler.to_response(
-            app=scope["litestar_app"], data=response_data, request=request
-        )
+        # mypy cannot infer that 'stack' never swallows exceptions, therefore it thinks
+        # this method is potentially missing a 'return' statement
+        async with stack:
+            if not response_data:
+                parsed_kwargs: dict[str, Any] = {}
 
-        if cleanup_group:
-            await cleanup_group.cleanup()
+                if parameter_model.has_kwargs and route_handler.signature_model:
+                    try:
+                        kwargs = await parameter_model.to_kwargs(connection=request)
+                    except SerializationException as e:
+                        raise ClientException(str(e)) from e
 
-        return response
+                    if kwargs.get("data") is Empty:
+                        del kwargs["data"]
 
-    @staticmethod
-    async def _get_response_data(
-        route_handler: HTTPRouteHandler, parameter_model: KwargsModel, request: Request
-    ) -> tuple[Any, DependencyCleanupGroup | None]:
-        """Determine what kwargs are required for the given route handler's ``fn`` and calls it."""
-        parsed_kwargs: dict[str, Any] = {}
-        cleanup_group: DependencyCleanupGroup | None = None
-
-        if parameter_model.has_kwargs and route_handler.signature_model:
-            try:
-                kwargs = await parameter_model.to_kwargs(connection=request)
-            except SerializationException as e:
-                raise ClientException(str(e)) from e
-
-            if kwargs.get("data") is Empty:
-                del kwargs["data"]
-
-            if parameter_model.dependency_batches:
-                cleanup_group = await parameter_model.resolve_dependencies(request, kwargs)
-
-            parsed_kwargs = route_handler.signature_model.parse_values_from_connection_kwargs(
-                connection=request, kwargs=kwargs
-            )
+                    if parameter_model.dependency_batches:
+                        cleanup_group = await parameter_model.resolve_dependencies(request, kwargs)
+                        await stack.enter_async_context(cleanup_group)
+
+                    parsed_kwargs = route_handler.signature_model.parse_values_from_connection_kwargs(
+                        connection=request, kwargs=kwargs
+                    )
 
-        if cleanup_group:
-            async with cleanup_group:
-                data = (
+                response_data = (
                     route_handler.fn(**parsed_kwargs)
                     if route_handler.has_sync_callable
                     else await route_handler.fn(**parsed_kwargs)
                 )
-        elif route_handler.has_sync_callable:
-            data = route_handler.fn(**parsed_kwargs)
-        else:
-            data = await route_handler.fn(**parsed_kwargs)
 
-        return data, cleanup_group
+            return await route_handler.to_response(app=scope["litestar_app"], data=response_data, request=request)
 
     @staticmethod
     async def _get_cached_response(request: Request, route_handler: HTTPRouteHandler) -> ASGIApp | None:

litestar/routes/websocket.py

@@ -81,6 +81,5 @@ async def handle(self, scope: WebSocketScope, receive: Receive, send: Send) -> N
         if cleanup_group:
             async with cleanup_group:
                 await self.route_handler.fn(**parsed_kwargs)
-            await cleanup_group.cleanup()
         else:
             await self.route_handler.fn(**parsed_kwargs)

pyproject.toml

@@ -46,6 +46,7 @@ dependencies = [
   "rich>=13.0.0",
   "rich-click",
   "multipart>=1.2.0",
+  "exceptiongroup>=1.2.2; python_version < \"3.11\"",
   # default litestar plugins
   "litestar-htmx>=0.4.0",
 ]
@@ -397,8 +398,6 @@ lint.select = [
   "W",   # pycodestyle - warning
   "YTT", # flake8-2020
 ]
-
-line-length = 120
 lint.ignore = [
   "A003",    # flake8-builtins - class attribute {name} is shadowing a python builtin
   "B010",    # flake8-bugbear - do not call setattr with a constant attribute value
@@ -419,6 +418,8 @@ lint.ignore = [
   "ISC001",  # Ruff formatter incompatible
   "CPY001",  # ruff - copyright notice at the top of the file
 ]
+
+line-length = 120
 src = ["litestar", "tests", "docs/examples"]
 target-version = "py38"