elstat/blueprints/streaming.py

50 lines
1.1 KiB
Python

import json
import uuid
import logging
import asyncio
import websockets
from sanic import Blueprint
bp = Blueprint(__name__)
log = logging.getLogger(__name__)
class OP:
SUBSCRIBED = 1
class ErrorCodes:
TOO_MUCH = 4420
@bp.websocket('/api/streaming')
async def streaming_ws(request, ws):
opening_msg = await ws.recv()
if len(opening_msg) > 256:
await ws.close(code=ErrorCodes.TOO_MUCH, reason='Too much data')
return
open_payload = json.loads(opening_msg)
# subscribe the websocket to all channels it wants
ws.client_id = uuid.uuid4()
try:
channels = open_payload['channels']
subscribed = request.app.manager.subscribe(channels, ws)
await ws.send(json.dumps({
'op': OP.SUBSCRIBED,
'channels': subscribed,
}))
# keep websocket alive
while True:
await ws.ping()
await asyncio.sleep(1)
except websockets.exceptions.ConnectionClosed as wc:
log.warning(f'connection {ws.client_id} closed: {wc!r}')
request.app.manager.unsub_all(ws)