-
-
Save timhughes/313c89a0d587a25506e204573c8017e4 to your computer and use it in GitHub Desktop.
| import asyncio | |
| import logging | |
| from fastapi import FastAPI | |
| from fastapi.responses import HTMLResponse | |
| from fastapi.websockets import WebSocket, WebSocketDisconnect | |
| from aioredis import create_connection, Channel | |
| logging.basicConfig(level=logging.INFO) | |
| logger = logging.getLogger(__name__) | |
| app = FastAPI() | |
| html = """ | |
| <!DOCTYPE html> | |
| <html> | |
| <head> | |
| <title>Chat</title> | |
| </head> | |
| <body> | |
| <h1>WebSocket Chat</h1> | |
| <form action="" onsubmit="sendMessage(event)"> | |
| <input type="text" id="messageText" autocomplete="off"/> | |
| <button>Send</button> | |
| </form> | |
| <ul id='messages'> | |
| </ul> | |
| <script> | |
| var ws = new WebSocket("ws://localhost:8000/ws"); | |
| ws.onmessage = function(event) { | |
| var messages = document.getElementById('messages') | |
| var message = document.createElement('li') | |
| var content = document.createTextNode(event.data) | |
| message.appendChild(content) | |
| messages.appendChild(message) | |
| }; | |
| function sendMessage(event) { | |
| var input = document.getElementById("messageText") | |
| ws.send(input.value) | |
| input.value = '' | |
| event.preventDefault() | |
| } | |
| </script> | |
| </body> | |
| </html> | |
| """ | |
| @app.get("/") | |
| async def get(): | |
| return HTMLResponse(html) | |
| @app.websocket("/ws") | |
| async def websocket_endpoint(websocket: WebSocket): | |
| await websocket.accept() | |
| await asyncio.gather(handle_websocket(websocket), handle_redis(websocket)) | |
| async def handle_redis(websocket): | |
| redis_conn = await create_connection(("localhost", 6379)) | |
| channel = Channel("chat", is_pattern=False) | |
| await redis_conn.execute_pubsub("subscribe", channel) | |
| try: | |
| while True: | |
| message = await channel.get() | |
| if message: | |
| await websocket.send_text(message.decode("utf-8")) | |
| except Exception as exc: | |
| logger.error(exc) | |
| async def handle_websocket(websocket): | |
| redis_conn = await create_connection(("localhost", 6379)) | |
| try: | |
| while True: | |
| message = await websocket.receive_text() | |
| if message: | |
| await redis_conn.execute("publish", "chat", message) | |
| except WebSocketDisconnect as exc: | |
| logger.error(exc) |
Thank you very much! Only needed to change one thing. When installing aioredis I had to specify aioredis==1.3.1 since aioredis 2.0 has since been released and has breaking changes with the code you provided.
@ericrdgz what was the change and i will update the gist
@timhughes Hello! Thanks for your solution! I tried to adapt your code for aioredis 2.0. Messages sent and receive successful. Please, if you know how you can improve this, then do it 👍
package versions
aioredis==2.0.0
websockets==10.1
fastapi==0.70.0
main.py
import asyncio
import logging
import aioredis
from aioredis.client import PubSub, Redis
from fastapi import FastAPI
from fastapi.websockets import WebSocket, WebSocketDisconnect
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
app = FastAPI()
@app.websocket('/ws')
async def ws_voting_endpoint(websocket: WebSocket):
await websocket.accept()
await redis_connector(websocket)
async def redis_connector(websocket: WebSocket):
async def consumer_handler(conn: Redis, ws: WebSocket):
try:
while True:
message = await ws.receive_text()
if message:
await conn.publish("chat:c", message)
except WebSocketDisconnect as exc:
# TODO this needs handling better
logger.error(exc)
async def producer_handler(pubsub: PubSub, ws: WebSocket):
await pubsub.subscribe("chat:c")
# assert isinstance(channel, PubSub)
try:
while True:
message = await pubsub.get_message(ignore_subscribe_messages=True)
if message:
await ws.send_text(message.get('data'))
except Exception as exc:
# TODO this needs handling better
logger.error(exc)
conn = await get_redis_pool()
pubsub = conn.pubsub()
consumer_task = consumer_handler(conn=conn, ws=websocket)
producer_task = producer_handler(pubsub=pubsub, ws=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()
async def get_redis_pool():
return await aioredis.from_url(f'redis://your_redis_host', encoding="utf-8", decode_responses=True)
How can I use this code then to send data over the socket to specific clients?
@FaisalJulaidan have a look at aioredis pubsub documentation. You want to have each client subscribe to it's own personal channel. then other clients can publish to that channel
So I'm trying to use aioredis through redis-py. When I implement the pattern you descrbie with python 3.11.0, the application doesn't manage to shut down.
Is this a problem introduced with python 3.11.0, or has this always been a problem?
@wholmen it was a long time ago and I don't remember. These days I would advise using https://github.com/encode/broadcaster instead of this gist and contribute any fixes to that project. In the example there they have an on_shutdown callback which closes the connection. Maybe that is what my code is missing
Thank you soo much..! really helpful 👍