from app.db.redis import get_redis_database
from app.services.kline_service import KLineService
from fastapi import APIRouter
from loguru import logger
from starlette.endpoints import WebSocketEndpoint
from starlette.types import Scope, Receive, Send
router = APIRouter()
async def init():
redis_client = await get_redis_database()
await KLineService.init_kline(redis_client, [60, 3600, 86400])
router.add_event_handler("startup", init)
@router.websocket_route("/subscribe/market/{market}/interval/{interval}")
class WebsocketConsumer(WebSocketEndpoint):
def __init__(self, scope: Scope, receive: Receive, send: Send):
super().__init__(scope, receive, send)
loop = asyncio.get_event_loop()
bootstrap_server = settings.KAFKA_INTERNAL_HOST_PORT
market = self._get_market(scope["path"])
interval = self._get_interval(scope["path"])
topic = f"^kline-{market}.*-{interval}.*-latest$"
logger.info(f"subscribing to topic:: {topic}")
self.kafka_consumer = create_kafka_consumer(loop, f"{topic}_consumer",
bootstrap_server)
self.kafka_consumer.subscribe(pattern=topic)
async def shut():
global kafka_producer
await close_kafka_producer(kafka_producer)
def get_kafka_producer() -> AIOKafkaProducer:
loop = asyncio.get_event_loop()
client_id = "balanced_kafka_producer"
bootstrap_server = settings.KAFKA_INTERNAL_HOST_PORT
producer = create_kafka_producer(loop=loop,
client_id=client_id,
bootstrap_server=bootstrap_server)
return producer
router.add_event_handler("startup", init)
router.add_event_handler("shutdown", shut)
@router.post("/event")
async def event(
_event_or_trade: Union[models.EventLog, models.TradeLog],
mongodb_client: AsyncIOMotorClient = Depends(get_mongodb_database),
redis_client: Redis = Depends(get_redis_database)):
# save raw input in mongodb
# logger.info(f"mongodb_client db => {mongodb_client}")
# is_saved = await TradeService.save_raw_event(redis_client, _event_or_trade)
# logger.info(f"TradeService.save_raw_event - is_saved ? {is_saved}")
# 1. update "depth", then send new depth value to kafka topic
# 2. update "kline", then send new kline(add 1min) to kafka topic