Files
aitbc/apps/blockchain-node/src/aitbc_chain/rpc/websocket.py

35 lines
909 B
Python

from __future__ import annotations
import asyncio
from typing import AsyncIterator, Dict
from fastapi import APIRouter, WebSocket, WebSocketDisconnect
from ..gossip import gossip_broker
router = APIRouter(prefix="/ws", tags=["ws"])
async def _stream_topic(topic: str, websocket: WebSocket) -> None:
subscription = await gossip_broker.subscribe(topic)
try:
while True:
message = await subscription.get()
await websocket.send_json(message)
except WebSocketDisconnect:
pass
finally:
subscription.close()
@router.websocket("/blocks")
async def blocks_stream(websocket: WebSocket) -> None:
await websocket.accept()
await _stream_topic("blocks", websocket)
@router.websocket("/transactions")
async def transactions_stream(websocket: WebSocket) -> None:
await websocket.accept()
await _stream_topic("transactions", websocket)