feat: add realtime capturing
This commit is contained in:
hh committed
1 parent
920a0235e2
commit
3c1a12750c
29 files changed
+967
-47
No files matched your search
@@ -0,0 +1,223 @@
|
||||
from datetime import datetime
|
||||
|
||||
import asyncpg
|
||||
|
||||
CHANNEL_ID_THRESHOLD = -1000000000000
|
||||
|
||||
_UPSERT_MESSAGE = """
|
||||
INSERT INTO messages
|
||||
(account_id, chat_id, message_id, date, sender_id, text, raw,
|
||||
has_media, is_self_destruct)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7::jsonb, $8, $9)
|
||||
ON CONFLICT (account_id, chat_id, message_id, date) DO UPDATE SET
|
||||
sender_id = EXCLUDED.sender_id,
|
||||
text = EXCLUDED.text,
|
||||
raw = EXCLUDED.raw,
|
||||
has_media = EXCLUDED.has_media,
|
||||
is_self_destruct = EXCLUDED.is_self_destruct
|
||||
"""
|
||||
|
||||
_TOUCH_EDITED = """
|
||||
INSERT INTO messages
|
||||
(account_id, chat_id, message_id, date, sender_id, text, raw,
|
||||
has_media, is_self_destruct, edited_at)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7::jsonb, $8, $9, now())
|
||||
ON CONFLICT (account_id, chat_id, message_id, date) DO UPDATE SET
|
||||
edited_at = now()
|
||||
"""
|
||||
|
||||
_INSERT_VERSION = """
|
||||
INSERT INTO message_versions
|
||||
(account_id, chat_id, message_id, observed_at, edit_date, text, raw)
|
||||
VALUES ($1, $2, $3, now(), $4, $5, $6::jsonb)
|
||||
ON CONFLICT DO NOTHING
|
||||
"""
|
||||
|
||||
_INSERT_MEDIA = """
|
||||
INSERT INTO media
|
||||
(account_id, chat_id, message_id, kind, storage_key, file_size, mime,
|
||||
ttl_seconds, downloaded)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)
|
||||
ON CONFLICT (account_id, chat_id, message_id) DO UPDATE SET
|
||||
kind = EXCLUDED.kind,
|
||||
storage_key = EXCLUDED.storage_key,
|
||||
file_size = EXCLUDED.file_size,
|
||||
mime = EXCLUDED.mime,
|
||||
ttl_seconds = EXCLUDED.ttl_seconds,
|
||||
downloaded = EXCLUDED.downloaded
|
||||
"""
|
||||
|
||||
|
||||
async def upsert_message( # noqa: PLR0913
|
||||
pool: asyncpg.Pool,
|
||||
account_id: int,
|
||||
chat_id: int,
|
||||
message_id: int,
|
||||
date: datetime,
|
||||
sender_id: int | None,
|
||||
text: str | None,
|
||||
raw: str,
|
||||
*,
|
||||
has_media: bool,
|
||||
is_self_destruct: bool,
|
||||
) -> None:
|
||||
await pool.execute(
|
||||
_UPSERT_MESSAGE,
|
||||
account_id,
|
||||
chat_id,
|
||||
message_id,
|
||||
date,
|
||||
sender_id,
|
||||
text,
|
||||
raw,
|
||||
has_media,
|
||||
is_self_destruct,
|
||||
)
|
||||
|
||||
|
||||
async def mark_deleted_box(
|
||||
pool: asyncpg.Pool, account_id: int, message_ids: list[int]
|
||||
) -> None:
|
||||
await pool.execute(
|
||||
"UPDATE messages SET deleted_at = now() "
|
||||
"WHERE account_id = $1 AND message_id = ANY($2::bigint[]) "
|
||||
"AND chat_id > $3 AND deleted_at IS NULL",
|
||||
account_id,
|
||||
message_ids,
|
||||
CHANNEL_ID_THRESHOLD,
|
||||
)
|
||||
|
||||
|
||||
async def mark_deleted_channel(
|
||||
pool: asyncpg.Pool, account_id: int, chat_id: int, message_ids: list[int]
|
||||
) -> None:
|
||||
await pool.execute(
|
||||
"UPDATE messages SET deleted_at = now() "
|
||||
"WHERE account_id = $1 AND chat_id = $2 AND message_id = ANY($3::bigint[]) "
|
||||
"AND deleted_at IS NULL",
|
||||
account_id,
|
||||
chat_id,
|
||||
message_ids,
|
||||
)
|
||||
|
||||
|
||||
async def add_version( # noqa: PLR0913
|
||||
pool: asyncpg.Pool,
|
||||
account_id: int,
|
||||
chat_id: int,
|
||||
message_id: int,
|
||||
date: datetime,
|
||||
sender_id: int | None,
|
||||
text: str | None,
|
||||
raw: str,
|
||||
edit_date: datetime | None,
|
||||
*,
|
||||
has_media: bool,
|
||||
is_self_destruct: bool,
|
||||
) -> None:
|
||||
async with pool.acquire() as conn, conn.transaction():
|
||||
await conn.execute(
|
||||
_TOUCH_EDITED,
|
||||
account_id,
|
||||
chat_id,
|
||||
message_id,
|
||||
date,
|
||||
sender_id,
|
||||
text,
|
||||
raw,
|
||||
has_media,
|
||||
is_self_destruct,
|
||||
)
|
||||
await conn.execute(
|
||||
_INSERT_VERSION, account_id, chat_id, message_id, edit_date, text, raw
|
||||
)
|
||||
|
||||
|
||||
async def insert_media( # noqa: PLR0913
|
||||
pool: asyncpg.Pool,
|
||||
account_id: int,
|
||||
chat_id: int,
|
||||
message_id: int,
|
||||
kind: str,
|
||||
storage_key: str | None,
|
||||
file_size: int | None,
|
||||
mime: str | None,
|
||||
ttl_seconds: int | None,
|
||||
*,
|
||||
downloaded: bool,
|
||||
) -> None:
|
||||
await pool.execute(
|
||||
_INSERT_MEDIA,
|
||||
account_id,
|
||||
chat_id,
|
||||
message_id,
|
||||
kind,
|
||||
storage_key,
|
||||
file_size,
|
||||
mime,
|
||||
ttl_seconds,
|
||||
downloaded,
|
||||
)
|
||||
|
||||
|
||||
async def insert_callbacks(
|
||||
pool: asyncpg.Pool,
|
||||
account_id: int,
|
||||
chat_id: int,
|
||||
message_id: int,
|
||||
buttons: list[tuple[int, str | None, bytes | None]],
|
||||
) -> None:
|
||||
async with pool.acquire() as conn, conn.transaction():
|
||||
await conn.execute(
|
||||
"DELETE FROM callbacks "
|
||||
"WHERE account_id = $1 AND chat_id = $2 AND message_id = $3",
|
||||
account_id,
|
||||
chat_id,
|
||||
message_id,
|
||||
)
|
||||
await conn.executemany(
|
||||
"INSERT INTO callbacks "
|
||||
"(account_id, chat_id, message_id, position, label, data) "
|
||||
"VALUES ($1, $2, $3, $4, $5, $6)",
|
||||
[
|
||||
(account_id, chat_id, message_id, position, label, data)
|
||||
for position, label, data in buttons
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
async def sync_reactions(
|
||||
pool: asyncpg.Pool,
|
||||
account_id: int,
|
||||
chat_id: int,
|
||||
message_id: int,
|
||||
current: list[tuple[int, str]],
|
||||
) -> None:
|
||||
async with pool.acquire() as conn, conn.transaction():
|
||||
rows = await conn.fetch(
|
||||
"SELECT id, peer_id, reaction FROM reactions "
|
||||
"WHERE account_id = $1 AND chat_id = $2 AND message_id = $3 "
|
||||
"AND removed_at IS NULL",
|
||||
account_id,
|
||||
chat_id,
|
||||
message_id,
|
||||
)
|
||||
existing = {(row["peer_id"], row["reaction"]): row["id"] for row in rows}
|
||||
current_set = set(current)
|
||||
stale = [pid for key, pid in existing.items() if key not in current_set]
|
||||
if stale:
|
||||
await conn.execute(
|
||||
"UPDATE reactions SET removed_at = now() WHERE id = ANY($1::bigint[])",
|
||||
stale,
|
||||
)
|
||||
fresh = [key for key in current_set if key not in existing]
|
||||
if fresh:
|
||||
await conn.executemany(
|
||||
"INSERT INTO reactions "
|
||||
"(account_id, chat_id, message_id, peer_id, reaction) "
|
||||
"VALUES ($1, $2, $3, $4, $5)",
|
||||
[
|
||||
(account_id, chat_id, message_id, peer_id, reaction)
|
||||
for peer_id, reaction in fresh
|
||||
],
|
||||
)
|
||||
Reference in new issue
Block a user