-
-
Save timhughes/313c89a0d587a25506e204573c8017e4 to your computer and use it in GitHub Desktop.
| """ | |
| Usage: | |
| Make sure that redis is running on localhost (or adjust the url) | |
| Install uvicorn or some other asgi server https://asgi.readthedocs.io/en/latest/implementations.html | |
| pip install -u uvicorn | |
| Install dependencies | |
| pip install -u aioredis fastapi | |
| Start the application, this will depend on the asgi server | |
| uvicorn fastapi_websocket_redis_pubsub:app | |
| Open two browser windows to the web interface http://127.0.0.1:8000 | |
| Enter some data in one window and it should appear in the other window. | |
| """ | |
| import asyncio | |
| import logging | |
| from fastapi import FastAPI | |
| from fastapi.responses import HTMLResponse | |
| from fastapi.websockets import WebSocket, WebSocketDisconnect | |
| from aioredis | |
| 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 redis_connector(websocket) | |
| async def redis_connector( | |
| websocket: WebSocket, redis_uri: str = "redis://localhost:6379" | |
| ): | |
| async def consumer_handler(ws: WebSocket, r): | |
| try: | |
| while True: | |
| message = await ws.receive_text() | |
| if message: | |
| await r.publish("chat:c", message) | |
| 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() |
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 👍