317 lines
11 KiB
Python
317 lines
11 KiB
Python
import asyncpg
|
|
|
|
from utils.policy.models import FolderSpec
|
|
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,
|
|
MediaRef,
|
|
MessageVersionView,
|
|
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, "
|
|
"edited_at, deleted_at, raw, raw->>'media_group_id' AS media_group_id"
|
|
)
|
|
|
|
|
|
async def _media_map(
|
|
pool: asyncpg.Pool, account_id: int, rows: list[asyncpg.Record]
|
|
) -> dict[tuple[int, int], asyncpg.Record]:
|
|
message_ids = list({row["message_id"] for row in rows})
|
|
if not message_ids:
|
|
return {}
|
|
media_rows = await pool.fetch(
|
|
"SELECT id, chat_id, message_id, kind, downloaded, mime, file_size, "
|
|
"ttl_seconds, extracted_text FROM media "
|
|
"WHERE account_id = $1 AND message_id = ANY($2::bigint[])",
|
|
account_id,
|
|
message_ids,
|
|
)
|
|
return {(row["chat_id"], row["message_id"]): row for row in media_rows}
|
|
|
|
|
|
def _single_media(
|
|
row: asyncpg.Record, raw: dict, media_by_key: dict[tuple[int, int], asyncpg.Record]
|
|
) -> list[MediaRef]:
|
|
media_row = media_by_key.get((row["chat_id"], row["message_id"]))
|
|
if not (row["has_media"] or media_row):
|
|
return []
|
|
ref = media_ref_from(row["message_id"], raw, media_row)
|
|
return [ref] if ref else []
|
|
|
|
|
|
_ALL_IDS = """
|
|
SELECT chat_id FROM chat_stats WHERE account_id = $1
|
|
UNION
|
|
SELECT chat_id FROM dialogs WHERE account_id = $1
|
|
UNION
|
|
SELECT scope_id AS chat_id FROM capture_policy
|
|
WHERE account_id = $1 AND scope_type = 'chat' AND scope_id IS NOT NULL
|
|
"""
|
|
|
|
_ONE_ID = """
|
|
SELECT chat_id FROM chat_stats WHERE account_id = $1 AND chat_id = $2
|
|
UNION
|
|
SELECT chat_id FROM dialogs WHERE account_id = $1 AND chat_id = $2
|
|
UNION
|
|
SELECT scope_id AS chat_id FROM capture_policy
|
|
WHERE account_id = $1 AND scope_type = 'chat' AND scope_id = $2
|
|
"""
|
|
|
|
_CHAT_ROWS = """
|
|
WITH ids AS ({ids})
|
|
SELECT ids.chat_id,
|
|
COALESCE(cs.message_count, 0) AS message_count,
|
|
cs.last_date, cs.last_text, cs.last_sender_id,
|
|
COALESCE(named.title,
|
|
NULLIF(trim(concat_ws(' ', p.first_name, p.last_name)), ''),
|
|
p.username) AS title,
|
|
COALESCE(typed.is_broadcast, false) AS is_broadcast,
|
|
COALESCE((p.raw->>'is_bot')::bool, (p.raw->>'bot')::bool,
|
|
p.raw->>'type' = 'ChatType.BOT', false) AS is_bot,
|
|
COALESCE((p.raw->>'is_contact')::bool, (p.raw->>'contact')::bool,
|
|
false) AS is_contact,
|
|
EXISTS (SELECT 1 FROM avatars a
|
|
WHERE a.account_id = $1 AND a.owner_id = ids.chat_id) AS has_avatar
|
|
FROM ids
|
|
LEFT JOIN chat_stats cs ON cs.account_id = $1 AND cs.chat_id = ids.chat_id
|
|
LEFT JOIN peers p ON p.account_id = $1 AND p.peer_id = ids.chat_id
|
|
LEFT JOIN LATERAL (
|
|
SELECT ch.title FROM chat_history ch
|
|
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
|
|
) named ON true
|
|
LEFT JOIN LATERAL (
|
|
SELECT COALESCE(ch.raw->'chat'->>'type', ch.raw->>'type')
|
|
= 'ChatType.CHANNEL' AS is_broadcast
|
|
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
|
|
) typed ON true
|
|
"""
|
|
|
|
|
|
def _folder_filter(base: int) -> str:
|
|
return (
|
|
f"NOT (chat.chat_id = ANY(${base + 1}::bigint[])) "
|
|
f"AND (chat.chat_id = ANY(${base + 2}::bigint[]) "
|
|
f"OR (NOT ${base + 3}::bool AND CASE "
|
|
f"WHEN chat.is_broadcast THEN ${base + 4}::bool "
|
|
f"WHEN chat.chat_id < 0 THEN ${base + 5}::bool "
|
|
f"WHEN chat.is_bot THEN ${base + 6}::bool "
|
|
f"WHEN chat.is_contact THEN ${base + 7}::bool "
|
|
f"ELSE ${base + 8}::bool END))"
|
|
)
|
|
|
|
|
|
def _folder_params(folder: FolderSpec) -> list[object]:
|
|
return [
|
|
sorted(folder.exclude_ids),
|
|
sorted(folder.include_ids | folder.pinned_ids),
|
|
folder.is_chatlist,
|
|
folder.broadcasts,
|
|
folder.groups,
|
|
folder.bots,
|
|
folder.contacts,
|
|
folder.non_contacts,
|
|
]
|
|
|
|
|
|
def _chat_item(row: asyncpg.Record) -> ChatListItem:
|
|
return ChatListItem(
|
|
chat_id=row["chat_id"],
|
|
title=row["title"],
|
|
kind="private" if row["chat_id"] > 0 else "group",
|
|
has_avatar=row["has_avatar"],
|
|
is_bot=bool(row["is_bot"]),
|
|
is_contact=bool(row["is_contact"]),
|
|
is_broadcast=bool(row["is_broadcast"]),
|
|
message_count=row["message_count"],
|
|
last_date=row["last_date"],
|
|
last_text=row["last_text"],
|
|
last_sender_id=row["last_sender_id"],
|
|
)
|
|
|
|
|
|
async def list_chats(
|
|
pool: asyncpg.Pool,
|
|
account_id: int,
|
|
page: Page,
|
|
*,
|
|
folder: FolderSpec | None = None,
|
|
search: str | None = None,
|
|
) -> list[ChatListItem]:
|
|
params: list[object] = [account_id, page.capped_limit, page.offset]
|
|
rows_sql = _CHAT_ROWS.format(ids=_ALL_IDS)
|
|
clauses: list[str] = []
|
|
if folder is not None:
|
|
clauses.append(_folder_filter(len(params)))
|
|
params.extend(_folder_params(folder))
|
|
if search:
|
|
params.append(f"%{search}%")
|
|
clauses.append(f"chat.title ILIKE ${len(params)}")
|
|
where = f" WHERE {' AND '.join(clauses)}" if clauses else ""
|
|
query = (
|
|
f"SELECT chat.* FROM ({rows_sql}) chat{where} " # noqa: S608
|
|
"ORDER BY last_date DESC NULLS LAST, chat_id DESC LIMIT $2 OFFSET $3"
|
|
)
|
|
rows = await pool.fetch(query, *params)
|
|
return [_chat_item(row) for row in rows]
|
|
|
|
|
|
async def get_chat(
|
|
pool: asyncpg.Pool, account_id: int, chat_id: int
|
|
) -> ChatListItem | None:
|
|
row = await pool.fetchrow(_CHAT_ROWS.format(ids=_ONE_ID), account_id, chat_id)
|
|
return _chat_item(row) if row is not None else None
|
|
|
|
|
|
async def get_chat_history( # noqa: PLR0913
|
|
pool: asyncpg.Pool,
|
|
account_id: int,
|
|
chat_id: int,
|
|
page: Page,
|
|
*,
|
|
include_deleted: bool = True,
|
|
before_id: int | None = None,
|
|
after_id: int | None = None,
|
|
) -> list[MessageView]:
|
|
where = "account_id = $1 AND chat_id = $2"
|
|
if not include_deleted:
|
|
where += " AND deleted_at IS NULL"
|
|
params: list[object] = [account_id, chat_id]
|
|
if after_id is not None:
|
|
params.append(after_id)
|
|
where += f" AND message_id > ${len(params)}"
|
|
order = "date ASC, message_id ASC"
|
|
elif before_id is not None:
|
|
params.append(before_id)
|
|
where += f" AND message_id < ${len(params)}"
|
|
order = "date DESC, message_id DESC"
|
|
else:
|
|
order = "date DESC, message_id DESC"
|
|
params.append(page.capped_limit)
|
|
query = (
|
|
f"SELECT {_MESSAGE_COLS} FROM messages WHERE {where} " # noqa: S608
|
|
f"ORDER BY {order} LIMIT ${len(params)}"
|
|
)
|
|
if before_id is None and after_id is None:
|
|
params.append(page.offset)
|
|
query += f" OFFSET ${len(params)}"
|
|
rows = await pool.fetch(query, *params)
|
|
media_by_key = await _media_map(pool, account_id, rows)
|
|
parsed = [(row, load_raw(row["raw"])) for row in rows]
|
|
views: list[MessageView] = []
|
|
index = 0
|
|
while index < len(parsed):
|
|
group_id = parsed[index][0]["media_group_id"]
|
|
end = index + 1
|
|
if group_id is not None:
|
|
while end < len(parsed) and parsed[end][0]["media_group_id"] == group_id:
|
|
end += 1
|
|
members = parsed[index:end]
|
|
if len(members) == 1:
|
|
row, raw = members[0]
|
|
views.append(
|
|
build_message_view(row, raw, _single_media(row, raw, media_by_key))
|
|
)
|
|
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],
|
|
) -> MessageView:
|
|
ordered = sorted(members, key=lambda m: m[0]["message_id"])
|
|
media: list[MediaRef] = []
|
|
for row, raw in ordered:
|
|
media_row = media_by_key.get((row["chat_id"], row["message_id"]))
|
|
ref = media_ref_from(row["message_id"], raw, media_row)
|
|
if ref:
|
|
media.append(ref)
|
|
primary_row, primary_raw = next(
|
|
((row, raw) for row, raw in ordered if row["text"]), ordered[0]
|
|
)
|
|
return build_message_view(primary_row, primary_raw, media)
|
|
|
|
|
|
async def get_deleted_messages(
|
|
pool: asyncpg.Pool, account_id: int, page: Page, *, chat_id: int | None = None
|
|
) -> list[MessageView]:
|
|
params: list[object] = [account_id]
|
|
where = "account_id = $1 AND deleted_at IS NOT NULL"
|
|
if chat_id is not None:
|
|
params.append(chat_id)
|
|
where += f" AND chat_id = ${len(params)}"
|
|
params.append(page.capped_limit)
|
|
params.append(page.offset)
|
|
rows = await pool.fetch(
|
|
f"SELECT {_MESSAGE_COLS} FROM messages WHERE {where} " # noqa: S608
|
|
f"ORDER BY deleted_at DESC LIMIT ${len(params) - 1} OFFSET ${len(params)}",
|
|
*params,
|
|
)
|
|
media_by_key = await _media_map(pool, account_id, rows)
|
|
views: list[MessageView] = []
|
|
for row in rows:
|
|
raw = load_raw(row["raw"])
|
|
views.append(
|
|
build_message_view(row, raw, _single_media(row, raw, media_by_key))
|
|
)
|
|
return views
|
|
|
|
|
|
async def get_message_versions(
|
|
pool: asyncpg.Pool, account_id: int, chat_id: int, message_id: int
|
|
) -> list[MessageVersionView]:
|
|
rows = await pool.fetch(
|
|
"SELECT observed_at, edit_date, text FROM message_versions "
|
|
"WHERE account_id = $1 AND chat_id = $2 AND message_id = $3 "
|
|
"ORDER BY observed_at",
|
|
account_id,
|
|
chat_id,
|
|
message_id,
|
|
)
|
|
return [MessageVersionView(**dict(row)) for row in rows]
|