Starting the slow migration to FastAPI...

app/routers/websockets.py

@@ -2,6 +2,11 @@ from fastapi import APIRouter, FastAPI, WebSocket, WebSocketDisconnect
 from fastapi.responses import HTMLResponse
 from typing import List
 
+from ..lib_general import *
+from ..log import *
+from app.config import settings
+from app.db import *
+
 router = APIRouter()
 
 
@@ -23,7 +28,9 @@ html = """
         <script>
             var client_id = Date.now()
             document.querySelector("#ws-id").textContent = client_id;
-            var ws = new WebSocket(`ws://localhost:5005/ws/${client_id}`);
+            //var ws = new WebSocket(`ws://localhost:5005/ws/${client_id}`);
+            var ws = new WebSocket("ws://localhost:8000/ws_redis");
+            //var ws = new WebSocket("ws://fastapi.localhost/ws_redis");
             ws.onmessage = function(event) {
                 var messages = document.getElementById('messages')
                 var message = document.createElement('li')
@@ -43,41 +50,64 @@ html = """
 """
 
 
-class ConnectionManager:
-    def __init__(self):
-        self.active_connections: List[WebSocket] = []
-
-    async def connect(self, websocket: WebSocket):
-        await websocket.accept()
-        self.active_connections.append(websocket)
-
-    def disconnect(self, websocket: WebSocket):
-        self.active_connections.remove(websocket)
-
-    async def send_personal_message(self, message: str, websocket: WebSocket):
-        await websocket.send_text(message)
-
-    async def broadcast(self, message: str):
-        for connection in self.active_connections:
-            await connection.send_text(message)
-
-
-manager = ConnectionManager()
-
-
 @router.get("/ws_test")
-async def websocket_root():
+async def get():
+    log.setLevel(logging.DEBUG)
+    log.debug(locals())
     return HTMLResponse(html)
 
 
-@router.websocket("/ws/{client_id}")
-async def websocket_endpoint(websocket: WebSocket, client_id: int):
-    await manager.connect(websocket)
-    try:
-        while True:
-            data = await websocket.receive_text()
-            await manager.send_personal_message(f"You wrote: {data}", websocket)
-            await manager.broadcast(f"Client #{client_id} says: {data}")
-    except WebSocketDisconnect:
-        manager.disconnect(websocket)
-        await manager.broadcast(f"Client #{client_id} left the chat")
+@router.websocket("/ws")
+async def websocket_endpoint(websocket: WebSocket):
+    log.setLevel(logging.DEBUG)
+    log.debug(locals())
+    await websocket.accept()
+    await redis_connector(websocket)
+
+
+async def redis_connector(
+    websocket: WebSocket, redis_uri: str = "redis://localhost:6379"
+):
+    log.setLevel(logging.DEBUG)
+    log.debug(locals())
+    async def consumer_handler(ws: WebSocket, r):
+        try:
+            while True:
+                message = await ws.receive_text()
+                if message:
+                    #logging.info(ws)
+                    #logging.info(dir(message))
+                    data = json.loads(message)
+                    #await r.publish("chat:c", message)
+                    #await r.publish("chat:c", str(data['message']))
+                    await r.publish("chat:c", str(data['client_id']))
+                    await r.publish("chat:c", str(data))
+        except WebSocketDisconnect as exc:
+            # TODO this needs handling better
+            logger.error(exc)
+
+    async def producer_handler(r, ws: WebSocket):
+        (channel,) = await r.subscribe("chat:c")
+        assert isinstance(channel, aioredis.Channel)
+        try:
+            while True:
+                message = await channel.get()
+                if message:
+                    await ws.send_text(message.decode("utf-8"))
+        except Exception as exc:
+            # TODO this needs handling better
+            logger.error(exc)
+
+    redis = await aioredis.create_redis_pool(redis_uri)
+
+    consumer_task = consumer_handler(websocket, redis)
+    producer_task = producer_handler(redis, websocket)
+    done, pending = await asyncio.wait(
+        [consumer_task, producer_task], return_when=asyncio.FIRST_COMPLETED,
+    )
+    logger.debug(f"Done task: {done}")
+    for task in pending:
+        logger.debug(f"Canceling task: {task}")
+        task.cancel()
+    redis.close()
+    await redis.wait_closed()