import json from datetime import UTC, datetime from typing import Any import asyncpg from utils.read import chats as chats_read from utils.read import peers as peers_read from utils.read.models import ( AlertItem, AlertMessage, AlertPage, AlertView, AlertWatch, ChatListItem, ChatRef, PeerView, PersonRef, WatchItem, WatchView, ) WATCHES_CHANGED_CHANNEL = "watches_changed" ALERTS_CHANGED_CHANNEL = "alerts_changed" ALERT_PAGE = 50 MAX_ALERT_PAGE = 200 _WATCH_COLS = "id, account_id, kind, name, params, enabled, created_at, updated_at" _ALERT_COLS = "id, account_id, watch_id, ts, payload, seen, created_at" _LEGACY_EVENTS = {"keyword": "message", "peer_online": "online"} _INSERT_ALERT = f""" WITH alert AS ( INSERT INTO alerts (account_id, watch_id, ts, payload, seen, dedup_key) VALUES ($1, $2, $3, $4::jsonb, false, $5) ON CONFLICT (watch_id, dedup_key) WHERE dedup_key IS NOT NULL DO NOTHING RETURNING {_ALERT_COLS} ), hooks AS ( INSERT INTO webhook_deliveries (webhook_id, alert_id, event) SELECT w.id, alert.id, 'alert' FROM alert JOIN webhooks w ON w.account_id = alert.account_id AND w.enabled AND (w.watch_ids IS NULL OR alert.watch_id = ANY(w.watch_ids)) ) SELECT * FROM alert """ # noqa: S608 _WATCH_STATS = """ SELECT watch_id, count(*) AS count, max(ts) AS last FROM alerts WHERE watch_id = ANY($1::int[]) GROUP BY watch_id """ _ALERT_ROWS = """ SELECT a.id, a.account_id, a.ts, a.seen, a.payload, w.id AS watch_id, w.kind AS watch_kind, w.name AS watch_name, w.params AS watch_params, m.text AS stored_text, m.sender_id AS stored_sender, m.date AS stored_date, m.deleted_at IS NOT NULL AS stored_deleted, m.edited_at IS NOT NULL AS stored_edited FROM alerts a JOIN watches w ON w.id = a.watch_id LEFT JOIN LATERAL ( SELECT text, sender_id, date, deleted_at, edited_at FROM messages WHERE account_id = a.account_id AND chat_id = (a.payload->>'chat_id')::bigint AND message_id = (a.payload->>'message_id')::bigint ORDER BY date DESC LIMIT 1 ) m ON true """ def _to_watch(row: asyncpg.Record) -> WatchView: data = dict(row) data["params"] = json.loads(data["params"]) return WatchView(**data) def _to_alert(row: asyncpg.Record) -> AlertView: data = dict(row) data["payload"] = json.loads(data["payload"]) return AlertView(**data) async def _notify(pool: asyncpg.Pool, channel: str, payload: dict[str, Any]) -> None: await pool.execute("SELECT pg_notify($1, $2)", channel, json.dumps(payload)) async def list_watches(pool: asyncpg.Pool, account_id: int) -> list[WatchView]: rows = await pool.fetch( f"SELECT {_WATCH_COLS} FROM watches WHERE account_id = $1 " # noqa: S608 "ORDER BY id DESC", account_id, ) return [_to_watch(row) for row in rows] async def get_watch(pool: asyncpg.Pool, watch_id: int) -> WatchView | None: row = await pool.fetchrow( f"SELECT {_WATCH_COLS} FROM watches WHERE id = $1", # noqa: S608 watch_id, ) return _to_watch(row) if row else None async def create_watch( # noqa: PLR0913 pool: asyncpg.Pool, account_id: int, kind: str, params: dict[str, Any], *, name: str | None = None, enabled: bool = True, ) -> WatchView: row = await pool.fetchrow( "INSERT INTO watches (account_id, kind, name, params, enabled) " # noqa: S608 f"VALUES ($1, $2, $3, $4::jsonb, $5) RETURNING {_WATCH_COLS}", account_id, kind, name, json.dumps(params), enabled, ) await pool.execute(f"NOTIFY {WATCHES_CHANGED_CHANNEL}") return _to_watch(row) async def update_watch( # noqa: PLR0913 pool: asyncpg.Pool, watch_id: int, kind: str, params: dict[str, Any], *, name: str | None, enabled: bool, ) -> WatchView | None: row = await pool.fetchrow( "UPDATE watches SET kind = $2, name = $3, params = $4::jsonb, " # noqa: S608 "enabled = $5, updated_at = now() " f"WHERE id = $1 RETURNING {_WATCH_COLS}", watch_id, kind, name, json.dumps(params), enabled, ) if row is None: return None await pool.execute(f"NOTIFY {WATCHES_CHANGED_CHANNEL}") return _to_watch(row) async def delete_watch(pool: asyncpg.Pool, watch_id: int) -> bool: account_id = await pool.fetchval( "DELETE FROM watches WHERE id = $1 RETURNING account_id", watch_id ) if account_id is None: return False await pool.execute(f"NOTIFY {WATCHES_CHANGED_CHANNEL}") await _notify(pool, ALERTS_CHANGED_CHANNEL, {"account_id": account_id}) return True def _chat_ref(chat_id: int, chat: ChatListItem | None, name: str | None) -> ChatRef: if chat is None: kind = "private" if chat_id > 0 else "group" return ChatRef(id=chat_id, title=name, kind=kind, has_avatar=False) kind = "channel" if chat.is_broadcast else chat.kind return ChatRef( id=chat_id, title=chat.title or name, kind=kind, has_avatar=chat.has_avatar ) def _peer_name(peer: PeerView) -> str | None: full = " ".join(part for part in (peer.first_name, peer.last_name) if part) return full or peer.username class Refs: def __init__(self, chats: list[ChatListItem], people: list[PeerView]) -> None: self._chats = {chat.chat_id: chat for chat in chats} self._people = {peer.peer_id: peer for peer in people} def chat(self, chat_id: int | None) -> ChatRef | None: if chat_id is None: return None peer = self._people.get(chat_id) name = _peer_name(peer) if peer else None return _chat_ref(chat_id, self._chats.get(chat_id), name) def person(self, person_id: int | None) -> PersonRef | None: if person_id is None: return None peer = self._people.get(person_id) chat = self._chats.get(person_id) name = (_peer_name(peer) if peer else None) or (chat.title if chat else None) has_avatar = (peer.has_avatar if peer else False) or bool( chat and chat.has_avatar ) return PersonRef(id=person_id, name=name, has_avatar=has_avatar) async def load_refs( pool: asyncpg.Pool, account_id: int, chat_ids: set[int], person_ids: set[int] ) -> Refs: chats = await chats_read.get_chats(pool, account_id, sorted(chat_ids | person_ids)) people = await peers_read.get_peers( pool, account_id, sorted(person_ids | {i for i in chat_ids if i > 0}) ) return Refs(chats, people) def _int(value: object) -> int | None: return value if isinstance(value, int) and not isinstance(value, bool) else None def _watch_person(params: dict[str, Any]) -> int | None: return _int(params.get("sender_id")) or _int(params.get("peer_id")) async def describe_watches( pool: asyncpg.Pool, account_id: int, watches: list[WatchView] ) -> list[WatchItem]: targets = [ (_int(watch.params.get("chat_id")), _watch_person(watch.params)) for watch in watches ] refs = await load_refs( pool, account_id, {chat for chat, _ in targets if chat is not None}, {person for _, person in targets if person is not None}, ) stats = { row["watch_id"]: row for row in await pool.fetch(_WATCH_STATS, [watch.id for watch in watches]) } return [ WatchItem( **watch.model_dump(), chat=refs.chat(chat), person=refs.person(person), alert_count=stats[watch.id]["count"] if watch.id in stats else 0, last_alert_at=stats[watch.id]["last"] if watch.id in stats else None, ) for watch, (chat, person) in zip(watches, targets, strict=True) ] def _path(chat_id: int | None, message_id: int | None, peer_id: int | None) -> str: if chat_id is not None and message_id is not None: return f"/app/{chat_id}?message={message_id}" if peer_id is not None: return f"/app/{peer_id}" return "/app" def _alert_item(row: asyncpg.Record, refs: Refs) -> AlertItem: payload = json.loads(row["payload"]) event = payload.get("event") or _LEGACY_EVENTS.get(row["watch_kind"], "message") chat_id = _int(payload.get("chat_id")) message_id = _int(payload.get("message_id")) peer_id = _int(payload.get("peer_id")) message = None if message_id is not None: date = row["stored_date"] or ( datetime.fromisoformat(payload["date"]) if payload.get("date") else None ) message = AlertMessage( id=message_id, text=payload.get("text") or row["stored_text"], sender=refs.person(_int(payload.get("sender_id")) or row["stored_sender"]), date=date, deleted=bool(row["stored_deleted"]) or event == "delete", edited=bool(row["stored_edited"]) or event == "edit", ) return AlertItem( id=row["id"], account_id=row["account_id"], ts=row["ts"], seen=row["seen"], event=event, watch=AlertWatch( id=row["watch_id"], kind=row["watch_kind"], name=row["watch_name"], params=json.loads(row["watch_params"]), ), chat=refs.chat(chat_id), message=message, peer=refs.person(peer_id), away_since=payload.get("since"), path=_path(chat_id, message_id, peer_id), ) async def _alert_items( pool: asyncpg.Pool, account_id: int, rows: list[asyncpg.Record] ) -> list[AlertItem]: chat_ids: set[int] = set() person_ids: set[int] = set() for row in rows: payload = json.loads(row["payload"]) for key, bucket in ( ("chat_id", chat_ids), ("sender_id", person_ids), ("peer_id", person_ids), ): value = _int(payload.get(key)) if value is not None: bucket.add(value) if row["stored_sender"] is not None: person_ids.add(row["stored_sender"]) refs = await load_refs(pool, account_id, chat_ids, person_ids) return [_alert_item(row, refs) for row in rows] async def list_alerts( pool: asyncpg.Pool, account_id: int, *, seen: bool | None = None, before: int | None = None, limit: int = ALERT_PAGE, ) -> AlertPage: size = max(1, min(limit, MAX_ALERT_PAGE)) params: list[object] = [account_id] where = "a.account_id = $1" if seen is not None: params.append(seen) where += f" AND a.seen = ${len(params)}" if before is not None: params.append(before) where += f" AND a.id < ${len(params)}" params.append(size + 1) rows = await pool.fetch( f"{_ALERT_ROWS} WHERE {where} ORDER BY a.id DESC LIMIT ${len(params)}", *params ) items = await _alert_items(pool, account_id, rows[:size]) return AlertPage( items=items, next_before=items[-1].id if len(rows) > size else None, unseen=await unseen_count(pool, account_id), ) async def get_alert(pool: asyncpg.Pool, alert_id: int) -> AlertItem | None: row = await pool.fetchrow(f"{_ALERT_ROWS} WHERE a.id = $1", alert_id) if row is None: return None return (await _alert_items(pool, row["account_id"], [row]))[0] async def unseen_count(pool: asyncpg.Pool, account_id: int) -> int: return await pool.fetchval( "SELECT count(*) FROM alerts WHERE account_id = $1 AND NOT seen", account_id ) async def insert_alert( pool: asyncpg.Pool, account_id: int, watch_id: int, payload: dict[str, Any], *, dedup_key: str | None = None, ) -> AlertView | None: row = await pool.fetchrow( _INSERT_ALERT, account_id, watch_id, datetime.now(UTC), json.dumps(payload), dedup_key, ) if row is None: return None await _notify( pool, ALERTS_CHANGED_CHANNEL, {"account_id": account_id, "alert_id": row["id"]} ) return _to_alert(row) async def mark_alert_seen(pool: asyncpg.Pool, alert_id: int) -> bool: account_id = await pool.fetchval( "UPDATE alerts SET seen = true WHERE id = $1 RETURNING account_id", alert_id ) if account_id is None: return False await _notify(pool, ALERTS_CHANGED_CHANNEL, {"account_id": account_id}) return True async def mark_all_seen( pool: asyncpg.Pool, account_id: int, up_to: int | None = None ) -> int: result = await pool.execute( "UPDATE alerts SET seen = true WHERE account_id = $1 AND NOT seen " "AND ($2::int IS NULL OR id <= $2)", account_id, up_to, ) marked = int(result.split()[-1]) if marked: await _notify(pool, ALERTS_CHANGED_CHANNEL, {"account_id": account_id}) return marked