yunke/yunke-server/app/db/fts.py

"""FTS5 message-search index — Chinese substring search at index speed.

The plain ``LIKE '%q%'`` scan over 675k messages costs ~4s (measured
2026-08-21). SQLite's FTS5 with the **trigram** tokenizer matches Chinese
substrings (no word segmentation needed) using an index — same results,
milliseconds instead of seconds.

Design:

* ``message_fts`` is an **external-content** FTS5 table over ``message``
  (``content='message', content_rowid='rowid'``) so the text isn't stored
  twice; only the trigram index adds size.
* Three triggers keep the index in sync with ORM upserts transparently —
  no scraper changes needed.
* MySQL (prod target) has no FTS5: ``fts_available`` returns False there
  and the search router falls back to the LIKE scan (``FULLTEXT`` index is
  the MySQL-side upgrade path, deferred).
* Trigram MATCH needs >= 3 characters; shorter queries fall back to LIKE
  at the call site (see :mod:`app.api.routers.search`).
"""

from __future__ import annotations

import logging

import sqlalchemy as sa
from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine

logger = logging.getLogger(__name__)

# DDL — idempotent (IF NOT EXISTS). Triggers mirror the classic FTS5
# external-content pattern from the SQLite docs.
_FTS_DDL = [
    # trigram 分词器:中文子串可 MATCH(≥3字),无需分词库
    """
    CREATE VIRTUAL TABLE IF NOT EXISTS message_fts USING fts5(
        text,
        content='message',
        content_rowid='rowid',
        tokenize='trigram'
    )
    """,
    """
    CREATE TRIGGER IF NOT EXISTS message_fts_ai AFTER INSERT ON message BEGIN
        INSERT INTO message_fts(rowid, text) VALUES (new.rowid, new.text);
    END
    """,
    """
    CREATE TRIGGER IF NOT EXISTS message_fts_au AFTER UPDATE ON message BEGIN
        INSERT INTO message_fts(message_fts, rowid, text)
            VALUES ('delete', old.rowid, old.text);
        INSERT INTO message_fts(rowid, text) VALUES (new.rowid, new.text);
    END
    """,
    """
    CREATE TRIGGER IF NOT EXISTS message_fts_ad AFTER DELETE ON message BEGIN
        INSERT INTO message_fts(message_fts, rowid, text)
            VALUES ('delete', old.rowid, old.text);
    END
    """,
]


async def ensure_fts(holder) -> None:
    """Create the FTS5 virtual table + sync triggers (idempotent, no-op on MySQL).

    Accepts an ``AsyncSession`` / ``AsyncConnection`` (anything with async
    ``execute``) or an ``AsyncEngine``. Raises whatever the dialect raises
    for unsupported DDL — callers on non-SQLite engines should catch and
    treat FTS as unavailable.
    """
    if hasattr(holder, "execute"):
        for ddl in _FTS_DDL:
            await holder.execute(sa.text(ddl))
    else:  # AsyncEngine
        async with holder.begin() as conn:
            for ddl in _FTS_DDL:
                await conn.execute(sa.text(ddl))


async def rebuild_fts(holder) -> int:
    """(Re)populate the index from the current ``message`` rows.

    Returns the rowcount. One-time backfill for existing data — triggers
    only cover writes made after ``ensure_fts``.
    """
    result = await holder.execute(
        sa.text("INSERT INTO message_fts(message_fts) VALUES('rebuild')")
    )
    return result.rowcount or 0


async def fts_available(holder) -> bool:
    """True iff the ``message_fts`` virtual table exists (session/conn/engine)."""
    try:
        if hasattr(holder, "execute"):
            row = (
                await holder.execute(
                    sa.text(
                        "SELECT count(*) FROM sqlite_master "
                        "WHERE type='table' AND name='message_fts'"
                    )
                )
            ).scalar_one()
            return bool(row)
        async with holder.connect() as conn:
            row = (
                await conn.execute(
                    sa.text(
                        "SELECT count(*) FROM sqlite_master "
                        "WHERE type='table' AND name='message_fts'"
                    )
                )
            ).scalar_one()
            return bool(row)
    except Exception:
        # MySQL 等 — sqlite_master 不存在
        return False


async def fts_match_rows(
    engine: AsyncEngine,
    query: str,
    *,
    limit: int = 200,
) -> list:
    """MATCH ``query`` against the trigram index, newest first.

    Returns Message ORM rows (joined back from rowid via raw SQL — FTS
    virtual tables can't join into ORM ``select()`` directly). Caller is
    responsible for ensuring the index exists (``fts_available``) and
    ``len(query) >= 3``.
    """
    from sqlalchemy.ext.asyncio import async_sessionmaker

    from app.db.models import Message

    maker = async_sessionmaker(engine, expire_on_commit=False)
    async with maker() as session:
        # 先取 MATCH 命中的 rowid(索引快),再按主键取 ORM 行
        rowids = (
            await session.execute(
                sa.text(
                    "SELECT rowid FROM message_fts WHERE message_fts MATCH :q "
                    "ORDER BY rowid LIMIT :lim"
                ),
                {"q": f'"{query}"', "lim": limit * 2},
            )
        ).scalars().all()
        if not rowids:
            return []
        # rowid → msg_id 再取行(消息表 rowid 与 ORM 无直接映射)
        rows = (
            await session.execute(
                sa.text(
                    "SELECT msg_id FROM message WHERE rowid IN :rids"
                ).bindparams(sa.bindparam("rids", expanding=True)),
                {"rids": list(rowids)},
            )
        ).scalars().all()
        if not rows:
            return []
        stmt = (
            sa.select(Message)
            .where(Message.msg_id.in_(rows))
            .order_by(sa.nullslast(Message.timestamp.desc()))
            .limit(limit)
        )
        result = await session.execute(stmt)
        return result.scalars().all()