164 lines
5.3 KiB
Python
164 lines
5.3 KiB
Python
import asyncio
|
|
import json
|
|
import logging
|
|
from datetime import UTC, datetime
|
|
from typing import TYPE_CHECKING, 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
|
|
from utils.read import watches as watches_read
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Callable
|
|
|
|
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,
|
|
watches_read.ALERTS_CHANGED_CHANNEL: self._on_alert,
|
|
},
|
|
)
|
|
self._tasks: set[asyncio.Task] = set()
|
|
self.alert_listeners: list[Callable[[], None]] = []
|
|
|
|
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)
|
|
|
|
def _on_alert(self, payload: str) -> None:
|
|
for listener in self.alert_listeners:
|
|
listener()
|
|
task = asyncio.create_task(self._dispatch_alert(payload))
|
|
self._tasks.add(task)
|
|
task.add_done_callback(self._tasks.discard)
|
|
|
|
async def _dispatch_alert(self, payload: str) -> None:
|
|
try:
|
|
event = json.loads(payload)
|
|
except json.JSONDecodeError:
|
|
return
|
|
account_id = event.get("account_id")
|
|
targets = [
|
|
sub
|
|
for sub in self._subscribers
|
|
if sub.account_id == account_id and sub.chat_id is None
|
|
]
|
|
if not targets or self._pool is None:
|
|
return
|
|
alert_id = event.get("alert_id")
|
|
item = await watches_read.get_alert(self._pool, alert_id) if alert_id else None
|
|
frame = {
|
|
"type": "alert",
|
|
"alert": item.model_dump(mode="json") if item else None,
|
|
"unseen": await watches_read.unseen_count(self._pool, account_id),
|
|
}
|
|
for sub in targets:
|
|
try:
|
|
sub.queue.put_nowait(frame)
|
|
except asyncio.QueueFull:
|
|
logger.warning("Dropping alert for slow subscriber")
|
|
|
|
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()
|