454 lines
14 KiB
Python
454 lines
14 KiB
Python
from dataclasses import dataclass
|
|
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, CASE WHEN $10 THEN now() END)
|
|
ON CONFLICT (account_id, chat_id, message_id, date) DO UPDATE SET
|
|
text = EXCLUDED.text,
|
|
raw = EXCLUDED.raw,
|
|
has_media = EXCLUDED.has_media,
|
|
is_self_destruct = EXCLUDED.is_self_destruct,
|
|
edited_at = coalesce(EXCLUDED.edited_at, messages.edited_at)
|
|
"""
|
|
|
|
_LOCK_MESSAGE = """
|
|
SELECT pg_advisory_xact_lock(
|
|
hashtextextended(concat_ws(':', 'message', $1::int, $2::bigint, $3::bigint), 0)
|
|
)
|
|
"""
|
|
|
|
_FIND_DELETED = """
|
|
SELECT chat_id, message_id, date FROM messages
|
|
WHERE account_id = $1 AND message_id = ANY($2::bigint[]) AND deleted_at IS NULL
|
|
"""
|
|
|
|
_MARK_DELETED = """
|
|
UPDATE messages SET deleted_at = now()
|
|
WHERE account_id = $1 AND chat_id = $2 AND message_id = $3 AND date = $4
|
|
"""
|
|
|
|
_MARK_DELETED_LATE = """
|
|
UPDATE messages SET deleted_at = now()
|
|
WHERE account_id = $1 AND chat_id = $2 AND message_id = $3 AND date = $4
|
|
AND deleted_at IS NULL
|
|
"""
|
|
|
|
_INSERT_VERSION = """
|
|
INSERT INTO message_versions
|
|
(account_id, chat_id, message_id, observed_at, edit_date, text, raw)
|
|
VALUES ($1, $2, $3, clock_timestamp(), $4, $5, $6::jsonb)
|
|
ON CONFLICT DO NOTHING
|
|
"""
|
|
|
|
_SNAPSHOT_ORIGINAL = """
|
|
INSERT INTO message_versions
|
|
(account_id, chat_id, message_id, observed_at, edit_date, text, raw)
|
|
SELECT account_id, chat_id, message_id, clock_timestamp(), NULL, text, raw
|
|
FROM messages m
|
|
WHERE m.account_id = $1 AND m.chat_id = $2 AND m.message_id = $3 AND m.date = $4
|
|
AND NOT EXISTS (
|
|
SELECT 1 FROM message_versions v
|
|
WHERE v.account_id = m.account_id AND v.chat_id = m.chat_id
|
|
AND v.message_id = m.message_id
|
|
)
|
|
ON CONFLICT DO NOTHING
|
|
"""
|
|
|
|
_INSERT_PAST_VERSION = """
|
|
INSERT INTO message_versions
|
|
(account_id, chat_id, message_id, observed_at, edit_date, text, raw)
|
|
SELECT $1::int, $2::bigint, $3::bigint,
|
|
coalesce(
|
|
(SELECT min(v.observed_at) - interval '1 microsecond'
|
|
FROM message_versions v
|
|
WHERE v.account_id = $1 AND v.chat_id = $2 AND v.message_id = $3
|
|
AND v.edit_date > $4::timestamptz),
|
|
$4::timestamptz,
|
|
$5::timestamptz
|
|
),
|
|
$4::timestamptz, $6::text, $7::jsonb
|
|
WHERE NOT EXISTS (
|
|
SELECT 1 FROM message_versions v
|
|
WHERE v.account_id = $1 AND v.chat_id = $2 AND v.message_id = $3
|
|
AND v.edit_date IS NOT DISTINCT FROM $4::timestamptz
|
|
AND v.text IS NOT DISTINCT FROM $6::text
|
|
)
|
|
ON CONFLICT DO NOTHING
|
|
"""
|
|
|
|
_CURRENT_CONTENT = """
|
|
SELECT
|
|
m.text AS text,
|
|
m.raw->>'edit_date' AS edit_date,
|
|
(SELECT d.unique_id FROM media d
|
|
WHERE d.account_id = m.account_id AND d.chat_id = m.chat_id
|
|
AND d.message_id = m.message_id) AS media_unique_id
|
|
FROM messages m
|
|
WHERE m.account_id = $1 AND m.chat_id = $2 AND m.message_id = $3 AND m.date = $4
|
|
"""
|
|
|
|
_CURRENT_MEDIA = """
|
|
SELECT unique_id, storage_key, file_size, downloaded FROM media
|
|
WHERE account_id = $1 AND chat_id = $2 AND message_id = $3
|
|
"""
|
|
|
|
_INSERT_MEDIA_VERSION = """
|
|
INSERT INTO media_versions
|
|
(account_id, chat_id, message_id, observed_at, kind, storage_key,
|
|
file_size, mime, ttl_seconds, unique_id)
|
|
VALUES ($1, $2, $3, clock_timestamp(), $4, $5, $6, $7, $8, $9)
|
|
ON CONFLICT (account_id, chat_id, message_id, storage_key) DO NOTHING
|
|
"""
|
|
|
|
_INSERT_MEDIA = """
|
|
INSERT INTO media
|
|
(account_id, chat_id, message_id, kind, storage_key, file_size, mime,
|
|
ttl_seconds, downloaded, unique_id, evicted_at)
|
|
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, CASE WHEN $11 THEN now() END)
|
|
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,
|
|
unique_id = EXCLUDED.unique_id,
|
|
evicted_at = CASE
|
|
WHEN EXCLUDED.downloaded THEN NULL
|
|
WHEN media.unique_id IS DISTINCT FROM EXCLUDED.unique_id
|
|
THEN EXCLUDED.evicted_at
|
|
ELSE coalesce(media.evicted_at, EXCLUDED.evicted_at) END,
|
|
evicted_key = CASE WHEN EXCLUDED.downloaded
|
|
OR media.unique_id IS DISTINCT FROM EXCLUDED.unique_id
|
|
THEN NULL ELSE media.evicted_key END
|
|
"""
|
|
|
|
_IS_EVICTED = """
|
|
SELECT EXISTS (
|
|
SELECT 1 FROM media WHERE unique_id = $1 AND evicted_at IS NOT NULL
|
|
)
|
|
"""
|
|
|
|
_UNEVICT = """
|
|
UPDATE media SET storage_key = $2, file_size = $3, downloaded = true,
|
|
evicted_at = NULL, evicted_key = NULL
|
|
WHERE evicted_key = $2 OR (unique_id = $1 AND evicted_at IS NOT NULL)
|
|
"""
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class MessageRow:
|
|
chat_id: int
|
|
message_id: int
|
|
date: datetime
|
|
sender_id: int | None
|
|
text: str | None
|
|
raw: str
|
|
has_media: bool
|
|
is_self_destruct: bool
|
|
edit_date: datetime | None
|
|
media_unique_id: str | None
|
|
|
|
@property
|
|
def columns(self) -> tuple:
|
|
return (
|
|
self.chat_id,
|
|
self.message_id,
|
|
self.date,
|
|
self.sender_id,
|
|
self.text,
|
|
self.raw,
|
|
self.has_media,
|
|
self.is_self_destruct,
|
|
)
|
|
|
|
|
|
async def _lock_current(
|
|
conn: asyncpg.Connection, account_id: int, row: MessageRow
|
|
) -> asyncpg.Record | None:
|
|
await conn.execute(_LOCK_MESSAGE, account_id, row.chat_id, row.message_id)
|
|
return await conn.fetchrow(
|
|
_CURRENT_CONTENT, account_id, row.chat_id, row.message_id, row.date
|
|
)
|
|
|
|
|
|
def _is_stale(current: asyncpg.Record, row: MessageRow) -> bool:
|
|
stored = current["edit_date"]
|
|
if stored is None:
|
|
return False
|
|
if row.edit_date is None:
|
|
return True
|
|
try:
|
|
return row.edit_date < datetime.fromisoformat(stored)
|
|
except (TypeError, ValueError):
|
|
return False
|
|
|
|
|
|
async def _record_past(
|
|
conn: asyncpg.Connection, account_id: int, row: MessageRow
|
|
) -> None:
|
|
await conn.execute(
|
|
_INSERT_PAST_VERSION,
|
|
account_id,
|
|
row.chat_id,
|
|
row.message_id,
|
|
row.edit_date,
|
|
row.date,
|
|
row.text,
|
|
row.raw,
|
|
)
|
|
|
|
|
|
async def _record_edit(
|
|
conn: asyncpg.Connection, account_id: int, row: MessageRow, *, mark_edited: bool
|
|
) -> None:
|
|
await conn.execute(
|
|
_SNAPSHOT_ORIGINAL, account_id, row.chat_id, row.message_id, row.date
|
|
)
|
|
await conn.execute(_TOUCH_EDITED, account_id, *row.columns, mark_edited)
|
|
await conn.execute(
|
|
_INSERT_VERSION,
|
|
account_id,
|
|
row.chat_id,
|
|
row.message_id,
|
|
row.edit_date,
|
|
row.text,
|
|
row.raw,
|
|
)
|
|
|
|
|
|
async def upsert_message(pool: asyncpg.Pool, account_id: int, row: MessageRow) -> bool:
|
|
async with pool.acquire() as conn, conn.transaction():
|
|
current = await _lock_current(conn, account_id, row)
|
|
if current is None:
|
|
await conn.execute(_UPSERT_MESSAGE, account_id, *row.columns)
|
|
return True
|
|
stored_media = current["media_unique_id"]
|
|
changed = current["text"] != row.text or (
|
|
stored_media is not None and stored_media != row.media_unique_id
|
|
)
|
|
if _is_stale(current, row):
|
|
if changed:
|
|
await _record_past(conn, account_id, row)
|
|
return False
|
|
if changed:
|
|
await _record_edit(
|
|
conn, account_id, row, mark_edited=row.edit_date is not None
|
|
)
|
|
else:
|
|
await conn.execute(_UPSERT_MESSAGE, account_id, *row.columns)
|
|
return True
|
|
|
|
|
|
async def max_message_id(
|
|
pool: asyncpg.Pool, account_id: int, chat_id: int
|
|
) -> int | None:
|
|
return await pool.fetchval(
|
|
"SELECT last_message_id FROM chat_stats WHERE account_id = $1 AND chat_id = $2",
|
|
account_id,
|
|
chat_id,
|
|
)
|
|
|
|
|
|
async def _mark_deleted(
|
|
pool: asyncpg.Pool, account_id: int, rows: list[asyncpg.Record]
|
|
) -> None:
|
|
for row in rows:
|
|
await pool.execute(
|
|
_MARK_DELETED, account_id, row["chat_id"], row["message_id"], row["date"]
|
|
)
|
|
|
|
|
|
async def mark_deleted_box(
|
|
pool: asyncpg.Pool, account_id: int, message_ids: list[int]
|
|
) -> None:
|
|
rows = await pool.fetch(
|
|
f"{_FIND_DELETED} AND chat_id > $3",
|
|
account_id,
|
|
message_ids,
|
|
CHANNEL_ID_THRESHOLD,
|
|
)
|
|
await _mark_deleted(pool, account_id, rows)
|
|
|
|
|
|
async def mark_deleted_channel(
|
|
pool: asyncpg.Pool, account_id: int, chat_id: int, message_ids: list[int]
|
|
) -> None:
|
|
rows = await pool.fetch(
|
|
f"{_FIND_DELETED} AND chat_id = $3", account_id, message_ids, chat_id
|
|
)
|
|
await _mark_deleted(pool, account_id, rows)
|
|
|
|
|
|
async def add_version(pool: asyncpg.Pool, account_id: int, row: MessageRow) -> bool:
|
|
async with pool.acquire() as conn, conn.transaction():
|
|
current = await _lock_current(conn, account_id, row)
|
|
if current is not None:
|
|
changed = (
|
|
current["text"] != row.text
|
|
or current["media_unique_id"] != row.media_unique_id
|
|
)
|
|
if not changed:
|
|
return False
|
|
if _is_stale(current, row):
|
|
await _record_past(conn, account_id, row)
|
|
return False
|
|
await _record_edit(conn, account_id, row, mark_edited=True)
|
|
return True
|
|
|
|
|
|
async def mark_deleted_late(
|
|
pool: asyncpg.Pool, account_id: int, chat_id: int, message_id: int, date: datetime
|
|
) -> None:
|
|
await pool.execute(_MARK_DELETED_LATE, account_id, chat_id, message_id, date)
|
|
|
|
|
|
async def current_media(
|
|
pool: asyncpg.Pool, account_id: int, chat_id: int, message_id: int
|
|
) -> asyncpg.Record | None:
|
|
return await pool.fetchrow(_CURRENT_MEDIA, account_id, chat_id, message_id)
|
|
|
|
|
|
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,
|
|
unique_id: str | None,
|
|
*,
|
|
downloaded: bool,
|
|
evicted: bool = False,
|
|
) -> None:
|
|
async with pool.acquire() as conn, conn.transaction():
|
|
await conn.execute(
|
|
_INSERT_MEDIA,
|
|
account_id,
|
|
chat_id,
|
|
message_id,
|
|
kind,
|
|
storage_key,
|
|
file_size,
|
|
mime,
|
|
ttl_seconds,
|
|
downloaded,
|
|
unique_id,
|
|
evicted,
|
|
)
|
|
if storage_key is not None:
|
|
await conn.execute(
|
|
_INSERT_MEDIA_VERSION,
|
|
account_id,
|
|
chat_id,
|
|
message_id,
|
|
kind,
|
|
storage_key,
|
|
file_size,
|
|
mime,
|
|
ttl_seconds,
|
|
unique_id,
|
|
)
|
|
|
|
|
|
async def is_evicted(pool: asyncpg.Pool, unique_id: str | None) -> bool:
|
|
return unique_id is not None and await pool.fetchval(_IS_EVICTED, unique_id)
|
|
|
|
|
|
async def unevict(
|
|
pool: asyncpg.Pool, unique_id: str | None, storage_key: str, file_size: int | None
|
|
) -> bool:
|
|
async with pool.acquire() as conn, conn.transaction():
|
|
await conn.execute(_UNEVICT, unique_id, storage_key, file_size)
|
|
dropped = await conn.fetchval(
|
|
"DELETE FROM media_trash WHERE storage_key = $1 RETURNING storage_key",
|
|
storage_key,
|
|
)
|
|
return dropped is not None
|
|
|
|
|
|
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
|
|
],
|
|
)
|