Support for CorrelationIdPlugin and CorrelationIdPlugin still executing when application encounters an exception

[derekbekoe]

Issue

When a route raises an exception and so returns a 500 Internal Server Error, x-request-id and x-correlation-id are not set.

The culprit seems to be

starlette-context/starlette_context/middleware.py

Lines 46 to 52 in 79aa8ff

	try:
	response = await call_next(request)
	for plugin in self.plugins:
	await plugin.enrich_response(response)
	finally:
	_request_scope_context_storage.reset(_starlette_context_token)

This is similar to this comment - tiangolo/fastapi#397 (comment)

Expectation
I'd expect those headers to be set for all responses. One benefit of correlation ids is when there are errors, we can use the ids to track down the issue.

Steps to reproduce
Sample.
The "/" endpoint returns the x-request-id and x-correlation-id headers.
The "/error" does not.

from starlette.applications import Starlette
from starlette.middleware import Middleware
from starlette.requests import Request
from starlette.responses import JSONResponse

import uvicorn
from starlette_context import context, plugins
from starlette_context.middleware import ContextMiddleware

middleware = [
    Middleware(
        ContextMiddleware,
        plugins=(plugins.RequestIdPlugin(), plugins.CorrelationIdPlugin()),
    )
]

app = Starlette(debug=True, middleware=middleware)


@app.route("/")
async def index(request: Request):
    return JSONResponse(context.data)

@app.route("/error")
async def index(request: Request):
    raise RuntimeError()
    return JSONResponse(context.data)

uvicorn.run(app, host="0.0.0.0")

[tomwojcik]

I removed my answers twice. At first I thought there's nothing to fix, then I've noticed what the issue was, but now I still think there's nothing to fix.

Originally I was confused because you wrote

I'd expect those headers to be set for all requests.

but I think you meant responses. So let's dig into it.

ExceptionMiddleware is added automatically as the last one, so the first one to be run on response and so response headers will be set after that.

But you ask about 500, which is handled by ServerErrorMiddleware.

ServerErrorMiddleware - Ensures that application exceptions may return a custom 500 page, or display an application traceback in DEBUG mode. This is always the outermost middleware layer.

As it's outermost there's nothing I can do about it.

So if you want to set response headers for 500 or Exception, you need to add another middleware only to handle these two or override build_middleware_stack method from Starlette.

I think the best solution for that problem would be to add another middleware which only does try catch for 500 and Exception.

Anyway, I will write some tests and examples for that over the weekend.

class Starlette:
    def __init__(
        self,
        debug: bool = False,
        #  ... more things ....
        middleware: typing.Sequence[Middleware] = None,
        exception_handlers: typing.Dict[
            typing.Union[int, typing.Type[Exception]], typing.Callable
        ] = None,
    ) -> None:
        self.exception_handlers = (
            {} if exception_handlers is None else dict(exception_handlers)
        )
        self.user_middleware = [] if middleware is None else list(middleware)
        self.middleware_stack = self.build_middleware_stack()

    def build_middleware_stack(self) -> ASGIApp:
        debug = self.debug
        error_handler = None
        exception_handlers = {}

        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,
                )
            ]

[derekbekoe]

Yes the headers x-request-id and x-correlation-id being set in responses was what I meant. Updated the initial issue comment to fix the typo.

[tomwojcik]

I was digging into it for some time and I don't think there's a clean solution for that.

I don't think you can do much by adding additional middleware. Subclassing Starlette and creating your own build_middleware_stack is a terrible idea.

For exceptions that are not 500 (basically those that are handled), you need to assign headers in exception function like in here.

starlette-context/tests/test_integration/test_server_error_middleware.py

Lines 23 to 24 in e8316a1

	async def custom_exception_handler(request: Request, exc: Exception):
	return JSONResponse({"exception": "handled"}, headers=context.data)

So you can handle that RuntimeError quite easily but there's no solution for all exceptions.

For 500 or anything that isn't handled by your handler_functions, it's impossible in the current state of Starlette.

Please consider closing this ticket or opening a PR. I will leave this ticket open as I don't know if you have any ideas. I run out of them.

[derekbekoe]

I will close this issue. Thanks for the explanation.

[tomwojcik]

@derekbekoe this case is now supported. https://github.com/tomwojcik/starlette-context/pull/79/files