Redis channel layer about nejma HOT 2 OPEN

import asyncio
import json
import logging

import aioredis
from aioredis import Channel
from aioredis.pubsub import Receiver
from starlette.endpoints import WebSocketEndpoint as BaseWebSocketEndpoint, status
from starlette.websockets import WebSocketState, WebSocket

logger = logging.getLogger("websockets")

class RedisLayer(ChannelLayer):
    def __init__(self, redis_host, default_redis_channel_path=None):
        self.redis_host = redis_host
        self.initialized = False
        self.default_path = default_redis_channel_path or "default_path"

    async def _initialize(self):
        self.pub = await aioredis.create_redis(self.redis_host)
        self.sub = await aioredis.create_redis(self.redis_host)
        self.mpsc = Receiver(loop=asyncio.get_event_loop())
        self.initialized = True

    async def publish_to_redis(self, msg, path=None):
        if not self.initialized:
            await self._initialize()
        _path = path or self.default_path
        await self.pub.execute("publish", _path, json.dumps(msg))

    async def subscribe_to_redis(self, websocket, path=None, receive_callback=None):
        _path = path or self.default_path
        channel = Channel(_path, is_pattern=False)
        if not self.initialized:
            await self._initialize()
        try:
            result = await self.sub.subscribe(self.mpsc.channel(_path))
            async for channel, msg in self.mpsc.iter():
                if websocket.client_state == WebSocketState.CONNECTED:
                    msg = json.loads(msg)
                    if receive_callback:
                        await receive_callback(websocket, msg)
                    else:
                        await websocket.send_json(msg)
        except:
            import traceback

            traceback.print_exc()
            await self.close_connections(_path)
        finally:
            logger.info("Connection closed")

    async def close_connections(self, channel, *channels):
        await self.sub.unsubscribe(channel, *channels)
        self.mpsc.stop()
        self.pub.close()
        self.sub.close()

class RedisWebSocketEndPoint(BaseWebSocketEndpoint):
    encoding = "json"

    async def dispatch(self) -> None:
        websocket = WebSocket(self.scope, receive=self.receive, send=self.send)
        await self.on_connect(websocket)
        await asyncio.gather(self._dispatch(websocket), self._redis_setup(websocket))

    async def _dispatch(self, websocket) -> None:
        close_code = status.WS_1000_NORMAL_CLOSURE

        try:
            while True:
                message = await websocket.receive()
                if message["type"] == "websocket.receive":
                    data = await self.decode(websocket, message)
                    await self.on_receive(websocket, data)
                elif message["type"] == "websocket.disconnect":
                    close_code = int(message.get("code", status.WS_1000_NORMAL_CLOSURE))
                    break
        except Exception as exc:
            close_code = status.WS_1011_INTERNAL_ERROR
            raise exc from None
        finally:
            await self.on_disconnect(websocket, close_code)

    async def _redis_setup(self, websocket) -> None:
        self.channel_layer = RedisLayer(REDIS_PATH) # could be read from the app scope.
        await self.channel_layer.subscribe_to_redis(
            websocket, receive_callback=self.on_redis_receive
        )

    async def on_redis_receive(self, websocket, data):
        # can be overidden by the user for preventing access to specific clients since all clients are subscribed to the channel.
        
        await websocket.send_json(data)

# Any client connected to this endpoint would automatically receive whatever data has been published.
@app.websocket_route("/redis")
class TestRedis(RedisWebSocketEndPoint):
    async def on_receive(self, ws, data):
        if data.pop("from_server", None):
            await self.channel_layer.publish_to_redis(data)