From e5a5edcaae53beb27717e7186ad5167988455b15 Mon Sep 17 00:00:00 2001 From: h Date: Mon, 19 Oct 2026 03:38:41 +0000 Subject: [PATCH] perf(backend): search stays in tens of milliseconds at archive size for rare words, channels, deleted and transcripts --- .../versions/d8f1a6c3e5b7_message_search.py | 14 +++- backend/src/utils/search/repository.py | 67 +++++++++++-------- 2 files changed, 50 insertions(+), 31 deletions(-) diff --git a/backend/migrations/versions/d8f1a6c3e5b7_message_search.py b/backend/migrations/versions/d8f1a6c3e5b7_message_search.py index cc0d6a7..34c6a12 100644 --- a/backend/migrations/versions/d8f1a6c3e5b7_message_search.py +++ b/backend/migrations/versions/d8f1a6c3e5b7_message_search.py @@ -1,7 +1,7 @@ """message search index outside the hypertable, versions searchable Revision ID: d8f1a6c3e5b7 -Revises: c9e2a7f4b1d8 +Revises: c5e7a9b1d3f6 Create Date: 2026-10-19 03:00:00.000000 """ @@ -11,7 +11,7 @@ from collections.abc import Sequence from alembic import op revision: str = "d8f1a6c3e5b7" -down_revision: str | None = "c9e2a7f4b1d8" +down_revision: str | None = "c5e7a9b1d3f6" branch_labels: str | Sequence[str] | None = None depends_on: str | Sequence[str] | None = None @@ -95,6 +95,13 @@ def upgrade() -> None: "CREATE INDEX ix_message_search_sender ON message_search " "(account_id, sender_id, date DESC)" ) + for flag in ("deleted", "edited"): + op.execute( + f"CREATE INDEX ix_message_search_{flag} ON message_search " + "(account_id, date DESC, chat_id DESC, message_id DESC) " + f"WHERE {flag}_at IS NOT NULL" + ) + op.execute("ALTER TABLE message_search ALTER COLUMN tsv SET STATISTICS 1000") op.execute(_SYNC) op.execute( "CREATE TRIGGER messages_search_insert AFTER INSERT ON messages " @@ -117,6 +124,9 @@ def upgrade() -> None: "CREATE INDEX ix_media_tsv ON media USING gin (search_vector(extracted_text))" ) op.execute("DROP INDEX IF EXISTS ix_messages_tsv") + for index in ("ix_message_versions_tsv", "ix_media_tsv"): + op.execute(f"ALTER INDEX {index} ALTER COLUMN 1 SET STATISTICS 1000") + op.execute("ANALYZE message_search, message_versions, media") def downgrade() -> None: diff --git a/backend/src/utils/search/repository.py b/backend/src/utils/search/repository.py index 5423b32..e5bb34f 100644 --- a/backend/src/utils/search/repository.py +++ b/backend/src/utils/search/repository.py @@ -2,6 +2,7 @@ import asyncpg from utils.search.models import ( MAX_LIMIT, + ChatKind, SearchFilters, SearchHit, SearchPage, @@ -20,14 +21,14 @@ FROM translate($1, 'Ёё', 'Ее') AS q """ _CHANNELS = """ -(SELECT chat_id FROM ( +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 = {account} + 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') +) typed WHERE type = 'ChatType.CHANNEL' """ _KEY_ORDER = "s.date DESC, s.chat_id DESC, s.message_id DESC" @@ -43,25 +44,27 @@ _TEXT = ( "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 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 " + "NULL::timestamptz FROM media md " + f"{_PROBE.replace('{source}', 'md')}" "WHERE md.account_id = {account} " - "AND search_vector(md.extracted_text) @@ {tsquery} " - "AND {where}" + "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 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 " + 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} " - "AND {where}" + "AND search_vector(v.text) @@ {tsquery}" ) _WITH_MEDIA = ( @@ -83,18 +86,16 @@ class _Binder: return f"${len(self.params)}" -def _kind(filters: SearchFilters, bind: _Binder) -> str | None: - if filters.kind is None: - return None - if filters.kind == "private": +def _kind(kind: ChatKind, channels: list[int], bind: _Binder) -> str: + if 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}" + 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, bind: _Binder) -> list[str]: +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), @@ -111,9 +112,8 @@ def _conditions(filters: SearchFilters, bind: _Binder) -> list[str]: for column, flag in flags if flag is not None ] - kind = _kind(filters, bind) - if kind is not None: - conds.append(kind) + 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( @@ -128,18 +128,27 @@ def _branch(sql: str, limit: str, order: str = _KEY_ORDER, **parts: str) -> str: 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 pool.fetchrow(_PARSE, 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, bind)) + where = " AND ".join(_conditions(filters, channels, bind)) size = min(filters.limit, MAX_LIMIT) limit = bind(size + 1) @@ -166,7 +175,7 @@ async def search_messages(pool: asyncpg.Pool, filters: SearchFilters) -> SearchP f"ORDER BY {_KEY_ORDER}, rank LIMIT {limit}" ) - rows = await pool.fetch(_WITH_MEDIA.format(hits=sql, account=account), *bind.params) + 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]