feat(backend): alerts come enriched and paged with an unseen count and a live alert frame
This commit is contained in:
1 parent
4469938f47
commit
2bb8ca45fb
7 files changed
+682
-65
No files matched your search
@@ -37,6 +37,7 @@ from utils.read.read_receipts import read_info
|
||||
from utils.search.models import SearchFilters
|
||||
from utils.search.repository import search_messages
|
||||
from utils.storage import ContentAddressedStorage
|
||||
from utils.watch_rules import WatchKind, WatchParamsError, clean_name, normalize_params
|
||||
|
||||
INSTRUCTIONS = """\
|
||||
beavergram archives Telegram data (chats, messages, media, presence, stories,
|
||||
@@ -76,8 +77,10 @@ Typical flows:
|
||||
returns a public link valid 24h that needs no auth (`curl -O <link>`), plus
|
||||
the content itself for small images and text files. `get_media` returns
|
||||
metadata only.
|
||||
- Monitoring: `set_watch` creates a local rule, `list_watches` lists rules, and
|
||||
`list_alerts` reads what those rules fired.
|
||||
- Monitoring: `set_watch` creates a local rule (kinds `keyword`,
|
||||
`new_message`, `message_deleted`, `message_edited`, `peer_online`; params in
|
||||
its description), `list_watches` lists rules, and `list_alerts` reads what
|
||||
those rules fired, with chat, message preview and a web client `path`.
|
||||
|
||||
Watched bots:
|
||||
- a bot whose token the user added is just another archived account: it has its
|
||||
@@ -420,13 +423,37 @@ async def get_annotations(
|
||||
@mcp.tool
|
||||
async def set_watch(
|
||||
account_id: int,
|
||||
kind: str,
|
||||
kind: WatchKind,
|
||||
params: dict[str, Any] | None = None,
|
||||
name: str | None = None,
|
||||
enabled: bool = True,
|
||||
) -> dict[str, Any]:
|
||||
"""Create a local watch rule (the only local write MCP is allowed)."""
|
||||
"""Create a local watch rule (the only local write MCP is allowed).
|
||||
|
||||
A watch fires an alert into the web client (live counter and list) and to
|
||||
the user's webhooks. The user's own messages never fire. Kinds and params:
|
||||
|
||||
- `keyword`: `{"pattern": str, "regex": bool = false, "chat_id"?: int,
|
||||
"sender_id"?: int}`. Case-insensitive substring, or a Python regex when
|
||||
`regex` is true (at most 200 characters, no nested repeats like
|
||||
`(a+)+`). Checked on new messages, on edits (the new text) and on voice
|
||||
transcripts. Without `chat_id` it watches every chat.
|
||||
- `new_message`: `{"chat_id"?: int, "sender_id"?: int}`, at least one: any
|
||||
new message in that chat, or from that person anywhere, or both.
|
||||
- `message_deleted`: `{"chat_id"?: int}`. A stored message deleted in that
|
||||
chat; without `chat_id`, in any private chat.
|
||||
- `message_edited`: `{"chat_id"?: int}`. Same scopes as `message_deleted`.
|
||||
- `peer_online`: `{"peer_id": int}`. The person comes online after at
|
||||
least 15 minutes away.
|
||||
|
||||
`name` is an optional human label shown in the list and in alerts.
|
||||
"""
|
||||
try:
|
||||
clean = normalize_params(kind, params or {})
|
||||
except WatchParamsError as exc:
|
||||
raise ToolError(str(exc)) from exc
|
||||
watch = await watches.create_watch(
|
||||
await _pool(), account_id, kind, params or {}, enabled=enabled
|
||||
await _pool(), account_id, kind, clean, name=clean_name(name), enabled=enabled
|
||||
)
|
||||
return watch.model_dump(mode="json")
|
||||
|
||||
@@ -441,12 +468,19 @@ async def list_watches(account_id: int) -> list[dict[str, Any]]:
|
||||
async def list_alerts(
|
||||
account_id: int,
|
||||
seen: bool | None = None,
|
||||
before: int | None = None,
|
||||
limit: int = DEFAULT_LIMIT,
|
||||
offset: int = 0,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""List fired alerts from watch rules."""
|
||||
return _dump(
|
||||
await watches.list_alerts(
|
||||
await _pool(), account_id, Page(limit=limit, offset=offset), seen=seen
|
||||
)
|
||||
) -> dict[str, Any]:
|
||||
"""List fired alerts from watch rules, newest first.
|
||||
|
||||
Each alert names its watch (id, kind, name), the event (`message`, `edit`,
|
||||
`transcript`, `delete`, `online`), the chat, the message preview with its
|
||||
sender and deleted/edited flags, the peer for `peer_online`, and `path`,
|
||||
the web client link. `seen=false` keeps unseen ones. Page older with
|
||||
`before` set to `next_before` from the previous call; `unseen` is the
|
||||
account's unseen count.
|
||||
"""
|
||||
page = await watches.list_alerts(
|
||||
await _pool(), account_id, seen=seen, before=before, limit=limit
|
||||
)
|
||||
return page.model_dump(mode="json")
|
||||
@@ -2,7 +2,7 @@ import asyncio
|
||||
import json
|
||||
import logging
|
||||
from datetime import UTC, datetime
|
||||
from typing import Any
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import asyncpg
|
||||
|
||||
@@ -10,6 +10,10 @@ 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__)
|
||||
|
||||
@@ -27,8 +31,15 @@ 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._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)
|
||||
@@ -52,6 +63,39 @@ class EventHub:
|
||||
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)
|
||||
|
||||
@@ -6,55 +6,88 @@ from fastapi import APIRouter, HTTPException, Query
|
||||
from pydantic import BaseModel
|
||||
|
||||
from utils.read import watches
|
||||
from utils.read.models import DEFAULT_LIMIT, AlertView, Page, WatchView
|
||||
from utils.read.models import AlertPage, WatchItem, WatchView
|
||||
from utils.watch_rules import WatchParamsError, clean_name, normalize_params
|
||||
|
||||
router = APIRouter(prefix="/api", tags=["watches"], route_class=DishkaRoute)
|
||||
|
||||
|
||||
class WatchCreate(BaseModel):
|
||||
class WatchUpdate(BaseModel):
|
||||
kind: str | None = None
|
||||
params: dict[str, Any] = {}
|
||||
name: str | None = None
|
||||
enabled: bool = True
|
||||
|
||||
|
||||
class WatchCreate(WatchUpdate):
|
||||
account_id: int
|
||||
kind: str
|
||||
params: dict[str, Any] = {}
|
||||
enabled: bool = True
|
||||
|
||||
|
||||
class WatchUpdate(BaseModel):
|
||||
params: dict[str, Any] = {}
|
||||
enabled: bool = True
|
||||
class SeenAll(BaseModel):
|
||||
account_id: int
|
||||
up_to: int | None = None
|
||||
|
||||
|
||||
def _params(kind: str, params: dict[str, Any]) -> dict[str, Any]:
|
||||
try:
|
||||
return normalize_params(kind, params)
|
||||
except WatchParamsError as exc:
|
||||
raise HTTPException(status_code=422, detail=str(exc)) from exc
|
||||
|
||||
|
||||
async def _item(pool: asyncpg.Pool, watch: WatchView) -> WatchItem:
|
||||
return (await watches.describe_watches(pool, watch.account_id, [watch]))[0]
|
||||
|
||||
|
||||
@router.get("/watches")
|
||||
async def list_watches(
|
||||
pool: FromDishka[asyncpg.Pool], account_id: Annotated[int, Query()]
|
||||
) -> list[WatchView]:
|
||||
return await watches.list_watches(pool, account_id)
|
||||
) -> list[WatchItem]:
|
||||
rows = await watches.list_watches(pool, account_id)
|
||||
return await watches.describe_watches(pool, account_id, rows)
|
||||
|
||||
|
||||
@router.post("/watches", status_code=201)
|
||||
async def create_watch(pool: FromDishka[asyncpg.Pool], body: WatchCreate) -> WatchView:
|
||||
return await watches.create_watch(
|
||||
pool, body.account_id, body.kind, body.params, enabled=body.enabled
|
||||
async def create_watch(pool: FromDishka[asyncpg.Pool], body: WatchCreate) -> WatchItem:
|
||||
watch = await watches.create_watch(
|
||||
pool,
|
||||
body.account_id,
|
||||
body.kind,
|
||||
_params(body.kind, body.params),
|
||||
name=clean_name(body.name),
|
||||
enabled=body.enabled,
|
||||
)
|
||||
return await _item(pool, watch)
|
||||
|
||||
|
||||
@router.get("/watches/{watch_id}")
|
||||
async def get_watch(pool: FromDishka[asyncpg.Pool], watch_id: int) -> WatchView:
|
||||
async def get_watch(pool: FromDishka[asyncpg.Pool], watch_id: int) -> WatchItem:
|
||||
watch = await watches.get_watch(pool, watch_id)
|
||||
if watch is None:
|
||||
raise HTTPException(status_code=404, detail="watch not found")
|
||||
return watch
|
||||
return await _item(pool, watch)
|
||||
|
||||
|
||||
@router.put("/watches/{watch_id}")
|
||||
async def update_watch(
|
||||
pool: FromDishka[asyncpg.Pool], watch_id: int, body: WatchUpdate
|
||||
) -> WatchView:
|
||||
) -> WatchItem:
|
||||
current = await watches.get_watch(pool, watch_id)
|
||||
if current is None:
|
||||
raise HTTPException(status_code=404, detail="watch not found")
|
||||
kind = body.kind or current.kind
|
||||
watch = await watches.update_watch(
|
||||
pool, watch_id, body.params, enabled=body.enabled
|
||||
pool,
|
||||
watch_id,
|
||||
kind,
|
||||
_params(kind, body.params),
|
||||
name=clean_name(body.name),
|
||||
enabled=body.enabled,
|
||||
)
|
||||
if watch is None:
|
||||
raise HTTPException(status_code=404, detail="watch not found")
|
||||
return watch
|
||||
return await _item(pool, watch)
|
||||
|
||||
|
||||
@router.delete("/watches/{watch_id}", status_code=204)
|
||||
@@ -68,14 +101,29 @@ async def list_alerts(
|
||||
pool: FromDishka[asyncpg.Pool],
|
||||
account_id: Annotated[int, Query()],
|
||||
seen: Annotated[bool | None, Query()] = None,
|
||||
limit: Annotated[int, Query()] = DEFAULT_LIMIT,
|
||||
offset: Annotated[int, Query()] = 0,
|
||||
) -> list[AlertView]:
|
||||
before: Annotated[int | None, Query()] = None,
|
||||
limit: Annotated[int, Query()] = watches.ALERT_PAGE,
|
||||
) -> AlertPage:
|
||||
return await watches.list_alerts(
|
||||
pool, account_id, Page(limit=limit, offset=offset), seen=seen
|
||||
pool, account_id, seen=seen, before=before, limit=limit
|
||||
)
|
||||
|
||||
|
||||
@router.get("/alerts/unseen")
|
||||
async def unseen_alerts(
|
||||
pool: FromDishka[asyncpg.Pool], account_id: Annotated[int, Query()]
|
||||
) -> dict[str, int]:
|
||||
return {"unseen": await watches.unseen_count(pool, account_id)}
|
||||
|
||||
|
||||
@router.post("/alerts/seen")
|
||||
async def mark_all_seen(
|
||||
pool: FromDishka[asyncpg.Pool], body: SeenAll
|
||||
) -> dict[str, int]:
|
||||
marked = await watches.mark_all_seen(pool, body.account_id, body.up_to)
|
||||
return {"marked": marked}
|
||||
|
||||
|
||||
@router.post("/alerts/{alert_id}/seen", status_code=204)
|
||||
async def mark_alert_seen(pool: FromDishka[asyncpg.Pool], alert_id: int) -> None:
|
||||
if not await watches.mark_alert_seen(pool, alert_id):
|
||||
|
||||
@@ -2,8 +2,17 @@ from datetime import datetime
|
||||
from enum import StrEnum
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import BigInteger, Column, DateTime, LargeBinary, Text, func
|
||||
from sqlalchemy.dialects.postgresql import JSONB
|
||||
from sqlalchemy import (
|
||||
BigInteger,
|
||||
Column,
|
||||
DateTime,
|
||||
ForeignKey,
|
||||
Integer,
|
||||
LargeBinary,
|
||||
Text,
|
||||
func,
|
||||
)
|
||||
from sqlalchemy.dialects.postgresql import ARRAY, JSONB
|
||||
from sqlmodel import Field, SQLModel
|
||||
|
||||
|
||||
@@ -445,6 +454,7 @@ class Watch(SQLModel, table=True):
|
||||
id: int | None = Field(default=None, primary_key=True)
|
||||
account_id: int
|
||||
kind: str
|
||||
name: str | None = Field(default=None, sa_column=Column(Text, nullable=True))
|
||||
params: dict[str, Any] = Field(
|
||||
default_factory=dict, sa_column=Column(JSONB, nullable=False)
|
||||
)
|
||||
@@ -469,7 +479,11 @@ class Alert(SQLModel, table=True):
|
||||
|
||||
id: int | None = Field(default=None, primary_key=True)
|
||||
account_id: int
|
||||
watch_id: int = Field(foreign_key="watches.id")
|
||||
watch_id: int = Field(
|
||||
sa_column=Column(
|
||||
Integer, ForeignKey("watches.id", ondelete="CASCADE"), nullable=False
|
||||
)
|
||||
)
|
||||
ts: datetime = Field(sa_column=Column(DateTime(timezone=True), nullable=False))
|
||||
payload: dict[str, Any] = Field(
|
||||
default_factory=dict, sa_column=Column(JSONB, nullable=False)
|
||||
@@ -546,3 +560,62 @@ class FileShareHit(SQLModel, table=True):
|
||||
ip: str | None = None
|
||||
user_agent: str | None = None
|
||||
counted: bool = False
|
||||
|
||||
|
||||
class Webhook(SQLModel, table=True):
|
||||
__tablename__ = "webhooks"
|
||||
|
||||
id: int | None = Field(default=None, primary_key=True)
|
||||
account_id: int
|
||||
url: str = Field(sa_column=Column(Text, nullable=False))
|
||||
secret: str = Field(sa_column=Column(Text, nullable=False))
|
||||
enabled: bool = True
|
||||
watch_ids: list[int] | None = Field(
|
||||
default=None, sa_column=Column(ARRAY(Integer), nullable=True)
|
||||
)
|
||||
created_at: datetime = Field(
|
||||
sa_column=Column(
|
||||
DateTime(timezone=True), nullable=False, server_default=func.now()
|
||||
)
|
||||
)
|
||||
updated_at: datetime = Field(
|
||||
sa_column=Column(
|
||||
DateTime(timezone=True), nullable=False, server_default=func.now()
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class WebhookDelivery(SQLModel, table=True):
|
||||
__tablename__ = "webhook_deliveries"
|
||||
|
||||
id: int | None = Field(default=None, sa_column=Column(BigInteger, primary_key=True))
|
||||
webhook_id: int = Field(
|
||||
sa_column=Column(
|
||||
Integer, ForeignKey("webhooks.id", ondelete="CASCADE"), nullable=False
|
||||
)
|
||||
)
|
||||
alert_id: int | None = Field(
|
||||
default=None,
|
||||
sa_column=Column(
|
||||
Integer, ForeignKey("alerts.id", ondelete="CASCADE"), nullable=True
|
||||
),
|
||||
)
|
||||
event: str = Field(sa_column=Column(Text, nullable=False))
|
||||
status: str = Field(default="pending", sa_column=Column(Text, nullable=False))
|
||||
attempts: int = 0
|
||||
next_attempt_at: datetime = Field(
|
||||
sa_column=Column(
|
||||
DateTime(timezone=True), nullable=False, server_default=func.now()
|
||||
)
|
||||
)
|
||||
response_code: int | None = None
|
||||
last_error: str | None = Field(default=None, sa_column=Column(Text))
|
||||
body: str | None = Field(default=None, sa_column=Column(Text))
|
||||
created_at: datetime = Field(
|
||||
sa_column=Column(
|
||||
DateTime(timezone=True), nullable=False, server_default=func.now()
|
||||
)
|
||||
)
|
||||
finished_at: datetime | None = Field(
|
||||
default=None, sa_column=Column(DateTime(timezone=True))
|
||||
)
|
||||
@@ -523,12 +523,31 @@ class WatchView(BaseModel):
|
||||
id: int
|
||||
account_id: int
|
||||
kind: str
|
||||
name: str | None
|
||||
params: dict[str, Any]
|
||||
enabled: bool
|
||||
created_at: datetime
|
||||
updated_at: datetime
|
||||
|
||||
|
||||
class ChatRef(BaseModel):
|
||||
id: int
|
||||
title: str | None
|
||||
kind: str
|
||||
has_avatar: bool
|
||||
|
||||
|
||||
class PersonRef(BaseModel):
|
||||
id: int
|
||||
name: str | None
|
||||
has_avatar: bool
|
||||
|
||||
|
||||
class WatchItem(WatchView):
|
||||
chat: ChatRef | None
|
||||
person: PersonRef | None
|
||||
|
||||
|
||||
class AlertView(BaseModel):
|
||||
id: int
|
||||
account_id: int
|
||||
@@ -539,6 +558,65 @@ class AlertView(BaseModel):
|
||||
created_at: datetime
|
||||
|
||||
|
||||
class AlertWatch(BaseModel):
|
||||
id: int
|
||||
kind: str
|
||||
name: str | None
|
||||
|
||||
|
||||
class AlertMessage(BaseModel):
|
||||
id: int
|
||||
text: str | None
|
||||
sender: PersonRef | None
|
||||
date: datetime | None
|
||||
deleted: bool
|
||||
edited: bool
|
||||
|
||||
|
||||
class AlertItem(BaseModel):
|
||||
id: int
|
||||
account_id: int
|
||||
ts: datetime
|
||||
seen: bool
|
||||
event: str
|
||||
watch: AlertWatch
|
||||
chat: ChatRef | None
|
||||
message: AlertMessage | None
|
||||
peer: PersonRef | None
|
||||
path: str
|
||||
|
||||
|
||||
class AlertPage(BaseModel):
|
||||
items: list[AlertItem]
|
||||
next_before: int | None
|
||||
unseen: int
|
||||
|
||||
|
||||
class WebhookView(BaseModel):
|
||||
id: int
|
||||
account_id: int
|
||||
url: str
|
||||
secret: str
|
||||
enabled: bool
|
||||
watch_ids: list[int] | None
|
||||
created_at: datetime
|
||||
updated_at: datetime
|
||||
|
||||
|
||||
class DeliveryView(BaseModel):
|
||||
id: int
|
||||
webhook_id: int
|
||||
alert_id: int | None
|
||||
event: str
|
||||
status: str
|
||||
attempts: int
|
||||
next_attempt_at: datetime
|
||||
response_code: int | None
|
||||
last_error: str | None
|
||||
created_at: datetime
|
||||
finished_at: datetime | None
|
||||
|
||||
|
||||
class FileShareView(BaseModel):
|
||||
id: int
|
||||
account_id: int
|
||||
|
||||
@@ -4,13 +4,64 @@ from typing import Any
|
||||
|
||||
import asyncpg
|
||||
|
||||
from utils.read.models import AlertView, Page, WatchView
|
||||
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, params, enabled, created_at, updated_at"
|
||||
_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
|
||||
|
||||
_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,
|
||||
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:
|
||||
@@ -25,6 +76,10 @@ def _to_alert(row: asyncpg.Record) -> AlertView:
|
||||
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
|
||||
@@ -42,19 +97,21 @@ async def get_watch(pool: asyncpg.Pool, watch_id: int) -> WatchView | None:
|
||||
return _to_watch(row) if row else None
|
||||
|
||||
|
||||
async def create_watch(
|
||||
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, params, enabled) " # noqa: S608
|
||||
f"VALUES ($1, $2, $3::jsonb, $4) RETURNING {_WATCH_COLS}",
|
||||
"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,
|
||||
)
|
||||
@@ -62,13 +119,22 @@ async def create_watch(
|
||||
return _to_watch(row)
|
||||
|
||||
|
||||
async def update_watch(
|
||||
pool: asyncpg.Pool, watch_id: int, params: dict[str, Any], *, enabled: bool
|
||||
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 params = $2::jsonb, enabled = $3, updated_at = now() " # noqa: S608
|
||||
"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,
|
||||
)
|
||||
@@ -79,29 +145,198 @@ async def update_watch(
|
||||
|
||||
|
||||
async def delete_watch(pool: asyncpg.Pool, watch_id: int) -> bool:
|
||||
result = await pool.execute("DELETE FROM watches WHERE id = $1", watch_id)
|
||||
if not result.endswith("1"):
|
||||
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},
|
||||
)
|
||||
return [
|
||||
WatchItem(
|
||||
**watch.model_dump(), chat=refs.chat(chat), person=refs.person(person)
|
||||
)
|
||||
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"]
|
||||
),
|
||||
chat=refs.chat(chat_id),
|
||||
message=message,
|
||||
peer=refs.person(peer_id),
|
||||
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, page: Page, *, seen: bool | None = None
|
||||
) -> list[AlertView]:
|
||||
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 = "account_id = $1"
|
||||
where = "a.account_id = $1"
|
||||
if seen is not None:
|
||||
params.append(seen)
|
||||
where += f" AND seen = ${len(params)}"
|
||||
params.append(page.capped_limit)
|
||||
params.append(page.offset)
|
||||
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"SELECT {_ALERT_COLS} FROM alerts WHERE {where} " # noqa: S608
|
||||
f"ORDER BY ts DESC LIMIT ${len(params) - 1} OFFSET ${len(params)}",
|
||||
*params,
|
||||
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
|
||||
)
|
||||
return [_to_alert(row) for row in rows]
|
||||
|
||||
|
||||
async def insert_alert(
|
||||
@@ -113,10 +348,7 @@ async def insert_alert(
|
||||
dedup_key: str | None = None,
|
||||
) -> AlertView | None:
|
||||
row = await pool.fetchrow(
|
||||
"INSERT INTO alerts (account_id, watch_id, ts, payload, seen, dedup_key) " # noqa: S608
|
||||
"VALUES ($1, $2, $3, $4::jsonb, false, $5) "
|
||||
"ON CONFLICT (watch_id, dedup_key) WHERE dedup_key IS NOT NULL "
|
||||
f"DO NOTHING RETURNING {_ALERT_COLS}",
|
||||
_INSERT_ALERT,
|
||||
account_id,
|
||||
watch_id,
|
||||
datetime.now(UTC),
|
||||
@@ -125,10 +357,32 @@ async def insert_alert(
|
||||
)
|
||||
if row is None:
|
||||
return None
|
||||
await pool.execute(f"NOTIFY {ALERTS_CHANGED_CHANNEL}")
|
||||
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:
|
||||
result = await pool.execute("UPDATE alerts SET seen = true WHERE id = $1", alert_id)
|
||||
return result.endswith("1")
|
||||
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
|
||||
Reference in new issue
Block a user