use std::path::Path;
use anyhow::{Context, Result};
use rusqlite::Connection;
use crate::model::CodeNode;
use crate::search::tokenize::extract_keywords;
const CREATE_ENTITIES_V2: &str = "CREATE VIRTUAL TABLE IF NOT EXISTS entities USING fts5(
name,
kind,
signature,
source,
file_path,
node_json,
tokens
);";
pub struct SearchStore {
conn: Connection,
}
impl SearchStore {
pub fn open(path: impl AsRef<Path>) -> Result<(Self, bool)> {
let conn = Connection::open(path.as_ref())
.context("打开 SQLite 数据库失败")?;
conn.pragma_update(None, "journal_mode", "WAL")
.context("设置 WAL 模式失败")?;
conn.busy_timeout(std::time::Duration::from_secs(5))
.context("设置 busy_timeout 失败")?;
let version: i64 = conn
.query_row("PRAGMA user_version", [], |row| row.get(0))
.context("读取 user_version 失败")?;
let table_exists: bool = conn
.query_row(
"SELECT EXISTS(SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = 'entities')",
[],
|row| row.get(0),
)
.context("查询 entities 表存在性失败")?;
if !table_exists {
conn.execute_batch(CREATE_ENTITIES_V2)
.context("创建 FTS5 表失败")?;
conn.pragma_update(None, "user_version", 2)
.context("写入 user_version 失败")?;
return Ok((Self { conn }, false));
}
if version < 2 {
conn.execute_batch("DROP TABLE IF EXISTS entities;")
.context("删除旧 FTS5 表失败")?;
conn.execute_batch(CREATE_ENTITIES_V2)
.context("重建 FTS5 表失败")?;
conn.pragma_update(None, "user_version", 2)
.context("写入 user_version 失败")?;
return Ok((Self { conn }, true));
}
Ok((Self { conn }, false))
}
fn cjk_tokens(parts: &[&str]) -> String {
let mut out: Vec<String> = Vec::new();
for part in parts {
for k in extract_keywords(part) {
if k.chars().all(|c| matches!(c as u32, 0x4E00..=0x9FFF | 0x3400..=0x4DBF | 0xF900..=0xFAFF)) {
out.push(k);
}
}
}
out.join(" ")
}
pub fn insert_entities_batch(&self, items: &[(CodeNode, String)]) -> Result<()> {
let mut stmt = self.conn.prepare(
"INSERT INTO entities (name, kind, signature, source, file_path, node_json, tokens)
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)"
).context("准备 FTS5 插入语句失败")?;
for (node, source) in items {
let node_json = serde_json::to_string(node)
.context("序列化 CodeNode 失败")?;
let tokens = Self::cjk_tokens(&[
&node.name,
node.signature.as_deref().unwrap_or(""),
source,
]);
stmt.execute(rusqlite::params![
node.name,
node.kind.as_str(),
node.signature.as_deref().unwrap_or(""),
source,
crate::incremental::norm_sep(node.file_path.as_deref().unwrap_or("")).as_str(),
node_json,
tokens,
]).context("插入 FTS5 条目失败")?;
}
Ok(())
}
fn build_match_terms(query: &str) -> Vec<String> {
extract_keywords(query)
}
pub fn search_fts(&self, query: &str, limit: usize) -> Result<Vec<(CodeNode, f64)>> {
let terms = Self::build_match_terms(query);
if terms.is_empty() {
return Ok(Vec::new());
}
let match_expr = format!("{{name signature source tokens}} : ({})", terms.join(" OR "));
let sql = format!(
"SELECT node_json, bm25(entities) as rank
FROM entities
WHERE entities MATCH ?1
ORDER BY rank
LIMIT {}",
limit
);
let mut stmt = self.conn.prepare(&sql)
.context("准备 FTS5 查询语句失败")?;
let rows = match stmt.query_map(rusqlite::params![match_expr], |row| {
let node_json: String = row.get(0)?;
let rank: f64 = row.get(1)?;
Ok((node_json, rank))
}) {
Ok(rows) => rows,
Err(e) => {
tracing::warn!("FTS5 查询语法错误,返回空结果: {} (query: {})", e, query);
return Ok(Vec::new());
}
};
let mut results = Vec::new();
for row in rows {
let (node_json, rank) = row.context("读取 FTS5 结果行失败")?;
if let Ok(node) = serde_json::from_str::<CodeNode>(&node_json) {
results.push((node, -rank));
}
}
Ok(results)
}
pub fn delete_entities_by_file(&self, file_path: &str) -> Result<usize> {
let count = self.conn.execute(
"DELETE FROM entities WHERE file_path = ?1",
rusqlite::params![crate::incremental::norm_sep(file_path)],
).context("删除 FTS5 条目失败")?;
Ok(count)
}
pub fn entity_count(&self) -> Result<usize> {
let count: usize = self.conn.query_row(
"SELECT COUNT(*) FROM entities",
[],
|row| row.get(0),
).context("查询 FTS5 文档数失败")?;
Ok(count)
}
pub fn clear_entities(&self) -> Result<()> {
self.conn.execute("DELETE FROM entities", [])
.context("清空 FTS5 表失败")?;
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::model::{NodeId, NodeKind};
fn tmp_db_path(label: &str) -> std::path::PathBuf {
let mut p = std::env::temp_dir();
p.push(format!("store_test_{}_{}.db", label, std::process::id()));
let _ = std::fs::remove_file(&p);
p
}
fn make_node(name: &str, file: &str) -> CodeNode {
CodeNode {
id: NodeId::new(0),
kind: NodeKind::Function,
name: name.into(),
file_path: Some(file.into()),
line_range: Some((1, 10)),
doc_comment: None,
signature: Some(format!("fn {}()", name)), visibility: None,
module_path: vec![],
}
}
#[test]
fn test_fts_insert_and_search() {
let (store, need_reindex) = SearchStore::open(tmp_db_path("fts")).unwrap();
assert!(!need_reindex, "新库无需重建");
let items = vec![
(make_node("authenticate", "src/auth.rs"), "fn authenticate(user: &str)".to_string()),
(make_node("save_session", "src/storage.rs"), "fn save_session(id: u64)".to_string()),
];
store.insert_entities_batch(&items).unwrap();
assert_eq!(store.entity_count().unwrap(), 2);
let results = store.search_fts("authenticate", 5).unwrap();
assert!(!results.is_empty());
assert_eq!(results[0].0.name, "authenticate");
}
#[test]
fn test_fts_delete_by_file() {
let (store, _) = SearchStore::open(tmp_db_path("fts_del")).unwrap();
let items = vec![
(make_node("alpha", "src/a.rs"), "alpha code".to_string()),
(make_node("beta", "src/b.rs"), "beta code".to_string()),
];
store.insert_entities_batch(&items).unwrap();
let removed = store.delete_entities_by_file("src/a.rs").unwrap();
assert_eq!(removed, 1);
assert_eq!(store.entity_count().unwrap(), 1);
}
#[test]
fn test_fts_cjk_substring_search() {
let (store, _) = SearchStore::open(tmp_db_path("fts_cjk")).unwrap();
let items = vec![
(make_node("提取配置", "src/config.rs"), "fn 提取配置() 读取合并后的配置".to_string()),
(make_node("save_session", "src/storage.rs"), "fn save_session(id: u64)".to_string()),
];
store.insert_entities_batch(&items).unwrap();
let results = store.search_fts("配置", 5).unwrap();
assert_eq!(results.len(), 1, "中文 2-gram 应命中实体");
assert_eq!(results[0].0.name, "提取配置");
let mixed = store.search_fts("配置 session", 5).unwrap();
assert_eq!(mixed.len(), 2, "混合查询应同时命中中文与英文实体");
}
#[test]
fn test_fts_legacy_schema_migration() {
let path = tmp_db_path("fts_migrate");
let _ = std::fs::remove_file(&path);
{
let conn = rusqlite::Connection::open(&path).unwrap();
conn.execute_batch(
"CREATE VIRTUAL TABLE entities USING fts5(
name, kind, signature, source, file_path, node_json
);"
).unwrap();
conn.pragma_update(None, "user_version", 1).unwrap();
conn.execute(
"INSERT INTO entities (name, kind, signature, source, file_path, node_json)
VALUES ('old_fn', 'function', 'fn old_fn()', 'old code', 'src/old.rs', '{}')",
[],
).unwrap();
}
let (store, need_reindex) = SearchStore::open(&path).unwrap();
assert!(need_reindex, "旧 schema 必须触发重建标记");
assert_eq!(store.entity_count().unwrap(), 0, "旧索引数据已清空");
let items = vec![
(make_node("验证迁移", "src/new.rs"), "fn 验证迁移()".to_string()),
];
store.insert_entities_batch(&items).unwrap();
let results = store.search_fts("迁移", 5).unwrap();
assert_eq!(results.len(), 1, "迁移后 v2 检索正常");
assert_eq!(results[0].0.name, "验证迁移");
let (_, need_reindex2) = SearchStore::open(&path).unwrap();
assert!(!need_reindex2, "二次 open 不应再触发重建");
}
#[test]
fn test_fts_punctuation_query_returns_empty() {
let (store, _) = SearchStore::open(tmp_db_path("fts_punct")).unwrap();
let items = vec![(make_node("alpha", "src/a.rs"), "alpha code".to_string())];
store.insert_entities_batch(&items).unwrap();
let empty = store.search_fts("", 5).unwrap();
assert!(empty.is_empty(), "空串查询返回空");
let punct = store.search_fts("!!!---", 5).unwrap();
assert!(punct.is_empty(), "纯标点查询返回空(v1 会语法错误上抛)");
}
}