From f31c39c40f101a36339c77bffed1df729ed98334 Mon Sep 17 00:00:00 2001 From: Symb0x76 <9667434@qq.com> Date: Fri, 10 Oct 2025 13:59:50 +0800 Subject: [PATCH] perf(search): Delete FTS fallback, use Bleve only. --- internal/chatlog/http/static/index.htm | 25 +- .../datasource/darwinv3/datasource.go | 4 - internal/wechatdb/datasource/datasource.go | 1 - internal/wechatdb/datasource/v4/datasource.go | 425 +++++++----------- .../datasource/windowsv3/datasource.go | 413 +++++------------ internal/wechatdb/repository/index.go | 315 +++++++++++++ internal/wechatdb/repository/repository.go | 36 +- internal/wechatdb/repository/search.go | 14 +- internal/wechatdb/wechatdb.go | 4 +- 9 files changed, 668 insertions(+), 569 deletions(-) create mode 100644 internal/wechatdb/repository/index.go diff --git a/internal/chatlog/http/static/index.htm b/internal/chatlog/http/static/index.htm index bf406a4..f2f6df1 100644 --- a/internal/chatlog/http/static/index.htm +++ b/internal/chatlog/http/static/index.htm @@ -640,8 +640,9 @@

- 使用全文索引在指定会话中查找消息。Bleve + 在指定会话中查找消息。GET /api/v1/search

