Ability to execute logic after the response is fully streamed

Right now, the middleware method process_response can be used to post-process responses after they are routed and handled by responders.

However, in case response body is a stream, one may want to execute logic when the stream iterable is exausted. We are already wrapping req.stream in some cases, so we could extend this to wrap .close() with code to execute the relevant method (finalize_response?) from middlewares added to the API.

Use cases could include cleaning up resources that were used to stream the response, such as releasing database connections back to the pool.

[vytas7]

Doh, sorry… My bad 😬 Anyway, finalize_response would very obviously be executed after everything else in the request-response cycle.

[numberoverzero]

+1 I’d love to have a way to run cleanup logic after the stream has been consumed, or after the stream raises, or the connection is closed.

I’m using the following to wrap the response stream and execute a callback when read raises or close is called[0].

import io

class StreamWithCallback(io.BufferedIOBase):
    def __init__(self, wrapped: io.BufferedIOBase, callback=None) -> None:
        self._wrapped = wrapped
        self._callback = callback

    def close(self):
        if self._wrapped:
            try:
                self._wrapped.close()
                self._wrapped = None
            finally:
                if self._callback:
                    self._callback()

    def read(self, n: int) -> bytes:
        try:
            return self._wrapped.read(n)
        except:
            self.close()
            raise

And I’m using it like so:

def on_get(self, req: falcon.Request, resp: falcon.Response, id: str) -> None:
    stream = self.controller.get_read_stream(id)
    def cleanup():
        self.controller.remove_by_id(id)
    resp.status = falcon.HTTP_200
    resp.stream = ReadWithCallback(stream, cleanup)

[0] Most wsgi frameworks should call close on the stream, although I only verified this for gunicorn sync threads.