perf(backend): search stays in tens of milliseconds at archive size for rare words, channels, deleted and transcripts

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:

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]