DontPanicO / fastapi-distributed-websocket

A library to implement websocket for distibuted system based on FastAPI.
MIT License
54 stars 10 forks source link

Decouple `Message` objects initialisation from `WebSocketManager` #9

Closed DontPanicO closed 1 year ago

DontPanicO commented 1 year ago

Feature or enhancement

Initialize Message objects in Connection.iter_json instead of WebSocketManager.receive

Pitch

Currently, Message objects are initialized from client data in WebSocketManager.receive after a validation step (which might raises a ValueError). Even if that's totally fine, a better way of doing it is decoupling the message handling from its initialisation. We can make Connection.iter_json yielding Message istances and WebSocketManager.receive accepting a Message instance (either manually instantiated by the user or through iter_json generator).

Example Usage:

# imports and WebSocketManager initialisation
# skippe for brevity

async def app(scope, receive, send):
    websocket = WebSocket(scope, receive, send)
    conn = await manager.new_connection(websocket, 'conn_id')
    async for msg in conn.iter_json():
        await manager.receive(msg)
    manager.remove_connection(conn)

Update

Users might want to not use WebSocketManager.receive but directly calls send, broadcast, etc... With the above implementation, that's start become expensive (users using iter_json, should then reconvert Message to dict). So we should leave iter_json as is and implement an async iterator for the Connection class. This involve a little change in Message.from_client_message, switching from dict.get(k) to dict.pop(k, None) that's surely out of scope for this issue, but it's ok to have it done here. Also message validation should happen in __anext__ so we can either return None to the caller and/or send a message to the client sayng that the data was not formatted as expected.

# ./_connection.py
...
class Connection:
    ...
    def __aiter__(self):
        return self

    async def __anext__(self):
        try:
            data = await self.receive_json()
            validate_incoming_message(data)
            return Message.from_client_message(data=data)
        except ValueError as exc:
            await self.send_json({'err': '{exc}'})  # maybe?
            return None
        except WebSocketDisconnect:
            raise StopAsyncIteration from None
    ...