@@ -649,18 +650,18 @@
-
- 可以从“最近会话”或“群聊”接口获取对应的 - ID；如留空将检索所有会话。 -
@@ -699,12 +696,12 @@
- 可选择开始日期和结束日期；留空表示不限制时间范围。 + 留空表示不限制时间范围。
@@ -718,7 +715,7 @@
diff --git a/internal/wechatdb/datasource/darwinv3/datasource.go b/internal/wechatdb/datasource/darwinv3/datasource.go index 9470c11..c476c95 100644 --- a/internal/wechatdb/datasource/darwinv3/datasource.go +++ b/internal/wechatdb/datasource/darwinv3/datasource.go @@ -511,10 +511,6 @@ func (ds *DataSource) GetMessages(ctx context.Context, startTime, endTime time.T return filteredMessages, nil } -func (ds *DataSource) SearchMessages(ctx context.Context, req *model.SearchRequest) (*model.SearchResponse, error) { - return nil, errors.SearchNotSupported("darwin", 3) -} - // 从表名中提取 talker func extractTalkerFromTableName(tableName string) string { diff --git a/internal/wechatdb/datasource/datasource.go b/internal/wechatdb/datasource/datasource.go index 6211d9b..ac54c25 100644 --- a/internal/wechatdb/datasource/datasource.go +++ b/internal/wechatdb/datasource/datasource.go @@ -17,7 +17,6 @@ type DataSource interface { // 消息 GetMessages(ctx context.Context, startTime, endTime time.Time, talker string, sender string, keyword string, limit, offset int) ([]*model.Message, error) - SearchMessages(ctx context.Context, req *model.SearchRequest) (*model.SearchResponse, error) GetDatasetFingerprint(ctx context.Context) (string, error) // 联系人 diff --git a/internal/wechatdb/datasource/v4/datasource.go b/internal/wechatdb/datasource/v4/datasource.go index 355c985..6a080c0 100644 --- a/internal/wechatdb/datasource/v4/datasource.go +++ b/internal/wechatdb/datasource/v4/datasource.go @@ -6,7 +6,6 @@ import ( "database/sql" "encoding/hex" "fmt" - "math" "regexp" "sort" "strings" @@ -369,279 +368,193 @@ func (ds *DataSource) GetMessages(ctx context.Context, startTime, endTime time.T return filteredMessages, nil } -func (ds *DataSource) SearchMessages(ctx context.Context, req *model.SearchRequest) (*model.SearchResponse, error) { - if req == nil { - return nil, errors.InvalidArg("request") +func (ds *DataSource) ListTalkers(ctx context.Context) ([]string, error) { + talkerSet := make(map[string]struct{}) + add := func(name string) { + name = strings.TrimSpace(name) + if name == "" { + return + } + if _, ok := talkerSet[name]; ok { + return + } + talkerSet[name] = struct{}{} } - ftsQuery := util.BuildFTSQuery(req.Query) - if ftsQuery == "" { - return &model.SearchResponse{ - Total: 0, - Hits: []*model.SearchHit{}, - DurationMs: 0, - Limit: req.Limit, - Offset: req.Offset, - Query: req.Query, - Talker: req.Talker, - Sender: req.Sender, - Start: req.Start, - End: req.End, - }, nil - } - - talkers := util.Str2List(req.Talker, ",") - if len(talkers) == 0 { - return nil, errors.ErrTalkerEmpty - } - - limit := req.Limit - if limit <= 0 { - limit = 20 - } - if limit > 200 { - limit = 200 - } - offset := req.Offset - if offset < 0 { - offset = 0 - } - req.Limit = limit - req.Offset = offset - - budget := limit + offset + 20 - if budget < 50 { - budget = 50 - } - if budget > 1000 { - budget = 1000 - } - - startRange := req.Start - endRange := req.End - if startRange.IsZero() { - startRange = time.Unix(0, 0) - } - if endRange.IsZero() { - endRange = time.Now().Add(24 * time.Hour) - } - if endRange.Before(startRange) { - startRange, endRange = endRange, startRange - } - - startUnix := int64(0) - if !req.Start.IsZero() { - startUnix = req.Start.Unix() - } - endUnix := int64(math.MaxInt64) - if !req.End.IsZero() { - endUnix = req.End.Unix() - } - if endUnix < startUnix { - startUnix, endUnix = endUnix, startUnix - } - - senders := util.Str2List(req.Sender, ",") - senderFilter := len(senders) > 0 - senderSet := make(map[string]struct{}, len(senders)) - for _, s := range senders { - senderSet[s] = struct{}{} - } - - startTs := time.Now() - allHits := make([]*model.SearchHit, 0, limit+offset) - seenSeq := make(map[int64]struct{}) - - dbInfos := ds.getDBInfosForTimeRange(startRange, endRange) - if len(dbInfos) == 0 { - return &model.SearchResponse{ - Total: 0, - Hits: []*model.SearchHit{}, - DurationMs: time.Since(startTs).Milliseconds(), - Limit: limit, - Offset: offset, - Query: req.Query, - Talker: req.Talker, - Sender: req.Sender, - Start: req.Start, - End: req.End, - }, nil - } - - for _, dbInfo := range dbInfos { - if err := ctx.Err(); err != nil { - return nil, err + db, err := ds.dbm.GetDB(Contact) + if err != nil { + log.Debug().Err(err).Msg("query contact usernames failed") + } else if db != nil { + rows, err := db.QueryContext(ctx, `SELECT username FROM contact WHERE IFNULL(username,'') <> ''`) + if err == nil { + for rows.Next() { + if err := ctx.Err(); err != nil { + rows.Close() + return nil, err + } + var username string + if scanErr := rows.Scan(&username); scanErr == nil { + add(username) + } else { + log.Debug().Err(scanErr).Msg("scan contact username failed") + } + } + if err := rows.Err(); err != nil { + log.Debug().Err(err).Msg("iterate contact usernames failed") + } + rows.Close() + } else { + log.Debug().Err(err).Msg("query contact usernames failed") } - db, err := ds.dbm.OpenDB(dbInfo.FilePath) + roomRows, err := db.QueryContext(ctx, `SELECT username FROM chat_room WHERE IFNULL(username,'') <> ''`) + if err == nil { + for roomRows.Next() { + if err := ctx.Err(); err != nil { + roomRows.Close() + return nil, err + } + var username string + if scanErr := roomRows.Scan(&username); scanErr == nil { + add(username) + } else { + log.Debug().Err(scanErr).Msg("scan chat_room username failed") + } + } + if err := roomRows.Err(); err != nil { + log.Debug().Err(err).Msg("iterate chat_room usernames failed") + } + roomRows.Close() + } else { + log.Debug().Err(err).Msg("query chat_room usernames failed") + } + } + + db, err = ds.dbm.GetDB(Session) + if err != nil { + log.Debug().Err(err).Msg("query session usernames failed") + } else if db != nil { + rows, err := db.QueryContext(ctx, `SELECT username FROM SessionTable WHERE IFNULL(username,'') <> ''`) + if err == nil { + for rows.Next() { + if err := ctx.Err(); err != nil { + rows.Close() + return nil, err + } + var username string + if scanErr := rows.Scan(&username); scanErr == nil { + add(username) + } else { + log.Debug().Err(scanErr).Msg("scan session username failed") + } + } + if err := rows.Err(); err != nil { + log.Debug().Err(err).Msg("iterate session usernames failed") + } + rows.Close() + } else { + log.Debug().Err(err).Msg("query session usernames failed") + } + } + + talkers := make([]string, 0, len(talkerSet)) + for username := range talkerSet { + talkers = append(talkers, username) + } + sort.Strings(talkers) + return talkers, nil +} + +func (ds *DataSource) IterateMessages(ctx context.Context, talkers []string, handler func(*model.Message) error) error { + if handler == nil { + return errors.InvalidArg("handler") + } + + if len(talkers) == 0 { + var err error + talkers, err = ds.ListTalkers(ctx) + if err != nil { + return err + } + } + if len(talkers) == 0 { + return nil + } + + tableNames := make(map[string]string, len(talkers)) + for _, talker := range talkers { + hash := md5.Sum([]byte(talker)) + tableNames[talker] = "Msg_" + hex.EncodeToString(hash[:]) + } + + for _, info := range ds.messageInfos { + if err := ctx.Err(); err != nil { + return err + } + + db, err := ds.dbm.OpenDB(info.FilePath) if err != nil { - log.Err(err).Msgf("数据库 %s 未打开", dbInfo.FilePath) continue } for _, talker := range talkers { if err := ctx.Err(); err != nil { - return nil, err + return err } + tableName := tableNames[talker] - err := func() error { - tx, err := db.BeginTx(ctx, nil) - if err != nil { - return errors.QueryFailed("begin tx", err) - } - defer func() { - _ = tx.Rollback() - }() - - talkerHash := md5.Sum([]byte(talker)) - tableName := "Msg_" + hex.EncodeToString(talkerHash[:]) - - var exists int - if err := tx.QueryRowContext(ctx, `SELECT 1 FROM sqlite_master WHERE type='table' AND name=?`, tableName).Scan(&exists); err != nil { - if err == sql.ErrNoRows { - return nil - } - return errors.QueryFailed("check table", err) - } - - if _, err := tx.ExecContext(ctx, `CREATE TEMP TABLE IF NOT EXISTS chatlog_msg_base(seq INTEGER PRIMARY KEY, content TEXT)`); err != nil { - return errors.QueryFailed("create temp table", err) - } - if _, err := tx.ExecContext(ctx, `DELETE FROM chatlog_msg_base`); err != nil { - return errors.QueryFailed("cleanup temp table", err) - } - - insertSQL := fmt.Sprintf(`INSERT OR IGNORE INTO chatlog_msg_base(seq, content) - SELECT sort_seq, CAST(message_content AS TEXT) - FROM %s - WHERE create_time >= ? AND create_time <= ? - AND message_content IS NOT NULL AND LENGTH(message_content) > 0`, tableName) - if _, err := tx.ExecContext(ctx, insertSQL, startUnix, endUnix); err != nil { - return errors.QueryFailed("insert temp base", err) - } - - if _, err := tx.ExecContext(ctx, `CREATE VIRTUAL TABLE IF NOT EXISTS temp.chatlog_msg_fts USING fts5(content, tokenize='unicode61')`); err != nil { - return errors.QueryFailed("create fts", err) - } - if _, err := tx.ExecContext(ctx, `DELETE FROM temp.chatlog_msg_fts`); err != nil { - return errors.QueryFailed("cleanup fts", err) - } - if _, err := tx.ExecContext(ctx, `INSERT INTO temp.chatlog_msg_fts(rowid, content) SELECT seq, content FROM chatlog_msg_base`); err != nil { - return errors.QueryFailed("insert fts", err) - } - - querySQL := fmt.Sprintf(` - SELECT m.sort_seq, m.server_id, m.local_type, n.user_name, - m.create_time, m.message_content, m.packed_info_data, m.status, - highlight(chatlog_msg_fts, 0, '', '') AS snippet, - bm25(chatlog_msg_fts) AS score - FROM temp.chatlog_msg_fts - JOIN %s AS m ON m.sort_seq = chatlog_msg_fts.rowid - LEFT JOIN Name2Id n ON m.real_sender_id = n.rowid - WHERE chatlog_msg_fts MATCH ? - ORDER BY score - LIMIT ? - `, tableName) - - rows, err := tx.QueryContext(ctx, querySQL, ftsQuery, budget) - if err != nil { - return errors.QueryFailed("fts query", err) - } - defer rows.Close() - - for rows.Next() { - var msg model.MessageV4 - var messageContent []byte - var snippet sql.NullString - var score sql.NullFloat64 - if err := rows.Scan(&msg.SortSeq, &msg.ServerID, &msg.LocalType, &msg.UserName, - &msg.CreateTime, &messageContent, &msg.PackedInfoData, &msg.Status, &snippet, &score); err != nil { - return errors.ScanRowFailed(err) - } - msg.MessageContent = messageContent - - message := msg.Wrap(talker) - if senderFilter { - if _, ok := senderSet[message.Sender]; !ok { - continue - } - } - if _, dup := seenSeq[message.Seq]; dup { - continue - } - seenSeq[message.Seq] = struct{}{} - - hitSnippet := "" - if snippet.Valid { - hitSnippet = snippet.String - } - hitScore := 0.0 - if score.Valid { - hitScore = score.Float64 - } - allHits = append(allHits, &model.SearchHit{Message: message, Snippet: hitSnippet, Score: hitScore}) - } - - if err := rows.Err(); err != nil { - return errors.QueryFailed("fts rows", err) - } - - if err := tx.Commit(); err != nil { - return errors.QueryFailed("commit tx", err) - } - return nil - }() + query := fmt.Sprintf(` + SELECT m.sort_seq, m.server_id, m.local_type, n.user_name, + m.create_time, m.message_content, m.packed_info_data, m.status + FROM %s AS m + LEFT JOIN Name2Id n ON m.real_sender_id = n.rowid + ORDER BY m.sort_seq ASC + `, tableName) + rows, err := db.QueryContext(ctx, query) if err != nil { - return nil, err + if strings.Contains(err.Error(), "no such table") { + continue + } + return errors.QueryFailed("iterate messages", err) } + + for rows.Next() { + if err := ctx.Err(); err != nil { + rows.Close() + return err + } + var msg model.MessageV4 + var messageContent []byte + if scanErr := rows.Scan( + &msg.SortSeq, + &msg.ServerID, + &msg.LocalType, + &msg.UserName, + &msg.CreateTime, + &messageContent, + &msg.PackedInfoData, + &msg.Status, + ); scanErr != nil { + rows.Close() + return errors.ScanRowFailed(scanErr) + } + msg.MessageContent = messageContent + message := msg.Wrap(talker) + if err := handler(message); err != nil { + rows.Close() + return err + } + } + if err := rows.Err(); err != nil { + rows.Close() + return errors.QueryFailed("iterate message rows", err) + } + rows.Close() } } - if len(allHits) == 0 { - return &model.SearchResponse{ - Total: 0, - Hits: []*model.SearchHit{}, - DurationMs: time.Since(startTs).Milliseconds(), - Limit: limit, - Offset: offset, - Query: req.Query, - Talker: req.Talker, - Sender: req.Sender, - Start: req.Start, - End: req.End, - }, nil - } - - sort.SliceStable(allHits, func(i, j int) bool { - if allHits[i].Score == allHits[j].Score { - return allHits[i].Message.Seq < allHits[j].Message.Seq - } - return allHits[i].Score < allHits[j].Score - }) - - total := len(allHits) - if offset > total { - offset = total - } - endIndex := offset + limit - if endIndex > total { - endIndex = total - } - paged := allHits[offset:endIndex] - - return &model.SearchResponse{ - Total: total, - Hits: paged, - DurationMs: time.Since(startTs).Milliseconds(), - Limit: limit, - Offset: offset, - Query: req.Query, - Talker: req.Talker, - Sender: req.Sender, - Start: req.Start, - End: req.End, - }, nil + return nil } func (ds *DataSource) GetDatasetFingerprint(context.Context) (string, error) { diff --git a/internal/wechatdb/datasource/windowsv3/datasource.go b/internal/wechatdb/datasource/windowsv3/datasource.go index 77e2b89..4bf5033 100644 --- a/internal/wechatdb/datasource/windowsv3/datasource.go +++ b/internal/wechatdb/datasource/windowsv3/datasource.go @@ -5,7 +5,6 @@ import ( "database/sql" "encoding/hex" "fmt" - "math" "regexp" "sort" "strings" @@ -423,294 +422,6 @@ func (ds *DataSource) GetMessages(ctx context.Context, startTime, endTime time.T return filteredMessages, nil } -func (ds *DataSource) SearchMessages(ctx context.Context, req *model.SearchRequest) (*model.SearchResponse, error) { - if req == nil { - return nil, errors.InvalidArg("request") - } - - ftsQuery := util.BuildFTSQuery(req.Query) - if ftsQuery == "" { - return &model.SearchResponse{ - Total: 0, - Hits: []*model.SearchHit{}, - DurationMs: 0, - Limit: req.Limit, - Offset: req.Offset, - Query: req.Query, - Talker: req.Talker, - Sender: req.Sender, - Start: req.Start, - End: req.End, - }, nil - } - - talkers := util.Str2List(req.Talker, ",") - if len(talkers) == 0 { - var err error - talkers, err = ds.collectAllTalkers(ctx) - if err != nil { - return nil, err - } - } - if len(talkers) == 0 { - return &model.SearchResponse{ - Total: 0, - Hits: []*model.SearchHit{}, - DurationMs: 0, - Limit: req.Limit, - Offset: req.Offset, - Query: req.Query, - Talker: req.Talker, - Sender: req.Sender, - Start: req.Start, - End: req.End, - }, nil - } - - limit := req.Limit - if limit <= 0 { - limit = 20 - } - if limit > 200 { - limit = 200 - } - offset := req.Offset - if offset < 0 { - offset = 0 - } - req.Limit = limit - req.Offset = offset - - budget := limit + offset + 20 - if budget < 50 { - budget = 50 - } - if budget > 1000 { - budget = 1000 - } - - startRange := req.Start - endRange := req.End - if startRange.IsZero() { - startRange = time.Unix(0, 0) - } - if endRange.IsZero() { - endRange = time.Now().Add(24 * time.Hour) - } - if endRange.Before(startRange) { - startRange, endRange = endRange, startRange - } - - startSeq := int64(0) - if !req.Start.IsZero() { - startSeq = req.Start.Unix() * 1000 - } - endSeq := int64(math.MaxInt64) - if !req.End.IsZero() { - endSeq = req.End.Unix() * 1000 - } - if endSeq < startSeq { - startSeq, endSeq = endSeq, startSeq - } - - senders := util.Str2List(req.Sender, ",") - senderFilter := len(senders) > 0 - senderSet := make(map[string]struct{}, len(senders)) - for _, s := range senders { - senderSet[s] = struct{}{} - } - - startTs := time.Now() - allHits := make([]*model.SearchHit, 0, limit+offset) - seenSeq := make(map[int64]struct{}) - - dbInfos := ds.getDBInfosForTimeRange(startRange, endRange) - if len(dbInfos) == 0 { - return &model.SearchResponse{ - Total: 0, - Hits: []*model.SearchHit{}, - DurationMs: time.Since(startTs).Milliseconds(), - Limit: limit, - Offset: offset, - Query: req.Query, - Talker: req.Talker, - Sender: req.Sender, - Start: req.Start, - End: req.End, - }, nil - } - - for _, dbInfo := range dbInfos { - if err := ctx.Err(); err != nil { - return nil, err - } - - db, err := ds.dbm.OpenDB(dbInfo.FilePath) - if err != nil { - log.Err(err).Msgf("数据库 %s 未打开", dbInfo.FilePath) - continue - } - - for _, talker := range talkers { - if err := ctx.Err(); err != nil { - return nil, err - } - - err := func() error { - tx, err := db.BeginTx(ctx, nil) - if err != nil { - return errors.QueryFailed("begin tx", err) - } - defer func() { - _ = tx.Rollback() - }() - - if _, err := tx.ExecContext(ctx, `CREATE TEMP TABLE IF NOT EXISTS chatlog_msg_base(seq INTEGER PRIMARY KEY, content TEXT)`); err != nil { - return errors.QueryFailed("create temp table", err) - } - if _, err := tx.ExecContext(ctx, `DELETE FROM chatlog_msg_base`); err != nil { - return errors.QueryFailed("cleanup temp table", err) - } - - conditions := []string{"Sequence >= ?", "Sequence <= ?", "StrContent IS NOT NULL", "LENGTH(StrContent) > 0"} - args := []interface{}{startSeq, endSeq} - if talkerID, ok := dbInfo.TalkerMap[talker]; ok { - conditions = append(conditions, "TalkerId = ?") - args = append(args, talkerID) - } else { - conditions = append(conditions, "StrTalker = ?") - args = append(args, talker) - } - - insertSQL := fmt.Sprintf(`INSERT OR IGNORE INTO chatlog_msg_base(seq, content) - SELECT Sequence, StrContent FROM MSG WHERE %s`, strings.Join(conditions, " AND ")) - if _, err := tx.ExecContext(ctx, insertSQL, args...); err != nil { - return errors.QueryFailed("insert temp base", err) - } - - if _, err := tx.ExecContext(ctx, `CREATE VIRTUAL TABLE IF NOT EXISTS temp.chatlog_msg_fts USING fts5(content, tokenize='unicode61')`); err != nil { - return errors.QueryFailed("create fts", err) - } - if _, err := tx.ExecContext(ctx, `DELETE FROM temp.chatlog_msg_fts`); err != nil { - return errors.QueryFailed("cleanup fts", err) - } - if _, err := tx.ExecContext(ctx, `INSERT INTO temp.chatlog_msg_fts(rowid, content) SELECT seq, content FROM chatlog_msg_base`); err != nil { - return errors.QueryFailed("insert fts", err) - } - - rows, err := tx.QueryContext(ctx, ` - SELECT m.MsgSvrID, m.Sequence, m.CreateTime, m.StrTalker, m.IsSender, - m.Type, m.SubType, m.StrContent, m.CompressContent, m.BytesExtra, - highlight(chatlog_msg_fts, 0, '', '') AS snippet, - bm25(chatlog_msg_fts) AS score - FROM temp.chatlog_msg_fts - JOIN MSG m ON m.Sequence = chatlog_msg_fts.rowid - WHERE chatlog_msg_fts MATCH ? - ORDER BY score - LIMIT ? - `, ftsQuery, budget) - if err != nil { - return errors.QueryFailed("fts query", err) - } - defer rows.Close() - - for rows.Next() { - var msg model.MessageV3 - var compressContent []byte - var bytesExtra []byte - var snippet sql.NullString - var score sql.NullFloat64 - if err := rows.Scan(&msg.MsgSvrID, &msg.Sequence, &msg.CreateTime, &msg.StrTalker, &msg.IsSender, - &msg.Type, &msg.SubType, &msg.StrContent, &compressContent, &bytesExtra, &snippet, &score); err != nil { - return errors.ScanRowFailed(err) - } - msg.CompressContent = compressContent - msg.BytesExtra = bytesExtra - - message := msg.Wrap() - if senderFilter { - if _, ok := senderSet[message.Sender]; !ok { - continue - } - } - if _, dup := seenSeq[message.Seq]; dup { - continue - } - seenSeq[message.Seq] = struct{}{} - - hitSnippet := "" - if snippet.Valid { - hitSnippet = snippet.String - } - hitScore := 0.0 - if score.Valid { - hitScore = score.Float64 - } - allHits = append(allHits, &model.SearchHit{Message: message, Snippet: hitSnippet, Score: hitScore}) - } - - if err := rows.Err(); err != nil { - return errors.QueryFailed("fts rows", err) - } - - if err := tx.Commit(); err != nil { - return errors.QueryFailed("commit tx", err) - } - return nil - }() - - if err != nil { - return nil, err - } - } - } - - if len(allHits) == 0 { - return &model.SearchResponse{ - Total: 0, - Hits: []*model.SearchHit{}, - DurationMs: time.Since(startTs).Milliseconds(), - Limit: limit, - Offset: offset, - Query: req.Query, - Talker: req.Talker, - Sender: req.Sender, - Start: req.Start, - End: req.End, - }, nil - } - - sort.SliceStable(allHits, func(i, j int) bool { - if allHits[i].Score == allHits[j].Score { - return allHits[i].Message.Seq < allHits[j].Message.Seq - } - return allHits[i].Score < allHits[j].Score - }) - - total := len(allHits) - if offset > total { - offset = total - } - endIndex := offset + limit - if endIndex > total { - endIndex = total - } - paged := allHits[offset:endIndex] - - return &model.SearchResponse{ - Total: total, - Hits: paged, - DurationMs: time.Since(startTs).Milliseconds(), - Limit: limit, - Offset: offset, - Query: req.Query, - Talker: req.Talker, - Sender: req.Sender, - Start: req.Start, - End: req.End, - }, nil -} - func (ds *DataSource) GetDatasetFingerprint(context.Context) (string, error) { return ds.dbm.FingerprintForGroups(Message) } @@ -822,6 +533,130 @@ func (ds *DataSource) collectAllTalkers(ctx context.Context) ([]string, error) { return talkers, nil } +// ListTalkers 返回所有对话用户名（联系人 + 群聊），供索引使用 +func (ds *DataSource) ListTalkers(ctx context.Context) ([]string, error) { + return ds.collectAllTalkers(ctx) +} + +// IterateMessages 按 talker 枚举消息并交给处理函数，供 Bleve 索引使用 +func (ds *DataSource) IterateMessages(ctx context.Context, talkers []string, handler func(*model.Message) error) error { + if handler == nil { + return errors.InvalidArg("handler") + } + + if len(talkers) == 0 { + var err error + talkers, err = ds.collectAllTalkers(ctx) + if err != nil { + return err + } + } + if len(talkers) == 0 { + return nil + } + + uniqueTalkers := make([]string, 0, len(talkers)) + seen := make(map[string]struct{}, len(talkers)) + for _, t := range talkers { + t = strings.TrimSpace(t) + if t == "" { + continue + } + if _, ok := seen[t]; ok { + continue + } + seen[t] = struct{}{} + uniqueTalkers = append(uniqueTalkers, t) + } + if len(uniqueTalkers) == 0 { + return nil + } + + for _, info := range ds.messageInfos { + if err := ctx.Err(); err != nil { + return err + } + + db, err := ds.dbm.OpenDB(info.FilePath) + if err != nil { + log.Debug().Err(err).Msgf("open message db failed: %s", info.FilePath) + continue + } + + for _, talker := range uniqueTalkers { + if err := ctx.Err(); err != nil { + return err + } + + conditions := []string{"StrContent IS NOT NULL"} + args := make([]interface{}, 0, 1) + if talkerID, ok := info.TalkerMap[talker]; ok { + conditions = append(conditions, "TalkerId = ?") + args = append(args, talkerID) + } else { + conditions = append(conditions, "StrTalker = ?") + args = append(args, talker) + } + + query := fmt.Sprintf(` + SELECT MsgSvrID, Sequence, CreateTime, StrTalker, IsSender, + Type, SubType, StrContent, CompressContent, BytesExtra + FROM MSG + WHERE %s + ORDER BY Sequence ASC + `, strings.Join(conditions, " AND ")) + + rows, err := db.QueryContext(ctx, query, args...) + if err != nil { + if strings.Contains(err.Error(), "no such table") { + continue + } + return errors.QueryFailed("iterate messages", err) + } + + for rows.Next() { + if err := ctx.Err(); err != nil { + rows.Close() + return err + } + var msg model.MessageV3 + var compressContent []byte + var bytesExtra []byte + if scanErr := rows.Scan( + &msg.MsgSvrID, + &msg.Sequence, + &msg.CreateTime, + &msg.StrTalker, + &msg.IsSender, + &msg.Type, + &msg.SubType, + &msg.StrContent, + &compressContent, + &bytesExtra, + ); scanErr != nil { + rows.Close() + return errors.ScanRowFailed(scanErr) + } + msg.CompressContent = compressContent + msg.BytesExtra = bytesExtra + + wrapped := msg.Wrap() + if err := handler(wrapped); err != nil { + rows.Close() + return err + } + } + if err := rows.Err(); err != nil { + rows.Close() + return errors.QueryFailed("iterate message rows", err) + } + rows.Close() + } + } + + return nil +} + // GetContacts 实现获取联系人信息的方法 func (ds *DataSource) GetContacts(ctx context.Context, key string, limit, offset int) ([]*model.Contact, error) { var query string diff --git a/internal/wechatdb/repository/index.go b/internal/wechatdb/repository/index.go new file mode 100644 index 0000000..63ca07e --- /dev/null +++ b/internal/wechatdb/repository/index.go @@ -0,0 +1,315 @@ +package repository + +import ( + "context" + "errors" + "fmt" + "sort" + "strings" + "time" + + "github.com/rs/zerolog/log" + + "github.com/sjzar/chatlog/internal/model" + "github.com/sjzar/chatlog/internal/wechatdb/indexer" + "github.com/sjzar/chatlog/pkg/util" +) + +type bleveIndexable interface { + ListTalkers(ctx context.Context) ([]string, error) + IterateMessages(ctx context.Context, talkers []string, fn func(*model.Message) error) error +} + +func (r *Repository) initIndex() error { + if r.indexPath == "" { + return nil + } + + idx, err := indexer.Open(r.indexPath) + if err != nil { + return err + } + + r.index = idx + r.indexCtx, r.indexCancel = context.WithCancel(context.Background()) + + go func() { + ready, err := r.ensureIndex(r.indexCtx) + if err != nil && !errors.Is(err, context.Canceled) { + log.Warn().Err(err).Msg("ensure bleve index failed") + return + } + if ready { + log.Info().Msg("bleve index ready") + } + }() + + return nil +} + +func (r *Repository) ensureIndex(ctx context.Context) (bool, error) { + if r.index == nil { + return false, nil + } + + if ctx == nil { + ctx = context.Background() + } + + versionMatched, err := r.index.EnsureVersion() + if err != nil { + return false, err + } + + if !versionMatched { + r.indexMu.Lock() + r.indexStatus.Ready = false + r.indexStatus.Progress = 0 + r.indexFingerprint = "" + r.indexMu.Unlock() + } + + fp, err := r.ds.GetDatasetFingerprint(ctx) + if err != nil { + return false, err + } + if strings.TrimSpace(fp) == "" { + return false, fmt.Errorf("dataset fingerprint is empty") + } + + r.indexMu.Lock() + if r.indexFingerprint == fp && r.indexStatus.Ready && !r.indexStatus.InProgress { + status := r.indexStatus + r.indexMu.Unlock() + if status.Progress < 1 { + r.updateIndexProgress(1) + } + return true, nil + } + if r.indexStatus.InProgress { + r.indexMu.Unlock() + return false, nil + } + r.indexStatus.InProgress = true + r.indexStatus.Ready = false + r.indexStatus.Progress = 0 + r.indexStatus.LastStartedAt = time.Now() + r.indexStatus.LastError = "" + r.indexMu.Unlock() + + err = r.rebuildIndex(ctx, fp) + if err != nil { + if err == context.Canceled || errors.Is(err, context.Canceled) { + r.indexMu.Lock() + r.indexStatus.InProgress = false + r.indexMu.Unlock() + return false, err + } + + r.indexMu.Lock() + r.indexStatus.InProgress = false + r.indexStatus.LastError = err.Error() + r.indexMu.Unlock() + return false, err + } + + r.indexMu.Lock() + r.indexFingerprint = fp + r.indexStatus.InProgress = false + r.indexStatus.Ready = true + r.indexStatus.Progress = 1 + r.indexStatus.LastCompletedAt = time.Now() + r.indexMu.Unlock() + + return true, nil +} + +func (r *Repository) rebuildIndex(ctx context.Context, fp string) error { + indexable, ok := r.ds.(bleveIndexable) + if !ok { + return fmt.Errorf("datasource does not support bleve indexing") + } + + if err := r.index.Reset(); err != nil { + return err + } + if _, err := r.index.EnsureVersion(); err != nil { + return err + } + + talkers, err := indexable.ListTalkers(ctx) + if err != nil { + return err + } + + if len(talkers) == 0 { + if err := r.index.UpdateFingerprint(fp); err != nil { + return err + } + return r.index.UpdateLastBuilt(time.Now()) + } + + sort.Strings(talkers) + + batch := make([]*model.Message, 0, 512) + flush := func() error { + if len(batch) == 0 { + return nil + } + if err := r.index.IndexMessages(batch); err != nil { + return err + } + batch = batch[:0] + return nil + } + + total := float64(len(talkers)) + for i, talker := range talkers { + if err := ctx.Err(); err != nil { + return err + } + + handler := func(msg *model.Message) error { + if msg == nil { + return nil + } + batch = append(batch, msg) + if len(batch) >= cap(batch) { + return flush() + } + return nil + } + + if err := indexable.IterateMessages(ctx, []string{talker}, handler); err != nil { + return err + } + + if err := flush(); err != nil { + return err + } + + r.updateIndexProgress(float64(i+1) / total) + } + + if err := r.index.UpdateFingerprint(fp); err != nil { + return err + } + if err := r.index.UpdateLastBuilt(time.Now()); err != nil { + return err + } + + return nil +} + +func (r *Repository) updateIndexProgress(progress float64) { + if r.index == nil { + return + } + + if progress < 0 { + progress = 0 + } + if progress > 1 { + progress = 1 + } + + r.indexMu.Lock() + r.indexStatus.Progress = progress + r.indexMu.Unlock() +} + +func (r *Repository) indexStatusSnapshot() *model.SearchIndexStatus { + if r.index == nil { + return nil + } + + r.indexMu.Lock() + status := r.indexStatus + r.indexMu.Unlock() + + copied := status + return &copied +} + +func (r *Repository) searchMessagesWithIndex(ctx context.Context, req *model.SearchRequest) (*model.SearchResponse, error) { + makeEmpty := func() *model.SearchResponse { + return &model.SearchResponse{ + Total: 0, + Hits: []*model.SearchHit{}, + DurationMs: 0, + Limit: req.Limit, + Offset: req.Offset, + Query: req.Query, + Talker: req.Talker, + Sender: req.Sender, + Start: req.Start, + End: req.End, + Index: r.indexStatusSnapshot(), + } + } + + if strings.TrimSpace(req.Query) == "" { + return makeEmpty(), nil + } + + if r.index == nil { + return makeEmpty(), nil + } + + ready, err := r.ensureIndex(ctx) + if err != nil { + return nil, err + } + if !ready { + return makeEmpty(), nil + } + + talkers := util.Str2List(req.Talker, ",") + senders := util.Str2List(req.Sender, ",") + + startUnix := int64(0) + if !req.Start.IsZero() { + startUnix = req.Start.Unix() + } + endUnix := int64(0) + if !req.End.IsZero() { + endUnix = req.End.Unix() + } + if startUnix > 0 && endUnix > 0 && endUnix < startUnix { + startUnix, endUnix = endUnix, startUnix + } + + begin := time.Now() + hits, total, err := r.index.Search(req, talkers, senders, startUnix, endUnix, req.Offset, req.Limit) + if err != nil { + return nil, err + } + + mapped := make([]*model.SearchHit, 0, len(hits)) + for _, hit := range hits { + if hit == nil || hit.Message == nil { + continue + } + mapped = append(mapped, &model.SearchHit{ + Message: hit.Message, + Snippet: hit.Snippet, + Score: hit.Score, + }) + } + + resp := &model.SearchResponse{ + Total: total, + Hits: mapped, + DurationMs: time.Since(begin).Milliseconds(), + Limit: req.Limit, + Offset: req.Offset, + Query: req.Query, + Talker: req.Talker, + Sender: req.Sender, + Start: req.Start, + End: req.End, + Index: r.indexStatusSnapshot(), + } + + return resp, nil +} diff --git a/internal/wechatdb/repository/repository.go b/internal/wechatdb/repository/repository.go index f1163f6..fbfd34f 100644 --- a/internal/wechatdb/repository/repository.go +++ b/internal/wechatdb/repository/repository.go @@ -2,6 +2,7 @@ package repository import ( "context" + "sync" "github.com/fsnotify/fsnotify" "github.com/rs/zerolog/log" @@ -9,12 +10,22 @@ import ( "github.com/sjzar/chatlog/internal/errors" "github.com/sjzar/chatlog/internal/model" "github.com/sjzar/chatlog/internal/wechatdb/datasource" + "github.com/sjzar/chatlog/internal/wechatdb/indexer" ) // Repository 实现了 repository.Repository 接口 type Repository struct { ds datasource.DataSource + indexPath string + index *indexer.Index + indexMu sync.Mutex + indexStatus model.SearchIndexStatus + indexFingerprint string + indexBuildMu sync.Mutex + indexCtx context.Context + indexCancel context.CancelFunc + // Cache for contact contactCache map[string]*model.Contact aliasToContact map[string][]*model.Contact @@ -39,9 +50,10 @@ type Repository struct { } // New 创建一个新的 Repository -func New(ds datasource.DataSource) (*Repository, error) { +func New(ds datasource.DataSource, indexPath string) (*Repository, error) { r := &Repository{ ds: ds, + indexPath: indexPath, contactCache: make(map[string]*model.Contact), aliasToContact: make(map[string][]*model.Contact), remarkToContact: make(map[string][]*model.Contact), @@ -67,6 +79,10 @@ func New(ds datasource.DataSource) (*Repository, error) { ds.SetCallback("contact", r.contactCallback) ds.SetCallback("chatroom", r.chatroomCallback) + if err := r.initIndex(); err != nil { + log.Warn().Err(err).Msg("init bleve index failed") + } + return r, nil } @@ -107,7 +123,23 @@ func (r *Repository) chatroomCallback(event fsnotify.Event) error { // Close 实现 Repository 接口的 Close 方法 func (r *Repository) Close() error { - return r.ds.Close() + if r.indexCancel != nil { + r.indexCancel() + r.indexCancel = nil + } + + var firstErr error + if r.index != nil { + if err := r.index.Close(); err != nil { + firstErr = err + } + r.index = nil + } + + if err := r.ds.Close(); err != nil && firstErr == nil { + firstErr = err + } + return firstErr } // GetAvatar proxies to datasource diff --git a/internal/wechatdb/repository/search.go b/internal/wechatdb/repository/search.go index 0367573..0f42918 100644 --- a/internal/wechatdb/repository/search.go +++ b/internal/wechatdb/repository/search.go @@ -25,12 +25,22 @@ func (r *Repository) SearchMessages(ctx context.Context, req *model.SearchReques nReq.Talker = normalizedTalker nReq.Sender = normalizedSender - resp, err := r.ds.SearchMessages(ctx, nReq) + if nReq.Limit <= 0 { + nReq.Limit = 20 + } + if nReq.Limit > 200 { + nReq.Limit = 200 + } + if nReq.Offset < 0 { + nReq.Offset = 0 + } + + resp, err := r.searchMessagesWithIndex(ctx, nReq) if err != nil { return nil, err } if resp == nil { - return &model.SearchResponse{Hits: []*model.SearchHit{}, Limit: nReq.Limit, Offset: nReq.Offset}, nil + resp = &model.SearchResponse{Hits: []*model.SearchHit{}, Limit: nReq.Limit, Offset: nReq.Offset} } // Enrich message metadata（头像、群昵称、显示名等） diff --git a/internal/wechatdb/wechatdb.go b/internal/wechatdb/wechatdb.go index d847fe3..a1d64c0 100644 --- a/internal/wechatdb/wechatdb.go +++ b/internal/wechatdb/wechatdb.go @@ -2,6 +2,7 @@ package wechatdb import ( "context" + "path/filepath" "time" "github.com/fsnotify/fsnotify" @@ -50,7 +51,8 @@ func (w *DB) Initialize() error { return err } - w.repo, err = repository.New(w.ds) + indexPath := filepath.Join(w.path, "indexes", "messages") + w.repo, err = repository.New(w.ds, indexPath) if err != nil { return err }