feat: web UI chat render, panels, presence + analytics

This commit is contained in:
h
2026-05-31 19:41:01 +02:00
parent 75425d1bee
commit ed469ba8dd
83 changed files with 6034 additions and 136 deletions
+10
View File
@@ -11,12 +11,16 @@ from starlette.applications import Starlette
from api.auth import BearerAuthMiddleware
from api.mcp.server import mcp
from api.realtime import hub
from api.routers import (
accounts,
analytics,
annotations,
avatars,
backfill,
chats,
custom_emoji,
events,
folders,
media,
peers,
@@ -39,7 +43,10 @@ mcp_app = mcp.http_app(path="/")
@asynccontextmanager
async def lifespan(app_: Starlette) -> AsyncGenerator[None]:
pool = await container.get(asyncpg.Pool)
await hub.start(pool)
yield
await hub.stop()
await app_.state.dishka_container.close()
@@ -59,6 +66,7 @@ async def health(pool: FromDishka[asyncpg.Pool]) -> dict[str, bool]:
app.include_router(accounts.router)
app.include_router(analytics.router)
app.include_router(policy.router)
app.include_router(folders.router)
app.include_router(backfill.router)
@@ -66,8 +74,10 @@ app.include_router(search.router)
app.include_router(chats.router)
app.include_router(media.router)
app.include_router(avatars.router)
app.include_router(custom_emoji.router)
app.include_router(social.router)
app.include_router(presence.router)
app.include_router(events.router)
app.include_router(peers.router)
app.include_router(annotations.router)
app.include_router(watches.router)
+123
View File
@@ -0,0 +1,123 @@
import asyncio
import json
import logging
from typing import Any
import asyncpg
from utils.env import env
from utils.events import BG_EVENTS_CHANNEL
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._conn: asyncpg.Connection | None = None
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
conn = await asyncpg.connect(dsn=env.db.connection_url)
await conn.add_listener(BG_EVENTS_CHANNEL, self._on_notify)
self._conn = conn
logger.info("Realtime hub listening on %s", BG_EVENTS_CHANNEL)
async def stop(self) -> None:
for task in self._tasks:
task.cancel()
if self._conn is not None:
await self._conn.close()
self._conn = None
def _on_notify(
self, _conn: asyncpg.Connection, _pid: int, _channel: str, 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")
targets = [
sub
for sub in self._subscribers
if sub.account_id == account_id
and (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"],
}
return None
hub = EventHub()
+30
View File
@@ -0,0 +1,30 @@
from typing import Annotated
import asyncpg
from dishka.integrations.fastapi import DishkaRoute, FromDishka
from fastapi import APIRouter, Query
from utils.read import analytics
from utils.read.models import ResponseStats, VolumeBucket
router = APIRouter(prefix="/api/analytics", tags=["analytics"], route_class=DishkaRoute)
AccountId = Annotated[int, Query()]
ChatId = Annotated[int, Query()]
@router.get("/volume")
async def volume(
pool: FromDishka[asyncpg.Pool],
account_id: AccountId,
chat_id: ChatId,
days: Annotated[int, Query()] = 90,
) -> list[VolumeBucket]:
return await analytics.message_volume(pool, account_id, chat_id, days=days)
@router.get("/response-time")
async def response_time(
pool: FromDishka[asyncpg.Pool], account_id: AccountId, chat_id: ChatId
) -> ResponseStats:
return await analytics.response_stats(pool, account_id, chat_id)
+12
View File
@@ -24,6 +24,10 @@ class FetchMediaRequest(BaseModel):
message_id: int
class SyncDialogsRequest(BaseModel):
account_id: int
class EnqueueResponse(BaseModel):
job_id: int
@@ -78,6 +82,14 @@ async def enqueue_fetch_media(
return EnqueueResponse(job_id=job_id)
@router.post("/dialogs/sync", status_code=201)
async def enqueue_sync_dialogs(
pool: FromDishka[asyncpg.Pool], body: SyncDialogsRequest
) -> EnqueueResponse:
job_id = await enqueue(pool, body.account_id, "sync_dialogs", {})
return EnqueueResponse(job_id=job_id)
@router.get("/jobs")
async def list_jobs(
pool: FromDishka[asyncpg.Pool],
+9
View File
@@ -13,7 +13,9 @@ from utils.read.models import (
MessageVersionView,
MessageView,
Page,
PinnedView,
)
from utils.read.pinned import get_pinned
router = APIRouter(prefix="/api", tags=["chats"], route_class=DishkaRoute)
@@ -55,6 +57,13 @@ async def chat_history(
)
@router.get("/chats/{chat_id}/pinned")
async def chat_pinned(
pool: FromDishka[asyncpg.Pool], chat_id: int, account_id: AccountId
) -> PinnedView | None:
return await get_pinned(pool, account_id, chat_id)
@router.post("/chats/{chat_id}/enrich")
async def enrich_chat(
pool: FromDishka[asyncpg.Pool], chat_id: int, body: EnrichRequest
+35
View File
@@ -0,0 +1,35 @@
from typing import Annotated
import asyncpg
from dishka.integrations.fastapi import DishkaRoute, FromDishka
from fastapi import APIRouter, HTTPException, Query
from fastapi.responses import FileResponse
from utils.jobs import enqueue
from utils.read.custom_emoji import current_custom_emoji
from utils.storage import ContentAddressedStorage
router = APIRouter(
prefix="/api/custom-emoji", tags=["custom-emoji"], route_class=DishkaRoute
)
@router.get("/{custom_emoji_id}")
async def serve_custom_emoji(
pool: FromDishka[asyncpg.Pool],
storage: FromDishka[ContentAddressedStorage],
custom_emoji_id: int,
account_id: Annotated[int, Query()],
) -> FileResponse:
emoji = await current_custom_emoji(pool, custom_emoji_id)
if emoji is None or not emoji.downloaded or emoji.storage_key is None:
await enqueue(
pool, account_id, "fetch_custom_emoji", {"custom_emoji_id": custom_emoji_id}
)
raise HTTPException(
status_code=409, detail="custom emoji not downloaded; fetching"
)
return FileResponse(
storage.url(emoji.storage_key),
media_type=emoji.mime or "application/octet-stream",
)
+42
View File
@@ -0,0 +1,42 @@
import asyncio
import json
from collections.abc import AsyncGenerator
from typing import Annotated
from fastapi import APIRouter, Query
from fastapi.responses import StreamingResponse
from api.realtime import Subscriber, hub
router = APIRouter(prefix="/api/events", tags=["events"])
HEARTBEAT_SECONDS = 15
AccountId = Annotated[int, Query()]
ChatId = Annotated[int | None, Query()]
async def _stream(sub: Subscriber) -> AsyncGenerator[str]:
try:
yield ": connected\n\n"
while True:
try:
frame = await asyncio.wait_for(
sub.queue.get(), timeout=HEARTBEAT_SECONDS
)
except TimeoutError:
yield ": keepalive\n\n"
continue
yield f"event: {frame['type']}\ndata: {json.dumps(frame)}\n\n"
finally:
hub.unsubscribe(sub)
@router.get("")
async def events(account_id: AccountId, chat_id: ChatId = None) -> StreamingResponse:
sub = hub.subscribe(account_id, chat_id)
return StreamingResponse(
_stream(sub),
media_type="text/event-stream",
headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
)
+7
View File
@@ -36,6 +36,13 @@ async def presence_history(
)
@router.get("/current")
async def current_presence(
pool: FromDishka[asyncpg.Pool], account_id: AccountId, peer_id: PeerId
) -> PresenceSample | None:
return await presence.current_presence(pool, account_id, peer_id)
@router.get("/hourly")
async def presence_hourly(
pool: FromDishka[asyncpg.Pool],
+5
View File
@@ -3,6 +3,7 @@ from pyrogram.types import Message
from userbot import PyroClient
from userbot.modules.capture import repository
from userbot.modules.capture.repository import CHANNEL_ID_THRESHOLD
from utils.events import notify_bg_event
@PyroClient.on_deleted_messages()
@@ -21,8 +22,12 @@ async def on_deleted_messages(client: PyroClient, messages: list[Message]) -> No
channels.setdefault(chat_id, []).append(message.id)
if box:
await repository.mark_deleted_box(ctx.pool, ctx.account_id, box)
await notify_bg_event(ctx.pool, "delete", ctx.account_id, message_ids=box)
for chat_id, ids in channels.items():
await repository.mark_deleted_channel(ctx.pool, ctx.account_id, chat_id, ids)
await notify_bg_event(
ctx.pool, "delete", ctx.account_id, chat_id=chat_id, message_ids=ids
)
handlers = on_deleted_messages.handlers
+7 -1
View File
@@ -5,6 +5,7 @@ from userbot.modules.capture import repository
from userbot.modules.capture.chat_meta import meta_from_chat
from userbot.modules.capture.message import sender_id
from userbot.modules.media import capture_media, media_unique_id, self_destruct_ttl
from utils.events import notify_bg_event
@PyroClient.on_edited_message()
@@ -32,8 +33,13 @@ async def on_edited_message(client: PyroClient, message: Message) -> None:
has_media=message.media is not None,
is_self_destruct=self_destruct_ttl(message) is not None,
)
if changed and message.media is not None:
if not changed:
return
if message.media is not None:
await capture_media(client, message, ctx, chat_id, message.id, toggles)
await notify_bg_event(
ctx.pool, "edit", ctx.account_id, chat_id=chat_id, message_id=message.id
)
handlers = on_edited_message.handlers
+4
View File
@@ -5,6 +5,7 @@ from userbot.modules.capture import capture_message
from userbot.modules.capture.chat_meta import meta_from_chat
from userbot.modules.stt import is_transcribable
from userbot.modules.stt.gate import safe_transcribe
from utils.events import notify_bg_event
@PyroClient.on_message()
@@ -18,6 +19,9 @@ async def on_message(client: PyroClient, message: Message) -> None:
if not toggles.messages:
return
await capture_message(client, message, ctx, toggles)
await notify_bg_event(
ctx.pool, "message", ctx.account_id, chat_id=meta.chat_id, message_id=message.id
)
if (
toggles.stt
and is_transcribable(message)
+2
View File
@@ -4,6 +4,7 @@ from pyrogram.types import User
from userbot import PyroClient
from userbot.modules.presence import repository
from utils.events import notify_bg_event
@PyroClient.on_user_status()
@@ -22,6 +23,7 @@ async def on_user_status(client: PyroClient, user: User) -> None:
str(user.raw),
)
await ctx.watches.on_status(user.id, is_online=user.status.name.lower() == "online")
await notify_bg_event(ctx.pool, "presence", ctx.account_id, chat_id=user.id)
handlers = on_user_status.handlers
@@ -3,6 +3,7 @@ from pyrogram import raw, utils
from userbot import PyroClient
from userbot.modules.capture import repository
from userbot.modules.capture.chat_meta import meta_from_peer
from utils.events import notify_bg_event
HANDLES = (raw.types.UpdateMessageReactions,)
@@ -47,3 +48,10 @@ async def handle(
await repository.sync_reactions(
ctx.pool, ctx.account_id, meta.chat_id, update.msg_id, current
)
await notify_bg_event(
ctx.pool,
"reaction",
ctx.account_id,
chat_id=meta.chat_id,
message_id=update.msg_id,
)
@@ -4,6 +4,7 @@ from pyrogram import raw, utils
from userbot import PyroClient
from userbot.modules.read_receipts import repository
from utils.events import notify_bg_event
HANDLES = (raw.types.UpdateReadHistoryOutbox,)
@@ -24,3 +25,6 @@ async def handle(
update.max_id,
str(update),
)
await notify_bg_event(
ctx.pool, "receipt", ctx.account_id, chat_id=chat_id, message_id=update.max_id
)
@@ -0,0 +1,34 @@
import asyncpg
_GET = """
SELECT downloaded FROM custom_emoji WHERE custom_emoji_id = $1
"""
_UPSERT_DOWNLOADED = """
INSERT INTO custom_emoji
(custom_emoji_id, storage_key, file_size, mime, kind, downloaded, raw)
VALUES ($1, $2, $3, $4, $5, true, '{}'::jsonb)
ON CONFLICT (custom_emoji_id) DO UPDATE SET
storage_key = EXCLUDED.storage_key,
file_size = EXCLUDED.file_size,
mime = EXCLUDED.mime,
kind = EXCLUDED.kind,
downloaded = true
"""
async def is_downloaded(pool: asyncpg.Pool, custom_emoji_id: int) -> bool:
return bool(await pool.fetchval(_GET, custom_emoji_id))
async def upsert_downloaded( # noqa: PLR0913
pool: asyncpg.Pool,
custom_emoji_id: int,
storage_key: str,
file_size: int | None,
mime: str | None,
kind: str,
) -> None:
await pool.execute(
_UPSERT_DOWNLOADED, custom_emoji_id, storage_key, file_size, mime, kind
)
@@ -2,8 +2,18 @@ from userbot.modules.jobs.handlers import (
backfill,
enrich_chat,
fetch_avatar,
fetch_custom_emoji,
fetch_media,
sync_dialogs,
transcribe,
)
__all__ = ["backfill", "enrich_chat", "fetch_avatar", "fetch_media", "transcribe"]
__all__ = [
"backfill",
"enrich_chat",
"fetch_avatar",
"fetch_custom_emoji",
"fetch_media",
"sync_dialogs",
"transcribe",
]
@@ -1,3 +1,5 @@
from pyrogram.errors import PeerIdInvalid
from userbot.modules.capture import capture_message
from userbot.modules.jobs.context import JobContext
from userbot.modules.jobs.registry import register
@@ -23,11 +25,15 @@ async def backfill(ctx: JobContext) -> None:
max_id = (ctx.job.cursor or {}).get("max_id", 0)
processed = ctx.job.progress.get("processed", 0)
kwargs = {"max_id": max_id} if max_id else {}
async for message in client.get_chat_history(chat_id, **kwargs):
await capture_message(client, message, capture, toggles)
processed += 1
if processed % SAVE_EVERY == 0:
next_max = message.id - 1
await ctx.save_cursor({"max_id": next_max})
await ctx.report_progress({"processed": processed, "max_id": next_max})
try:
async for message in client.get_chat_history(chat_id, **kwargs):
await capture_message(client, message, capture, toggles)
processed += 1
if processed % SAVE_EVERY == 0:
next_max = message.id - 1
await ctx.save_cursor({"max_id": next_max})
await ctx.report_progress({"processed": processed, "max_id": next_max})
except PeerIdInvalid:
await ctx.report_progress({"processed": processed, "error": "peer_id_invalid"})
return
await ctx.report_progress({"processed": processed, "done": True})
@@ -0,0 +1,45 @@
from io import BytesIO
from pyrogram.types import Sticker
from userbot.modules.custom_emoji.repository import is_downloaded, upsert_downloaded
from userbot.modules.jobs.context import JobContext
from userbot.modules.jobs.registry import register
def _kind(sticker: Sticker) -> str:
if sticker.is_animated:
return "animated"
if sticker.is_video:
return "video"
return "static"
@register("fetch_custom_emoji")
async def fetch_custom_emoji(ctx: JobContext) -> None:
client = ctx.client
if client is None:
return
capture = getattr(client, "capture", None)
if capture is None:
return
custom_emoji_id = int(ctx.job.params["custom_emoji_id"])
if await is_downloaded(ctx.pool, custom_emoji_id):
return
stickers = await client.get_custom_emoji_stickers([str(custom_emoji_id)])
if not stickers:
return
sticker = stickers[0]
buffer = await client.download_media(sticker.file_id, in_memory=True)
if not isinstance(buffer, BytesIO):
return
data = buffer.getvalue()
storage_key = capture.storage.put(data)
await upsert_downloaded(
ctx.pool,
custom_emoji_id,
storage_key,
len(data),
sticker.mime_type,
_kind(sticker),
)
@@ -0,0 +1,108 @@
from datetime import UTC, datetime
from pyrogram import Client
from pyrogram.errors import BadRequest, Forbidden
from pyrogram.types import Chat, User
from userbot.modules.avatars import note_avatar
from userbot.modules.capture.context import CaptureContext
from userbot.modules.groups.repository import insert_chat_history
from userbot.modules.jobs.context import JobContext
from userbot.modules.jobs.registry import register
from userbot.modules.profiles.parse import snapshot_from_chat, snapshot_from_high_level
from userbot.modules.profiles.repository import write_profile
SAVE_EVERY = 100
USERS_BATCH = 200
_UPSERT_DIALOG = """
INSERT INTO dialogs (account_id, chat_id) VALUES ($1, $2)
ON CONFLICT (account_id, chat_id) DO UPDATE SET updated_at = now()
"""
async def _save_private(ctx: CaptureContext, chat: Chat, chat_id: int) -> bool:
fields, photo_file_id, photo_unique_id = snapshot_from_chat(chat)
await write_profile(ctx.pool, ctx.account_id, chat_id, fields, str(chat))
if photo_file_id and photo_unique_id:
await note_avatar(
ctx.pool, ctx.account_id, chat_id, "peer", photo_unique_id, photo_file_id
)
return bool(fields.first_name or fields.last_name or fields.username)
async def _enrich_users(client: Client, ctx: CaptureContext, ids: list[int]) -> None:
for start in range(0, len(ids), USERS_BATCH):
batch = ids[start : start + USERS_BATCH]
try:
result = await client.get_users(batch)
except (BadRequest, Forbidden):
continue
users = result if isinstance(result, list) else [result]
for user in users:
if not isinstance(user, User):
continue
fields, photo_file_id, photo_unique_id = snapshot_from_high_level(user)
await write_profile(ctx.pool, ctx.account_id, user.id, fields, str(user))
if photo_file_id and photo_unique_id:
await note_avatar(
ctx.pool,
ctx.account_id,
user.id,
"peer",
photo_unique_id,
photo_file_id,
)
async def _save_group(ctx: CaptureContext, chat: Chat, chat_id: int) -> None:
photo = chat.photo
photo_unique_id = photo.big_photo_unique_id if photo else None
photo_file_id = photo.big_file_id if photo else None
await insert_chat_history(
ctx.pool,
ctx.account_id,
chat_id,
0,
"meta",
chat.title,
photo_unique_id,
None,
datetime.now(UTC),
str(chat),
)
if photo_file_id and photo_unique_id:
await note_avatar(
ctx.pool, ctx.account_id, chat_id, "chat", photo_unique_id, photo_file_id
)
@register("sync_dialogs")
async def sync_dialogs(ctx: JobContext) -> None:
client = ctx.client
if client is None:
return
capture = getattr(client, "capture", None)
if capture is None:
return
processed = ctx.job.progress.get("processed", 0)
nameless: list[int] = []
async for dialog in client.get_dialogs():
chat = dialog.chat
if chat is None or chat.id is None:
continue
chat_id = chat.id
try:
if chat_id > 0:
if not await _save_private(capture, chat, chat_id):
nameless.append(chat_id)
else:
await _save_group(capture, chat, chat_id)
except (BadRequest, Forbidden):
pass
await ctx.pool.execute(_UPSERT_DIALOG, ctx.account_id, chat_id)
processed += 1
if processed % SAVE_EVERY == 0:
await ctx.report_progress({"processed": processed})
await _enrich_users(client, capture, nameless)
await ctx.report_progress({"processed": processed, "done": True})
@@ -20,11 +20,20 @@ _MEDIA_ATTRS = (
)
_WEB_PAGE_ATTRS = ("photo", "video", "animation", "document", "audio")
def media_object(message: Message) -> tuple[str | None, Any]:
for attr in _MEDIA_ATTRS:
obj = getattr(message, attr, None)
if obj is not None:
return attr, obj
web_page = getattr(message, "web_page", None)
if web_page is not None:
for attr in _WEB_PAGE_ATTRS:
obj = getattr(web_page, attr, None)
if obj is not None:
return attr, obj
return None, None
@@ -70,7 +79,8 @@ async def capture_media( # noqa: PLR0913
file_size = existing["file_size"]
downloaded = True
else:
buffer = await client.download_media(message, in_memory=True)
target = message if getattr(message, kind or "", None) is obj else obj
buffer = await client.download_media(target, in_memory=True)
if isinstance(buffer, BytesIO):
data = buffer.getvalue()
storage_key = ctx.storage.put(data)
+16 -1
View File
@@ -1,7 +1,7 @@
from dataclasses import dataclass
from pyrogram import Client, raw
from pyrogram.types import User
from pyrogram.types import Chat, User
@dataclass(frozen=True)
@@ -54,3 +54,18 @@ def snapshot_from_high_level(
is_deleted_account=bool(user.is_deleted),
)
return fields, photo_file_id, photo_unique_id
def snapshot_from_chat(chat: Chat) -> tuple[ProfileFields, str | None, str | None]:
photo = chat.photo
photo_unique_id = photo.big_photo_unique_id if photo else None
photo_file_id = photo.big_file_id if photo else None
fields = ProfileFields(
first_name=chat.first_name,
last_name=chat.last_name,
username=chat.username,
phone=None,
photo_unique_id=photo_unique_id,
is_deleted_account=False,
)
return fields, photo_file_id, photo_unique_id
+13
View File
@@ -13,6 +13,7 @@ from userbot import PyroClient
from userbot.modules.capture import CaptureContext, build_capture_context
from userbot.modules.jobs import JobConsumer
from utils.env import env
from utils.jobs import enqueue
from utils.logging import logger, setup_logging
from utils.read.watches import WATCHES_CHANGED_CHANNEL
from utils.storage import ContentAddressedStorage
@@ -72,6 +73,17 @@ async def _setup_capture(
logger.info("[green]Capture context ready.[/]")
async def _enqueue_sync_dialogs(pool: asyncpg.Pool, account_id: int) -> None:
existing = await pool.fetchval(
"SELECT 1 FROM jobs WHERE account_id = $1 AND kind = 'sync_dialogs' "
"AND status IN ('pending', 'running') LIMIT 1",
account_id,
)
if existing is None:
await enqueue(pool, account_id, "sync_dialogs", {})
logger.info("[green]Queued sync_dialogs.[/]")
async def _listen_changes(
clients: list[PyroClient], tasks: set[asyncio.Task]
) -> asyncpg.Connection:
@@ -134,6 +146,7 @@ async def runner() -> None:
await _setup_capture(pool, client, account_id, storage)
consumer = JobConsumer(client, pool, account_id)
consumer_tasks.append(asyncio.create_task(consumer.run()))
await _enqueue_sync_dialogs(pool, account_id)
if clients:
listen_conn = await _listen_changes(clients, reload_tasks)
+15
View File
@@ -477,3 +477,18 @@ class Alert(SQLModel, table=True):
DateTime(timezone=True), nullable=False, server_default=func.now()
)
)
class Dialog(SQLModel, table=True):
__tablename__ = "dialogs"
account_id: int = Field(primary_key=True)
chat_id: int = Field(primary_key=True)
updated_at: datetime = Field(
sa_column=Column(
DateTime(timezone=True),
nullable=False,
server_default=func.now(),
onupdate=func.now(),
)
)
+29
View File
@@ -0,0 +1,29 @@
import json
from typing import Literal
import asyncpg
BG_EVENTS_CHANNEL = "bg_events"
EventKind = Literal["message", "edit", "delete", "reaction", "presence", "receipt"]
async def notify_bg_event( # noqa: PLR0913
pool: asyncpg.Pool,
kind: EventKind,
account_id: int,
*,
chat_id: int | None = None,
message_id: int | None = None,
message_ids: list[int] | None = None,
) -> None:
payload: dict[str, object] = {"kind": kind, "account_id": account_id}
if chat_id is not None:
payload["chat_id"] = chat_id
if message_id is not None:
payload["message_id"] = message_id
if message_ids is not None:
payload["message_ids"] = message_ids
await pool.execute(
"SELECT pg_notify($1, $2)", BG_EVENTS_CHANNEL, json.dumps(payload)
)
+6
View File
@@ -3,6 +3,12 @@ import asyncpg
from utils.read.models import AccountView
async def self_user_id(pool: asyncpg.Pool, account_id: int) -> int | None:
return await pool.fetchval(
"SELECT tg_user_id FROM accounts WHERE account_id = $1", account_id
)
async def list_accounts(pool: asyncpg.Pool) -> list[AccountView]:
rows = await pool.fetch(
"SELECT account_id, label, phone, tg_user_id, is_active FROM accounts "
+69
View File
@@ -0,0 +1,69 @@
from datetime import UTC, datetime, timedelta
import asyncpg
from utils.read.accounts import self_user_id
from utils.read.models import ResponseStats, VolumeBucket
async def message_volume(
pool: asyncpg.Pool, account_id: int, chat_id: int, *, days: int = 90
) -> list[VolumeBucket]:
self_id = await self_user_id(pool, account_id)
rows = await pool.fetch(
"SELECT date_trunc('day', date) AS bucket, count(*) AS total, "
"count(*) FILTER (WHERE sender_id = $3) AS outgoing "
"FROM messages "
"WHERE account_id = $1 AND chat_id = $2 AND date >= $4 "
"GROUP BY bucket ORDER BY bucket",
account_id,
chat_id,
self_id,
datetime.now(UTC) - timedelta(days=days),
)
return [
VolumeBucket(
bucket=row["bucket"],
total=row["total"],
outgoing=row["outgoing"],
incoming=row["total"] - row["outgoing"],
)
for row in rows
]
async def response_stats(
pool: asyncpg.Pool, account_id: int, chat_id: int
) -> ResponseStats:
self_id = await self_user_id(pool, account_id)
rows = await pool.fetch(
"WITH ordered AS ("
"SELECT sender_id, date, "
"lag(sender_id) OVER w AS prev_sender, "
"lag(date) OVER w AS prev_date "
"FROM messages "
"WHERE account_id = $1 AND chat_id = $2 AND sender_id IS NOT NULL "
"WINDOW w AS (ORDER BY date, message_id)), "
"resp AS ("
"SELECT (sender_id = $3) AS is_mine, "
"EXTRACT(EPOCH FROM (date - prev_date)) AS secs "
"FROM ordered "
"WHERE prev_sender IS NOT NULL AND prev_sender <> sender_id) "
"SELECT is_mine, count(*) AS n, "
"percentile_cont(0.5) WITHIN GROUP (ORDER BY secs) AS median_secs "
"FROM resp GROUP BY is_mine",
account_id,
chat_id,
self_id,
)
stats = ResponseStats(
mine_median_seconds=None, mine_count=0, their_median_seconds=None, their_count=0
)
for row in rows:
if row["is_mine"]:
stats.mine_median_seconds = row["median_secs"]
stats.mine_count = row["n"]
else:
stats.their_median_seconds = row["median_secs"]
stats.their_count = row["n"]
return stats
+60 -16
View File
@@ -1,5 +1,6 @@
import asyncpg
from utils.read.accounts import self_user_id
from utils.read.message_view import build_message_view, load_raw, media_ref_from
from utils.read.models import (
ChatListItem,
@@ -8,6 +9,7 @@ from utils.read.models import (
MessageView,
Page,
)
from utils.read.read_receipts import read_up_to
_MESSAGE_COLS = (
"chat_id, message_id, date, sender_id, text, has_media, is_self_destruct, "
@@ -52,35 +54,43 @@ async def list_chats(
pool: asyncpg.Pool, account_id: int, page: Page
) -> list[ChatListItem]:
rows = await pool.fetch(
"SELECT m.chat_id, count(*) AS message_count, max(m.date) AS last_date, "
"WITH ids AS ("
"SELECT DISTINCT chat_id FROM messages WHERE account_id = $1 "
"UNION SELECT chat_id FROM dialogs WHERE account_id = $1), "
"agg AS (SELECT chat_id, count(*) AS message_count, max(date) AS last_date "
"FROM messages WHERE account_id = $1 GROUP BY chat_id) "
"SELECT ids.chat_id, COALESCE(agg.message_count, 0) AS message_count, "
"agg.last_date AS last_date, "
"(SELECT p.first_name FROM peers p "
"WHERE p.account_id = $1 AND p.peer_id = m.chat_id) AS first_name, "
"WHERE p.account_id = $1 AND p.peer_id = ids.chat_id) AS first_name, "
"(SELECT p.last_name FROM peers p "
"WHERE p.account_id = $1 AND p.peer_id = m.chat_id) AS last_name, "
"WHERE p.account_id = $1 AND p.peer_id = ids.chat_id) AS last_name, "
"(SELECT p.username FROM peers p "
"WHERE p.account_id = $1 AND p.peer_id = m.chat_id) AS username, "
"WHERE p.account_id = $1 AND p.peer_id = ids.chat_id) AS username, "
"(SELECT ch.title FROM chat_history ch "
"WHERE ch.account_id = $1 AND ch.chat_id = m.chat_id "
"WHERE ch.account_id = $1 AND ch.chat_id = ids.chat_id "
"AND ch.title IS NOT NULL ORDER BY ch.ts DESC LIMIT 1) AS group_title, "
"EXISTS (SELECT 1 FROM avatars a "
"WHERE a.account_id = $1 AND a.owner_id = m.chat_id) AS has_avatar, "
"(SELECT COALESCE((p.raw->>'is_bot')::bool, (p.raw->>'bot')::bool, false) "
"FROM peers p WHERE p.account_id = $1 AND p.peer_id = m.chat_id) AS is_bot, "
"WHERE a.account_id = $1 AND a.owner_id = ids.chat_id) AS has_avatar, "
"(SELECT COALESCE((p.raw->>'is_bot')::bool, (p.raw->>'bot')::bool, "
"p.raw->>'type' = 'ChatType.BOT', false) "
"FROM peers p WHERE p.account_id = $1 AND p.peer_id = ids.chat_id) AS is_bot, "
"(SELECT COALESCE((p.raw->>'is_contact')::bool, (p.raw->>'contact')::bool, "
"false) FROM peers p "
"WHERE p.account_id = $1 AND p.peer_id = m.chat_id) AS is_contact, "
"(SELECT ch.raw->'chat'->>'type' = 'ChatType.CHANNEL' FROM chat_history ch "
"WHERE ch.account_id = $1 AND ch.chat_id = m.chat_id "
"AND ch.raw->'chat'->>'type' IS NOT NULL "
"WHERE p.account_id = $1 AND p.peer_id = ids.chat_id) AS is_contact, "
"(SELECT COALESCE(ch.raw->'chat'->>'type', ch.raw->>'type') "
"= 'ChatType.CHANNEL' FROM chat_history ch "
"WHERE ch.account_id = $1 AND ch.chat_id = ids.chat_id "
"AND COALESCE(ch.raw->'chat'->>'type', ch.raw->>'type') IS NOT NULL "
"ORDER BY ch.ts DESC LIMIT 1) AS is_broadcast, "
"(SELECT lm.text FROM messages lm "
"WHERE lm.account_id = $1 AND lm.chat_id = m.chat_id "
"WHERE lm.account_id = $1 AND lm.chat_id = ids.chat_id "
"ORDER BY lm.date DESC, lm.message_id DESC LIMIT 1) AS last_text, "
"(SELECT lm.sender_id FROM messages lm "
"WHERE lm.account_id = $1 AND lm.chat_id = m.chat_id "
"WHERE lm.account_id = $1 AND lm.chat_id = ids.chat_id "
"ORDER BY lm.date DESC, lm.message_id DESC LIMIT 1) AS last_sender_id "
"FROM messages m WHERE m.account_id = $1 "
"GROUP BY m.chat_id ORDER BY last_date DESC LIMIT $2 OFFSET $3",
"FROM ids LEFT JOIN agg ON agg.chat_id = ids.chat_id "
"ORDER BY last_date DESC NULLS LAST, ids.chat_id DESC LIMIT $2 OFFSET $3",
account_id,
page.capped_limit,
page.offset,
@@ -146,9 +156,43 @@ async def get_chat_history(
else:
views.append(_build_album(members, media_by_key))
index = end
await _apply_read_status(pool, account_id, chat_id, views)
return views
async def get_message(
pool: asyncpg.Pool, account_id: int, chat_id: int, message_id: int
) -> MessageView | None:
row = await pool.fetchrow(
f"SELECT {_MESSAGE_COLS} FROM messages " # noqa: S608
"WHERE account_id = $1 AND chat_id = $2 AND message_id = $3",
account_id,
chat_id,
message_id,
)
if row is None:
return None
media_by_key = await _media_map(pool, account_id, [row])
raw = load_raw(row["raw"])
view = build_message_view(row, raw, _single_media(row, raw, media_by_key))
await _apply_read_status(pool, account_id, chat_id, [view])
return view
async def _apply_read_status(
pool: asyncpg.Pool, account_id: int, chat_id: int, views: list[MessageView]
) -> None:
self_id = await self_user_id(pool, account_id)
if self_id is None:
return
marker = await read_up_to(pool, account_id, chat_id)
if marker is None:
return
for view in views:
if view.sender_id == self_id and view.message_id <= marker:
view.read = True
def _build_album(
members: list[tuple[asyncpg.Record, dict]],
media_by_key: dict[tuple[int, int], asyncpg.Record],
+15
View File
@@ -0,0 +1,15 @@
import asyncpg
from utils.read.models import CustomEmojiRef
_GET = """
SELECT storage_key, downloaded, mime, kind FROM custom_emoji
WHERE custom_emoji_id = $1
"""
async def current_custom_emoji(
pool: asyncpg.Pool, custom_emoji_id: int
) -> CustomEmojiRef | None:
row = await pool.fetchrow(_GET, custom_emoji_id)
return CustomEmojiRef(**dict(row)) if row else None
+43 -1
View File
@@ -1,5 +1,6 @@
import asyncpg
from utils.read.message_view import load_raw
from utils.read.models import MediaVersionView, MediaView
_MEDIA_COLS = (
@@ -9,6 +10,45 @@ _MEDIA_COLS = (
_VERSION_COLS = "id, kind, storage_key, file_size, mime, observed_at"
_WEB_PAGE_MEDIA_KINDS = ("photo", "video", "animation", "document", "audio")
async def _web_page_media_stub(
pool: asyncpg.Pool, account_id: int, chat_id: int, message_id: int
) -> MediaView | None:
row = await pool.fetchrow(
"SELECT date, raw FROM messages "
"WHERE account_id = $1 AND chat_id = $2 AND message_id = $3",
account_id,
chat_id,
message_id,
)
if row is None:
return None
web_page = load_raw(row["raw"]).get("web_page")
if not isinstance(web_page, dict):
return None
kind = next(
(k for k in _WEB_PAGE_MEDIA_KINDS if isinstance(web_page.get(k), dict)), None
)
if kind is None:
return None
obj = web_page[kind]
return MediaView(
id=0,
account_id=account_id,
chat_id=chat_id,
message_id=message_id,
kind=kind,
storage_key=None,
file_size=obj.get("file_size"),
mime=obj.get("mime_type"),
ttl_seconds=None,
downloaded=False,
extracted_text=None,
created_at=row["date"],
)
async def get_media(pool: asyncpg.Pool, media_id: int) -> MediaView | None:
row = await pool.fetchrow(
@@ -28,7 +68,9 @@ async def get_message_media(
chat_id,
message_id,
)
return MediaView(**dict(row)) if row else None
if row is not None:
return MediaView(**dict(row))
return await _web_page_media_stub(pool, account_id, chat_id, message_id)
async def get_media_versions(
+29
View File
@@ -141,6 +141,13 @@ class ServiceView(BaseModel):
duration: int | None = None
class PinnedView(BaseModel):
message_id: int
text: str | None = None
media_kind: str | None = None
sender_name: str | None = None
class StickerView(BaseModel):
emoji: str | None = None
set_name: str | None = None
@@ -178,6 +185,7 @@ class MessageView(BaseModel):
sticker: StickerView | None = None
is_sticker: bool = False
is_animated_emoji: bool = False
read: bool = False
class MessageVersionView(BaseModel):
@@ -217,6 +225,13 @@ class AvatarRef(BaseModel):
mime: str | None
class CustomEmojiRef(BaseModel):
storage_key: str | None
downloaded: bool
mime: str | None
kind: str | None
class CallbackView(BaseModel):
position: int
label: str | None
@@ -256,6 +271,20 @@ class PresenceHourly(BaseModel):
last_seen: datetime | None
class VolumeBucket(BaseModel):
bucket: datetime
total: int
outgoing: int
incoming: int
class ResponseStats(BaseModel):
mine_median_seconds: float | None
mine_count: int
their_median_seconds: float | None
their_count: int
class PeerView(BaseModel):
peer_id: int
first_name: str | None
+32
View File
@@ -0,0 +1,32 @@
import asyncpg
from utils.read.message_view import _media_kind, _peer_name, load_raw
from utils.read.models import PinnedView
_PINNED_SQL = (
"SELECT raw FROM messages "
"WHERE account_id = $1 AND chat_id = $2 "
"AND raw->>'service' LIKE '%PINNED_MESSAGE%' "
"ORDER BY date DESC, message_id DESC LIMIT 1"
)
async def get_pinned(
pool: asyncpg.Pool, account_id: int, chat_id: int
) -> PinnedView | None:
row = await pool.fetchrow(_PINNED_SQL, account_id, chat_id)
if row is None:
return None
pinned = load_raw(row["raw"]).get("pinned_message")
if not isinstance(pinned, dict):
return None
message_id = pinned.get("id")
if message_id is None:
return None
sender = pinned.get("from_user")
return PinnedView(
message_id=message_id,
text=pinned.get("text") or pinned.get("caption"),
media_kind=_media_kind(pinned),
sender_name=_peer_name(sender) if isinstance(sender, dict) else None,
)
+13
View File
@@ -33,6 +33,19 @@ async def presence_history( # noqa: PLR0913
return [PresenceSample(**dict(row)) for row in rows]
async def current_presence(
pool: asyncpg.Pool, account_id: int, peer_id: int
) -> PresenceSample | None:
row = await pool.fetchrow(
"SELECT peer_id, ts, status, last_online_date, next_offline_date "
"FROM presence WHERE account_id = $1 AND peer_id = $2 "
"ORDER BY ts DESC LIMIT 1",
account_id,
peer_id,
)
return PresenceSample(**dict(row)) if row is not None else None
async def presence_hourly(
pool: asyncpg.Pool,
account_id: int,
+10
View File
@@ -0,0 +1,10 @@
import asyncpg
async def read_up_to(pool: asyncpg.Pool, account_id: int, chat_id: int) -> int | None:
return await pool.fetchval(
"SELECT max(message_id) FROM read_receipts "
"WHERE account_id = $1 AND chat_id = $2 AND kind = 'read'",
account_id,
chat_id,
)