import asyncio import json import logging from datetime import UTC, datetime from typing import Any import asyncpg from utils.events import BG_EVENTS_CHANNEL from utils.listen import Listener from utils.read import chats as chats_read from utils.read import presence as presence_read logger = logging.getLogger(__name__) QUEUE_MAXSIZE = 256 class Subscriber: def __init__(self, account_id: int, chat_id: int | None) -> None: self.account_id = account_id self.chat_id = chat_id self.queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue(maxsize=QUEUE_MAXSIZE) class EventHub: def __init__(self) -> None: self._subscribers: set[Subscriber] = set() self._pool: asyncpg.Pool | None = None self._listener = Listener("realtime", {BG_EVENTS_CHANNEL: self._on_notify}) self._tasks: set[asyncio.Task] = set() def subscribe(self, account_id: int, chat_id: int | None) -> Subscriber: sub = Subscriber(account_id, chat_id) self._subscribers.add(sub) return sub def unsubscribe(self, sub: Subscriber) -> None: self._subscribers.discard(sub) async def start(self, pool: asyncpg.Pool) -> None: self._pool = pool await self._listener.start() async def stop(self) -> None: for task in self._tasks: task.cancel() await self._listener.stop() def _on_notify(self, payload: str) -> None: task = asyncio.create_task(self._dispatch(payload)) self._tasks.add(task) task.add_done_callback(self._tasks.discard) async def _dispatch(self, payload: str) -> None: try: event = json.loads(payload) except json.JSONDecodeError: return account_id = event.get("account_id") chat_id = event.get("chat_id") account_wide = event.get("kind") == "presence" targets = [ sub for sub in self._subscribers if sub.account_id == account_id and (account_wide or sub.chat_id is None or sub.chat_id == chat_id) ] if not targets: return frame = await self._build_frame(event) if frame is None: return for sub in targets: try: sub.queue.put_nowait(frame) except asyncio.QueueFull: logger.warning("Dropping event for slow subscriber") async def _build_frame( # noqa: PLR0911 self, event: dict[str, Any] ) -> dict[str, Any] | None: if self._pool is None: return None kind = event.get("kind") account_id = event["account_id"] if kind in {"message", "edit", "reaction"}: view = await chats_read.get_message( self._pool, account_id, event["chat_id"], event["message_id"] ) if view is None: return None return {"type": kind, "message": view.model_dump(mode="json")} if kind == "delete": return { "type": "delete", "chat_id": event.get("chat_id"), "message_ids": event.get("message_ids", []), } if kind == "presence": sample = await presence_read.current_presence( self._pool, account_id, event["chat_id"] ) return { "type": "presence", "peer_id": event["chat_id"], "sample": sample.model_dump(mode="json") if sample else None, } if kind == "receipt": return { "type": "receipt", "chat_id": event["chat_id"], "read_up_to": event["message_id"], "read_at": datetime.now(UTC).isoformat(), } return None hub = EventHub()