feat(backend): search finds word prefixes, earlier versions and deleted messages, with filters and cursor paging
This commit is contained in:
hh committed
1 parent
42b4e1d7e6
commit
d45c96c259
5 files changed
+377
-88
No files matched your search
@@ -1,23 +1,48 @@
|
||||
from datetime import datetime
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from typing import Literal
|
||||
|
||||
from pydantic import BaseModel
|
||||
from pydantic import BaseModel, field_validator
|
||||
|
||||
DEFAULT_LIMIT = 50
|
||||
MAX_LIMIT = 500
|
||||
|
||||
_EPOCH = datetime(1970, 1, 1, tzinfo=UTC)
|
||||
_MICROSECOND = timedelta(microseconds=1)
|
||||
|
||||
ChatKind = Literal["private", "group", "channel"]
|
||||
HitSource = Literal["text", "stt", "version"]
|
||||
|
||||
|
||||
def encode_cursor(date: datetime, chat_id: int, message_id: int) -> str:
|
||||
return f"{(date - _EPOCH) // _MICROSECOND}_{chat_id}_{message_id}"
|
||||
|
||||
|
||||
def decode_cursor(cursor: str) -> tuple[datetime, int, int]:
|
||||
micros, chat_id, message_id = (int(part) for part in cursor.split("_"))
|
||||
return _EPOCH + micros * _MICROSECOND, chat_id, message_id
|
||||
|
||||
|
||||
class SearchFilters(BaseModel):
|
||||
account_id: int
|
||||
query: str | None = None
|
||||
chat_id: int | None = None
|
||||
sender_id: int | None = None
|
||||
kind: ChatKind | None = None
|
||||
has_media: bool | None = None
|
||||
deleted: bool | None = None
|
||||
edited: bool | None = None
|
||||
date_from: datetime | None = None
|
||||
date_to: datetime | None = None
|
||||
regex: str | None = None
|
||||
cursor: str | None = None
|
||||
limit: int = DEFAULT_LIMIT
|
||||
offset: int = 0
|
||||
|
||||
@field_validator("cursor")
|
||||
@classmethod
|
||||
def _valid_cursor(cls, value: str | None) -> str | None:
|
||||
if value is not None:
|
||||
decode_cursor(value)
|
||||
return value
|
||||
|
||||
|
||||
class SearchHit(BaseModel):
|
||||
@@ -26,5 +51,17 @@ class SearchHit(BaseModel):
|
||||
date: datetime
|
||||
sender_id: int | None
|
||||
text: str | None
|
||||
source: Literal["text", "stt"]
|
||||
has_media: bool
|
||||
media_kind: str | None
|
||||
edited_at: datetime | None
|
||||
deleted_at: datetime | None
|
||||
source: HitSource
|
||||
extracted_text: str | None
|
||||
version_text: str | None
|
||||
version_at: datetime | None
|
||||
|
||||
|
||||
class SearchPage(BaseModel):
|
||||
hits: list[SearchHit]
|
||||
next_cursor: str | None
|
||||
terms: list[str]
|
||||
@@ -1,88 +1,178 @@
|
||||
import asyncpg
|
||||
|
||||
from utils.search.models import MAX_LIMIT, SearchFilters, SearchHit
|
||||
|
||||
_TEXT_TSV = "to_tsvector('russian', coalesce(m.text, ''))"
|
||||
_STT_TSV = "to_tsvector('russian', coalesce(md.extracted_text, ''))"
|
||||
_TSQUERY = "websearch_to_tsquery('russian', ${})"
|
||||
|
||||
_HITS = "{cols} FROM messages m WHERE {where}" # text branch
|
||||
_STT_HITS = (
|
||||
"{cols} FROM messages m "
|
||||
"JOIN media md ON md.account_id = m.account_id "
|
||||
"AND md.chat_id = m.chat_id AND md.message_id = m.message_id "
|
||||
"WHERE {where}"
|
||||
)
|
||||
_TEXT_COLS = (
|
||||
"SELECT m.chat_id, m.message_id, m.date, m.sender_id, m.text, "
|
||||
"'text' AS source, NULL::text AS extracted_text"
|
||||
)
|
||||
_STT_COLS = (
|
||||
"SELECT m.chat_id, m.message_id, m.date, m.sender_id, m.text, "
|
||||
"'stt' AS source, md.extracted_text"
|
||||
from utils.search.models import (
|
||||
MAX_LIMIT,
|
||||
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
|
||||
"""
|
||||
|
||||
def _common_conditions(filters: SearchFilters, params: list[object]) -> list[str]:
|
||||
params.append(filters.account_id)
|
||||
conds = [f"m.account_id = ${len(params)}"]
|
||||
if filters.chat_id is not None:
|
||||
params.append(filters.chat_id)
|
||||
conds.append(f"m.chat_id = ${len(params)}")
|
||||
if filters.sender_id is not None:
|
||||
params.append(filters.sender_id)
|
||||
conds.append(f"m.sender_id = ${len(params)}")
|
||||
if filters.has_media is not None:
|
||||
params.append(filters.has_media)
|
||||
conds.append(f"m.has_media = ${len(params)}")
|
||||
if filters.date_from is not None:
|
||||
params.append(filters.date_from)
|
||||
conds.append(f"m.date >= ${len(params)}")
|
||||
if filters.date_to is not None:
|
||||
params.append(filters.date_to)
|
||||
conds.append(f"m.date <= ${len(params)}")
|
||||
_CHANNELS = """
|
||||
(SELECT chat_id FROM (
|
||||
SELECT DISTINCT ON (chat_id) chat_id,
|
||||
coalesce(raw->'chat'->>'type', raw->>'type') AS type
|
||||
FROM chat_history
|
||||
WHERE account_id = {account}
|
||||
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}"
|
||||
)
|
||||
|
||||
_STT = (
|
||||
f"SELECT {_COLS}, 1 AS rank, md.extracted_text, NULL::text, " # noqa: S608
|
||||
"NULL::timestamptz "
|
||||
"FROM media md JOIN message_search s ON s.account_id = md.account_id "
|
||||
"AND s.chat_id = md.chat_id AND s.message_id = md.message_id "
|
||||
"WHERE md.account_id = {account} "
|
||||
"AND search_vector(md.extracted_text) @@ {tsquery} "
|
||||
"AND {where}"
|
||||
)
|
||||
|
||||
_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 JOIN message_search s ON s.account_id = v.account_id "
|
||||
"AND s.chat_id = v.chat_id AND s.message_id = v.message_id "
|
||||
"WHERE v.account_id = {account} "
|
||||
"AND search_vector(v.text) @@ {tsquery} "
|
||||
"AND {where}"
|
||||
)
|
||||
|
||||
_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(filters: SearchFilters, bind: _Binder) -> str | None:
|
||||
if filters.kind is None:
|
||||
return None
|
||||
if filters.kind == "private":
|
||||
return "s.chat_id > 0"
|
||||
channels = _CHANNELS.format(account=bind(filters.account_id))
|
||||
if filters.kind == "channel":
|
||||
return f"s.chat_id IN {channels}"
|
||||
return f"s.chat_id < 0 AND s.chat_id NOT IN {channels}"
|
||||
|
||||
|
||||
def _conditions(filters: SearchFilters, 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
|
||||
]
|
||||
kind = _kind(filters, bind)
|
||||
if kind is not None:
|
||||
conds.append(kind)
|
||||
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
|
||||
|
||||
|
||||
async def search_messages(
|
||||
pool: asyncpg.Pool, filters: SearchFilters
|
||||
) -> list[SearchHit]:
|
||||
query = filters.query.strip() if filters.query else None
|
||||
params: list[object] = []
|
||||
common = _common_conditions(filters, params)
|
||||
if filters.regex is not None:
|
||||
params.append(filters.regex)
|
||||
common.append(f"m.text ~ ${len(params)}")
|
||||
def _branch(sql: str, limit: str, order: str = _KEY_ORDER, **parts: str) -> str:
|
||||
return f"({sql.format(**parts)} ORDER BY {order} LIMIT {limit})"
|
||||
|
||||
q_idx: int | None = None
|
||||
if query is not None:
|
||||
params.append(query)
|
||||
q_idx = len(params)
|
||||
|
||||
text_conds = [*common]
|
||||
if q_idx is not None:
|
||||
text_conds.append(f"{_TEXT_TSV} @@ {_TSQUERY.format(q_idx)}")
|
||||
async def search_messages(pool: asyncpg.Pool, filters: SearchFilters) -> SearchPage:
|
||||
query = (filters.query or "").strip()
|
||||
terms: list[str] = []
|
||||
tsquery = None
|
||||
if query:
|
||||
parsed = await pool.fetchrow(_PARSE, query)
|
||||
terms = list(parsed["terms"])
|
||||
if not parsed["tsquery"]:
|
||||
return SearchPage(hits=[], next_cursor=None, terms=terms)
|
||||
tsquery = parsed["tsquery"]
|
||||
|
||||
branches = [_HITS.format(cols=_TEXT_COLS, where=" AND ".join(text_conds))]
|
||||
if q_idx is not None:
|
||||
stt_conds = [*common, f"{_STT_TSV} @@ {_TSQUERY.format(q_idx)}"]
|
||||
branches.append(_STT_HITS.format(cols=_STT_COLS, where=" AND ".join(stt_conds)))
|
||||
bind = _Binder()
|
||||
where = " AND ".join(_conditions(filters, bind))
|
||||
size = min(filters.limit, MAX_LIMIT)
|
||||
limit = bind(size + 1)
|
||||
|
||||
params.append(min(filters.limit, MAX_LIMIT))
|
||||
limit_idx = len(params)
|
||||
params.append(filters.offset)
|
||||
offset_idx = len(params)
|
||||
page = f"ORDER BY date DESC LIMIT ${limit_idx} OFFSET ${offset_idx}"
|
||||
|
||||
if len(branches) == 1:
|
||||
sql = f"SELECT * FROM ({branches[0]}) hits {page}" # noqa: S608
|
||||
account = bind(filters.account_id)
|
||||
if tsquery is None:
|
||||
sql = _branch(_TEXT, limit, where=where)
|
||||
else:
|
||||
union = " UNION ALL ".join(f"({b})" for b in branches)
|
||||
dedup = (
|
||||
"SELECT DISTINCT ON (chat_id, message_id) * " # noqa: S608
|
||||
f"FROM ({union}) hits ORDER BY chat_id, message_id, source DESC"
|
||||
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}"
|
||||
)
|
||||
sql = f"SELECT * FROM ({dedup}) d {page}" # noqa: S608
|
||||
|
||||
rows = await pool.fetch(sql, *params)
|
||||
return [SearchHit(**dict(row)) for row in rows]
|
||||
rows = await pool.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)
|
||||
Reference in new issue
Block a user