import json from dataclasses import dataclass from datetime import UTC, datetime, timedelta from enum import StrEnum import asyncpg from utils.read import chats, peers from utils.read.accounts import self_user_id from utils.read.models import ( DEFAULT_LIMIT, MAX_LIMIT, ChatListItem, FeedChat, FeedItem, FeedMessagesPage, FeedPage, FeedUnseen, MessageView, PeerView, ) CURSOR_PATTERN = r"^\d+\.\d+$" _EPOCH = datetime(1970, 1, 1, tzinfo=UTC) _MICROSECOND = timedelta(microseconds=1) _ITEM_COLS = ( "id, kind, chat_id, chat_kind, message_id, message_date, at, since, " "message_count, detail" ) _VERSIONS = """ SELECT v.chat_id, v.message_id, v.observed_at, v.edit_date, v.text, p.text AS before, p.observed_at IS NOT NULL AS has_before FROM unnest($2::bigint[], $3::bigint[], $4::timestamptz[]) AS k(chat_id, message_id, observed_at) JOIN message_versions v ON v.account_id = $1 AND v.chat_id = k.chat_id AND v.message_id = k.message_id AND v.observed_at = k.observed_at LEFT JOIN LATERAL ( SELECT text, observed_at FROM message_versions p WHERE p.account_id = $1 AND p.chat_id = v.chat_id AND p.message_id = v.message_id AND p.observed_at < v.observed_at ORDER BY p.observed_at DESC LIMIT 1 ) p ON true """ _UNSEEN = """ SELECT count(*) FILTER (WHERE kind = 'deleted') AS deleted, count(*) FILTER (WHERE kind = 'edited') AS edited FROM feed_items WHERE account_id = $1 AND at > coalesce($2, '-infinity'::timestamptz) """ _MARK_SEEN = """ INSERT INTO feed_seen (account_id, seen_at) VALUES ($1, coalesce($2, now())) ON CONFLICT (account_id) DO UPDATE SET seen_at = greatest(feed_seen.seen_at, EXCLUDED.seen_at) """ class FeedKind(StrEnum): ALL = "all" MESSAGES = "messages" DELETED = "deleted" EDITED = "edited" PROFILE = "profile" class FeedChatKind(StrEnum): DM = "dm" GROUP = "group" CHANNEL = "channel" @dataclass(slots=True) class FeedFilter: kind: FeedKind = FeedKind.ALL chat_id: int | None = None chat_kind: FeedChatKind | None = None since: datetime | None = None def encode_cursor(at: datetime, key: int) -> str: return f"{(at - _EPOCH) // _MICROSECOND}.{key}" def decode_cursor(cursor: str) -> tuple[datetime, int]: micros, _, key = cursor.partition(".") return _EPOCH + int(micros) * _MICROSECOND, int(key) def _where( account_id: int, feed: FeedFilter, cursor: str | None ) -> tuple[str, list[object]]: params: list[object] = [account_id] clauses = ["account_id = $1"] if feed.kind is FeedKind.MESSAGES: clauses.append("kind <> 'profile'") exact = feed.kind not in {FeedKind.ALL, FeedKind.MESSAGES} conditions = ( ("kind = {}", feed.kind.value if exact else None), ("chat_id = {}", feed.chat_id), ("chat_kind = {}", feed.chat_kind.value if feed.chat_kind else None), ("at >= {}", feed.since), ) for template, value in conditions: if value is not None: params.append(value) clauses.append(template.format(f"${len(params)}")) if cursor is not None: params.extend(decode_cursor(cursor)) clauses.append(f"(at, id) < (${len(params) - 1}, ${len(params)})") return " AND ".join(clauses), params def _peer_name(peer: PeerView) -> str | None: name = " ".join(part for part in (peer.first_name, peer.last_name) if part) return name or peer.username async def _versions( pool: asyncpg.Pool, account_id: int, rows: list[asyncpg.Record] ) -> dict[tuple[int, int, datetime], asyncpg.Record]: edited = [row for row in rows if row["kind"] == FeedKind.EDITED] if not edited: return {} found = await pool.fetch( _VERSIONS, account_id, [row["chat_id"] for row in edited], [row["message_id"] for row in edited], [row["at"] for row in edited], ) return { (row["chat_id"], row["message_id"], row["observed_at"]): row for row in found } async def _sender_names( pool: asyncpg.Pool, account_id: int, messages: dict[tuple[int, int], MessageView], chat_views: dict[int, ChatListItem], ) -> dict[int, str | None]: ids = {m.sender_id for m in messages.values() if m.sender_id is not None} found = await peers.get_peers(pool, account_id, list(ids)) names: dict[int, str | None] = { chat_id: chat.title for chat_id, chat in chat_views.items() } names.update({peer.peer_id: _peer_name(peer) for peer in found}) return names async def _build_items( pool: asyncpg.Pool, account_id: int, rows: list[asyncpg.Record] ) -> list[FeedItem]: if not rows: return [] messages = await chats.get_messages_at( pool, account_id, [ (row["chat_id"], row["message_id"], row["message_date"]) for row in rows if row["kind"] != FeedKind.PROFILE ], ) versions = await _versions(pool, account_id, rows) chat_views = { chat.chat_id: chat for chat in await chats.get_chats( pool, account_id, list({row["chat_id"] for row in rows}) ) } names = await _sender_names(pool, account_id, messages, chat_views) self_id = await self_user_id(pool, account_id) items: list[FeedItem] = [] for row in rows: message = messages.get((row["chat_id"], row["message_id"])) chat = chat_views.get(row["chat_id"]) version = versions.get((row["chat_id"], row["message_id"], row["at"])) sender_id = message.sender_id if message else None items.append( FeedItem( id=row["id"], kind=row["kind"], at=row["at"], since=row["since"], message_count=row["message_count"], chat=FeedChat( chat_id=row["chat_id"], title=chat.title if chat else None, kind=row["chat_kind"], has_avatar=chat.has_avatar if chat else False, ), message_id=row["message_id"], message=message, sender_id=sender_id, sender_name=names.get(sender_id) if sender_id is not None else None, outgoing=self_id is not None and sender_id == self_id, before=version["before"] if version else None, after=version["text"] if version else None, has_before=bool(version and version["has_before"]), edit_date=version["edit_date"] if version else None, detail=json.loads(row["detail"]) if row["detail"] else None, ) ) return items async def list_feed( pool: asyncpg.Pool, account_id: int, feed: FeedFilter, *, limit: int = DEFAULT_LIMIT, cursor: str | None = None, ) -> FeedPage: size = max(1, min(limit, MAX_LIMIT)) where, params = _where(account_id, feed, cursor) params.append(size + 1) rows = await pool.fetch( f"SELECT {_ITEM_COLS} FROM feed_items WHERE {where} " # noqa: S608 f"ORDER BY at DESC, id DESC LIMIT ${len(params)}", *params, ) page = rows[:size] more = len(rows) > size return FeedPage( items=await _build_items(pool, account_id, page), next_cursor=encode_cursor(page[-1]["at"], page[-1]["id"]) if more else None, ) async def unseen(pool: asyncpg.Pool, account_id: int) -> FeedUnseen: seen_at = await pool.fetchval( "SELECT seen_at FROM feed_seen WHERE account_id = $1", account_id ) counts = await pool.fetchrow(_UNSEEN, account_id, seen_at) latest = await list_feed( pool, account_id, FeedFilter(kind=FeedKind.MESSAGES), limit=1 ) return FeedUnseen( count=counts["deleted"] + counts["edited"], deleted=counts["deleted"], edited=counts["edited"], seen_at=seen_at, latest=latest.items[0] if latest.items else None, ) async def mark_seen( pool: asyncpg.Pool, account_id: int, at: datetime | None ) -> FeedUnseen: await pool.execute(_MARK_SEEN, account_id, at) return await unseen(pool, account_id) async def item_messages( pool: asyncpg.Pool, account_id: int, item_id: int, *, limit: int = DEFAULT_LIMIT, cursor: str | None = None, ) -> FeedMessagesPage | None: row = await pool.fetchrow( "SELECT chat_id, since, at FROM feed_items " "WHERE id = $1 AND account_id = $2 AND kind = 'deleted'", item_id, account_id, ) if row is None: return None size = max(1, min(limit, MAX_LIMIT)) views = await chats.get_deleted_between( pool, account_id, row["chat_id"], row["since"], row["at"], size + 1, decode_cursor(cursor) if cursor else None, ) page = views[:size] more = len(views) > size return FeedMessagesPage( messages=page, next_cursor=encode_cursor(page[-1].date, page[-1].message_id) if more else None, )