Server websocket conflicts with middleware stack

[al-dpopowich]

I have a mature aiohttp server with the following middlewares defined:

middlewares = [
   capture_exc,      # runs handler in try/except/else to convert app exceptions to proper json response
   task_collection,  # each request may kick off long-running background tasks; this manages these tasks
   db_trx,           # run request in database transaction (trx); see below
   authentication,   # within the database trx, confirm user credentials / session
]

The db_trx middleware does this:

@web.middleware
async def db_transaction(request, handler):
   """Manage db transaction"""
   with new_dbsession() as sess:
      async with sess:
         async with sess.begin():
            return await handler(request)

For those familiar with it, this is async SQLAlchemy (2.x). The handler is run within the context manager and if it succeeds the transaction is committed. If there's an exception, the trx is rolled back. This works beautifully across my application. Run any SQL inside my requests and I never need worry about committing vs rolling back; it's handled automatically.

Now, my question regarding websockets...

We want to add websockets to our app to handle "live" communication with the clients. Think chat or dashboard updates.

The stock websocket implementation as described in the Quickstart guide doesn't return. (Well, not until the connection is closed.) This is in conflict with my middleware stack which is designed for per-request common tasks. My websocket connection needs to "skip" the middleware stack and do its own version of it per-message. For example, I may want/need this:

async def websocket_handler(request):

    ws = web.WebSocketResponse()
    await ws.prepare(request)

    async for msg in ws:
        with new_dbsession() as sess:
            async with sess:
                async with sess.begin():
                   # handle each message in a db trx

All of which would be in conflict with my common application middlewares.

Surely, per-request middleware not being compatible with long-lived websocket connections can't be a new discovery, so what I think I'm asking is: does anyone have a recipe for skipping the middlewares for a given route?

All I've come up with seems un-aiohttp-ic:

@middleware
async def ws(request, handler):
   """Skip the middlware stack if creating websocket"""
   if request.path == '/ws':
      return await websocket_handler(request)  # my websocket handler
   else:
      return await handler(request)  # my middleware stack

middelwares = [
   ws,   # new middleware on top of the stack
   ...,  # the rest of my middleware
]

Essentially, it requires one to munge the stack within the stack. What would be ideal is a way to customize the middlewares stack based on the route, with a default middlewares stack if there's no match against the custom definitions.

Thanks in advance!

[Dreamsorcerer]

Middlewares are just functions that call the next function in the list. So, the simple approach would be to do:

async def db_middleware(...):
    if request.path == ...:
        return await handler(...)

    async with db():
          ...

You could easily make that into some kind of decorator or similar:

def skip_ws(middleware):
    async def wrapped(request, handler):
        if request.path == ...:
            return await handler(request)
        return await middleware(request, handler)
    return wrapped

@skip_ws
async def db_session(...):
    ...

[al-dpopowich]

Right. This is essentially my work-around as described in my original message. I just wanted to make sure I wasn't missing some standard aiohttp way of doing this.

IMHO, it's rather ugly, having to modify the stack of middleware within the definitions of middleware. It can become gnarly fast. I'd love to see a way to customize the stack per request based on the path.

[Dreamsorcerer]

Right. This is essentially my work-around as described in my original message. I just wanted to make sure I wasn't missing some standard aiohttp way of doing this.

Right, I didn't realise because you missed the handler argument, which needs to be passed in to websocket_handler(), then it should work the same.

IMHO, it's rather ugly, having to modify the stack of middleware within the definitions of middleware. It can become gnarly fast. I'd love to see a way to customize the stack per request based on the path.

I disagree on that. You are asking for something that dynamically changes the stack of middlewares, which would be complicated and unclear what functions are being called each time. Instead, you have a middleware which simply calls a custom function or doesn't. It doesn't modify the stack of middlewares, it is a middleware which runs some code on a given condition.

If you add a yield to the db code instead of calling the handler, then you could even simplify it to something like:

async def middleware(request, handler):
    g = contextlib.nullcontext() if request.path == "/ws" else db_handler(request)
    async with g:
        return await handler(request)