188 lines
6.2 KiB
Python
188 lines
6.2 KiB
Python
import asyncpg
|
|
|
|
from utils.search.models import (
|
|
MAX_LIMIT,
|
|
ChatKind,
|
|
SearchFilters,
|
|
SearchHit,
|
|
SearchPage,
|
|
decode_cursor,
|
|
encode_cursor,
|
|
)
|
|
|
|
_PARSE = """
|
|
SELECT (regexp_replace(websearch_to_tsquery('russian', q)::text,
|
|
'''((?:[^'']|''''){3,})''', '''\\1'':*', 'g')::tsquery
|
|
|| regexp_replace(websearch_to_tsquery('simple', q)::text,
|
|
'''((?:[^'']|''''){3,})''', '''\\1'':*', 'g')::tsquery
|
|
)::text AS tsquery,
|
|
tsvector_to_array(search_vector(q)) AS terms
|
|
FROM translate($1, 'Ёё', 'Ее') AS q
|
|
"""
|
|
|
|
_CHANNELS = """
|
|
SELECT coalesce(array_agg(chat_id), '{}'::bigint[]) FROM (
|
|
SELECT DISTINCT ON (chat_id) chat_id,
|
|
coalesce(raw->'chat'->>'type', raw->>'type') AS type
|
|
FROM chat_history
|
|
WHERE account_id = $1
|
|
AND coalesce(raw->'chat'->>'type', raw->>'type') IS NOT NULL
|
|
ORDER BY chat_id, ts DESC
|
|
) typed WHERE type = 'ChatType.CHANNEL'
|
|
"""
|
|
|
|
_KEY_ORDER = "s.date DESC, s.chat_id DESC, s.message_id DESC"
|
|
|
|
_COLS = (
|
|
"s.chat_id, s.message_id, s.date, s.sender_id, s.text, s.has_media, "
|
|
"s.edited_at, s.deleted_at"
|
|
)
|
|
|
|
_TEXT = (
|
|
f"SELECT {_COLS}, 0 AS rank, NULL::text AS extracted_text, " # noqa: S608
|
|
"NULL::text AS version_text, NULL::timestamptz AS version_at "
|
|
"FROM message_search s WHERE {where}"
|
|
)
|
|
|
|
_PROBE = (
|
|
"CROSS JOIN LATERAL (SELECT * FROM message_search s "
|
|
"WHERE s.chat_id = {source}.chat_id AND s.message_id = {source}.message_id "
|
|
"AND {where} LIMIT 1) s "
|
|
)
|
|
|
|
_STT = (
|
|
f"SELECT {_COLS}, 1 AS rank, md.extracted_text, NULL::text, " # noqa: S608
|
|
"NULL::timestamptz FROM media md "
|
|
f"{_PROBE.replace('{source}', 'md')}"
|
|
"WHERE md.account_id = {account} "
|
|
"AND search_vector(md.extracted_text) @@ {tsquery}"
|
|
)
|
|
|
|
_VERSION = (
|
|
"SELECT DISTINCT ON (s.date, s.chat_id, s.message_id) " # noqa: S608
|
|
f"{_COLS}, 2 AS rank, NULL::text, v.text, v.observed_at "
|
|
"FROM message_versions v "
|
|
f"{_PROBE.replace('{source}', 'v')}"
|
|
"WHERE v.account_id = {account} "
|
|
"AND search_vector(v.text) @@ {tsquery}"
|
|
)
|
|
|
|
_WITH_MEDIA = (
|
|
"SELECT h.*, md.kind AS media_kind FROM ({hits}) h "
|
|
"LEFT JOIN media md ON md.account_id = {account} "
|
|
"AND md.chat_id = h.chat_id AND md.message_id = h.message_id "
|
|
"ORDER BY h.date DESC, h.chat_id DESC, h.message_id DESC"
|
|
)
|
|
|
|
_SOURCES = ("text", "stt", "version")
|
|
|
|
|
|
class _Binder:
|
|
def __init__(self) -> None:
|
|
self.params: list[object] = []
|
|
|
|
def __call__(self, value: object) -> str:
|
|
self.params.append(value)
|
|
return f"${len(self.params)}"
|
|
|
|
|
|
def _kind(kind: ChatKind, channels: list[int], bind: _Binder) -> str:
|
|
if kind == "private":
|
|
return "s.chat_id > 0"
|
|
member = f"s.chat_id = ANY({bind(channels)}::bigint[])"
|
|
return member if kind == "channel" else f"s.chat_id < 0 AND NOT {member}"
|
|
|
|
|
|
def _conditions(
|
|
filters: SearchFilters, channels: list[int], bind: _Binder
|
|
) -> list[str]:
|
|
conds = [f"s.account_id = {bind(filters.account_id)}"]
|
|
compared = (
|
|
("s.chat_id =", filters.chat_id),
|
|
("s.sender_id =", filters.sender_id),
|
|
("s.has_media =", filters.has_media),
|
|
("s.date >=", filters.date_from),
|
|
("s.date <=", filters.date_to),
|
|
("s.text ~", filters.regex),
|
|
)
|
|
conds += [f"{left} {bind(value)}" for left, value in compared if value is not None]
|
|
flags = (("s.deleted_at", filters.deleted), ("s.edited_at", filters.edited))
|
|
conds += [
|
|
f"{column} IS {'NOT ' if flag else ''}NULL"
|
|
for column, flag in flags
|
|
if flag is not None
|
|
]
|
|
if filters.kind is not None:
|
|
conds.append(_kind(filters.kind, channels, bind))
|
|
if filters.cursor is not None:
|
|
date, chat_id, message_id = decode_cursor(filters.cursor)
|
|
conds.append(
|
|
"(s.date, s.chat_id, s.message_id) < "
|
|
f"({bind(date)}, {bind(chat_id)}, {bind(message_id)})"
|
|
)
|
|
return conds
|
|
|
|
|
|
def _branch(sql: str, limit: str, order: str = _KEY_ORDER, **parts: str) -> str:
|
|
return f"({sql.format(**parts)} ORDER BY {order} LIMIT {limit})"
|
|
|
|
|
|
async def search_messages(pool: asyncpg.Pool, filters: SearchFilters) -> SearchPage:
|
|
async with pool.acquire() as conn, conn.transaction():
|
|
await conn.execute("SET LOCAL jit = off")
|
|
return await _search(conn, filters)
|
|
|
|
|
|
async def _search(conn: asyncpg.Connection, filters: SearchFilters) -> SearchPage:
|
|
query = (filters.query or "").strip()
|
|
terms: list[str] = []
|
|
tsquery = None
|
|
if query:
|
|
parsed = await conn.fetchrow(_PARSE, query)
|
|
terms = list(parsed["terms"])
|
|
if not parsed["tsquery"]:
|
|
return SearchPage(hits=[], next_cursor=None, terms=terms)
|
|
tsquery = parsed["tsquery"]
|
|
|
|
channels: list[int] = []
|
|
if filters.kind in {"group", "channel"}:
|
|
channels = await conn.fetchval(_CHANNELS, filters.account_id)
|
|
bind = _Binder()
|
|
where = " AND ".join(_conditions(filters, channels, bind))
|
|
size = min(filters.limit, MAX_LIMIT)
|
|
limit = bind(size + 1)
|
|
|
|
account = bind(filters.account_id)
|
|
if tsquery is None:
|
|
sql = _branch(_TEXT, limit, where=where)
|
|
else:
|
|
match = f"{bind(tsquery)}::tsquery"
|
|
branches = (
|
|
_branch(_TEXT, limit, where=f"{where} AND s.tsv @@ {match}"),
|
|
_branch(_STT, limit, where=where, account=account, tsquery=match),
|
|
_branch(
|
|
_VERSION,
|
|
limit,
|
|
f"{_KEY_ORDER}, v.observed_at DESC",
|
|
where=where,
|
|
account=account,
|
|
tsquery=match,
|
|
),
|
|
)
|
|
sql = (
|
|
"SELECT DISTINCT ON (date, chat_id, message_id) * FROM (" # noqa: S608
|
|
f"{' UNION ALL '.join(branches)}) s "
|
|
f"ORDER BY {_KEY_ORDER}, rank LIMIT {limit}"
|
|
)
|
|
|
|
rows = await conn.fetch(_WITH_MEDIA.format(hits=sql, account=account), *bind.params)
|
|
hits = [
|
|
SearchHit(**{**dict(row), "source": _SOURCES[row["rank"]]})
|
|
for row in rows[:size]
|
|
]
|
|
next_cursor = None
|
|
if len(rows) > size:
|
|
last = hits[-1]
|
|
next_cursor = encode_cursor(last.date, last.chat_id, last.message_id)
|
|
return SearchPage(hits=hits, next_cursor=next_cursor, terms=terms)
|