feat: add persistent CPU search engine and CLI with lifecycle tests
Core CI / core (ubuntu-latest) (push) Successful in 1m0s
Core CI / core (macos-latest) (push) Canceled after 0s
Core CI / core (windows-latest) (push) Canceled after 0s
Core CI / core (ubuntu-latest) (push) Successful in 1m0s
Core CI / core (macos-latest) (push) Canceled after 0s
Core CI / core (windows-latest) (push) Canceled after 0s
This commit is contained in:
karl committed
1 parent
aae803771c
commit
f77112db97
19 files changed
+1688
No files matched your search
@@ -0,0 +1,482 @@
|
||||
//! Offline lexical search engine shared by the CLI and desktop shell.
|
||||
//!
|
||||
//! Reconciliation is atomic. Incomplete directory walks never delete records.
|
||||
//! Text extraction is deliberately bounded; binary document extractors belong in
|
||||
//! isolated workers in the next milestone, not in this indexing transaction.
|
||||
use anyhow::{bail, Context, Result};
|
||||
use rusqlite::{params, Connection, OptionalExtension};
|
||||
use serde::Serialize;
|
||||
use std::{
|
||||
fs,
|
||||
io::Read,
|
||||
path::{Path, PathBuf},
|
||||
time::UNIX_EPOCH,
|
||||
};
|
||||
use walkdir::WalkDir;
|
||||
|
||||
const SCHEMA_VERSION: i64 = 2;
|
||||
const MAX_TEXT_BYTES: u64 = 2 * 1024 * 1024;
|
||||
const MAX_RESULTS: usize = 1000;
|
||||
|
||||
#[derive(Debug, Default, Serialize)]
|
||||
pub struct IndexReport {
|
||||
pub discovered: usize,
|
||||
pub indexed: usize,
|
||||
pub removed: usize,
|
||||
pub skipped: usize,
|
||||
pub unchanged: usize,
|
||||
}
|
||||
|
||||
/// Names of directories to prune anywhere below an explicitly chosen root.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct IndexOptions {
|
||||
pub excluded_directory_names: Vec<String>,
|
||||
}
|
||||
|
||||
impl Default for IndexOptions {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
excluded_directory_names: vec![".git".into(), "node_modules".into(), "target".into()],
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Filters are independent of query syntax and apply before the result limit.
|
||||
/// The folder is canonicalized and includes descendants, with a path-boundary check.
|
||||
#[derive(Debug, Default, Clone)]
|
||||
pub struct SearchOptions {
|
||||
pub folder: Option<PathBuf>,
|
||||
pub extension: Option<String>,
|
||||
pub min_size: Option<u64>,
|
||||
pub max_size: Option<u64>,
|
||||
pub modified_after_unix_ms: Option<i64>,
|
||||
pub modified_before_unix_ms: Option<i64>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct SearchResult {
|
||||
pub path: String,
|
||||
pub filename: String,
|
||||
pub snippet: String,
|
||||
pub score: f64,
|
||||
pub match_kind: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct IndexStats {
|
||||
pub files: usize,
|
||||
pub content_files: usize,
|
||||
}
|
||||
|
||||
pub struct SearchIndex {
|
||||
connection: Connection,
|
||||
database_path: Option<PathBuf>,
|
||||
}
|
||||
|
||||
impl SearchIndex {
|
||||
pub fn open(path: impl AsRef<Path>) -> Result<Self> {
|
||||
let path = path.as_ref();
|
||||
let mut connection = Connection::open(path).context("open search database")?;
|
||||
let database_path = if path == Path::new(":memory:") {
|
||||
None
|
||||
} else {
|
||||
Some(fs::canonicalize(path).context("resolve database path")?)
|
||||
};
|
||||
connection.busy_timeout(std::time::Duration::from_secs(5))?;
|
||||
let version: i64 = connection.query_row("PRAGMA user_version", [], |row| row.get(0))?;
|
||||
if !(0..=SCHEMA_VERSION).contains(&version) {
|
||||
bail!("unsupported database version {version}; expected {SCHEMA_VERSION}");
|
||||
}
|
||||
if version != SCHEMA_VERSION {
|
||||
connection.execute_batch("PRAGMA journal_mode=WAL;")?;
|
||||
let initialization =
|
||||
connection.transaction_with_behavior(rusqlite::TransactionBehavior::Immediate)?;
|
||||
// Another opener may have completed initialization while we waited.
|
||||
let version: i64 =
|
||||
initialization.query_row("PRAGMA user_version", [], |row| row.get(0))?;
|
||||
if !(0..=SCHEMA_VERSION).contains(&version) {
|
||||
bail!("unsupported database version {version}");
|
||||
}
|
||||
initialization.execute_batch(
|
||||
"
|
||||
CREATE TABLE IF NOT EXISTS files (
|
||||
id INTEGER PRIMARY KEY, path TEXT NOT NULL UNIQUE,
|
||||
root TEXT NOT NULL, filename TEXT NOT NULL, content TEXT NOT NULL,
|
||||
size INTEGER NOT NULL DEFAULT -1, mtime INTEGER NOT NULL DEFAULT -1,
|
||||
extension TEXT NOT NULL DEFAULT '');
|
||||
CREATE INDEX IF NOT EXISTS files_root ON files(root);
|
||||
CREATE VIRTUAL TABLE IF NOT EXISTS search USING fts5(filename, path, content,
|
||||
content='files', content_rowid='id', tokenize='unicode61');
|
||||
CREATE TRIGGER IF NOT EXISTS files_insert AFTER INSERT ON files BEGIN
|
||||
INSERT INTO search(rowid, filename, path, content)
|
||||
VALUES (new.id, new.filename, new.path, new.content); END;
|
||||
CREATE TRIGGER IF NOT EXISTS files_delete AFTER DELETE ON files BEGIN
|
||||
INSERT INTO search(search, rowid, filename, path, content)
|
||||
VALUES ('delete', old.id, old.filename, old.path, old.content); END;
|
||||
CREATE TRIGGER IF NOT EXISTS files_update AFTER UPDATE ON files BEGIN
|
||||
INSERT INTO search(search, rowid, filename, path, content)
|
||||
VALUES ('delete', old.id, old.filename, old.path, old.content);
|
||||
INSERT INTO search(rowid, filename, path, content)
|
||||
VALUES (new.id, new.filename, new.path, new.content); END;
|
||||
",
|
||||
)?;
|
||||
if version == 1 {
|
||||
initialization.execute_batch(
|
||||
"
|
||||
ALTER TABLE files ADD COLUMN size INTEGER NOT NULL DEFAULT -1;
|
||||
ALTER TABLE files ADD COLUMN mtime INTEGER NOT NULL DEFAULT -1;
|
||||
ALTER TABLE files ADD COLUMN extension TEXT NOT NULL DEFAULT '';
|
||||
PRAGMA user_version=2;",
|
||||
)?;
|
||||
} else {
|
||||
initialization.execute_batch("PRAGMA user_version=2;")?;
|
||||
}
|
||||
initialization.commit()?;
|
||||
}
|
||||
Ok(Self {
|
||||
connection,
|
||||
database_path,
|
||||
})
|
||||
}
|
||||
|
||||
/// Reconcile one explicitly selected root. Symlinks are never followed.
|
||||
/// Every file receives a filename record, including unsupported formats.
|
||||
pub fn index_folder(&mut self, root: impl AsRef<Path>) -> Result<IndexReport> {
|
||||
self.index_folder_with_options(root, &IndexOptions::default())
|
||||
}
|
||||
|
||||
pub fn index_folder_with_options(
|
||||
&mut self,
|
||||
root: impl AsRef<Path>,
|
||||
options: &IndexOptions,
|
||||
) -> Result<IndexReport> {
|
||||
let root = fs::canonicalize(root).context("resolve index root")?;
|
||||
if !root.is_dir() {
|
||||
bail!("index root must be a directory");
|
||||
}
|
||||
let root_key = root.to_string_lossy().into_owned();
|
||||
let transaction = self.connection.transaction()?;
|
||||
transaction.execute_batch(
|
||||
"CREATE TEMP TABLE IF NOT EXISTS seen_paths(path TEXT PRIMARY KEY);
|
||||
DELETE FROM seen_paths;",
|
||||
)?;
|
||||
let mut report = IndexReport::default();
|
||||
for entry in WalkDir::new(&root)
|
||||
.follow_links(false)
|
||||
.into_iter()
|
||||
.filter_entry(|entry| {
|
||||
entry.depth() == 0
|
||||
|| !entry.file_type().is_dir()
|
||||
|| !options.excluded_directory_names.iter().any(|name| {
|
||||
entry
|
||||
.file_name()
|
||||
.to_string_lossy()
|
||||
.eq_ignore_ascii_case(name)
|
||||
})
|
||||
})
|
||||
{
|
||||
let entry = entry.context("directory traversal failed; reconciliation rolled back")?;
|
||||
if !entry.file_type().is_file() {
|
||||
continue;
|
||||
}
|
||||
if let Some(database) = &self.database_path {
|
||||
let path_key = entry.path().to_string_lossy();
|
||||
let db_key = database.to_string_lossy();
|
||||
if path_key == db_key
|
||||
|| path_key == format!("{db_key}-wal")
|
||||
|| path_key == format!("{db_key}-shm")
|
||||
{
|
||||
continue;
|
||||
}
|
||||
}
|
||||
let path = entry.path().to_string_lossy().into_owned();
|
||||
let filename = entry.file_name().to_string_lossy().into_owned();
|
||||
transaction.execute("INSERT INTO seen_paths(path) VALUES (?1)", [&path])?;
|
||||
report.discovered += 1;
|
||||
let metadata = entry.metadata()?;
|
||||
let size = i64::try_from(metadata.len()).context("file exceeds supported size")?;
|
||||
let modified = metadata
|
||||
.modified()?
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.context("file timestamp predates Unix epoch")?;
|
||||
let mtime = i64::try_from(modified.as_nanos())
|
||||
.context("file timestamp exceeds supported range")?;
|
||||
let extension = entry
|
||||
.path()
|
||||
.extension()
|
||||
.and_then(|value| value.to_str())
|
||||
.unwrap_or("")
|
||||
.to_ascii_lowercase();
|
||||
let previous: Option<(i64, i64, String)> = transaction
|
||||
.query_row(
|
||||
"SELECT size,mtime,root FROM files WHERE path=?1",
|
||||
[&path],
|
||||
|row| Ok((row.get(0)?, row.get(1)?, row.get(2)?)),
|
||||
)
|
||||
.optional()?;
|
||||
if previous == Some((size, mtime, root_key.clone())) {
|
||||
report.unchanged += 1;
|
||||
continue;
|
||||
}
|
||||
let content = extract_text(entry.path()).with_context(|| format!("read {path}"))?;
|
||||
if content.is_none() {
|
||||
report.skipped += 1;
|
||||
}
|
||||
transaction.execute(
|
||||
"INSERT INTO files(path,root,filename,content,size,mtime,extension) VALUES (?1,?2,?3,?4,?5,?6,?7)
|
||||
ON CONFLICT(path) DO UPDATE SET root=excluded.root, filename=excluded.filename,
|
||||
content=excluded.content, size=excluded.size, mtime=excluded.mtime,
|
||||
extension=excluded.extension",
|
||||
params![path, root_key, filename, content.unwrap_or_default(), size, mtime, extension],
|
||||
)?;
|
||||
report.indexed += 1;
|
||||
}
|
||||
report.removed = transaction.execute(
|
||||
"DELETE FROM files WHERE root=?1
|
||||
AND NOT EXISTS (SELECT 1 FROM seen_paths WHERE seen_paths.path=files.path)",
|
||||
[&root_key],
|
||||
)?;
|
||||
transaction.execute("DELETE FROM seen_paths", [])?;
|
||||
transaction.commit()?;
|
||||
Ok(report)
|
||||
}
|
||||
|
||||
/// Literal input is escaped before FTS parsing; user text cannot inject operators.
|
||||
/// Exact filenames precede partial filenames, which precede content matches.
|
||||
/// Filename substring matching currently scans metadata; a dedicated substring
|
||||
/// index is required before claiming million-file latency targets.
|
||||
pub fn search(&self, query: &str, limit: usize) -> Result<Vec<SearchResult>> {
|
||||
self.search_with_options(query, limit, &SearchOptions::default())
|
||||
}
|
||||
|
||||
pub fn search_with_options(
|
||||
&self,
|
||||
query: &str,
|
||||
limit: usize,
|
||||
options: &SearchOptions,
|
||||
) -> Result<Vec<SearchResult>> {
|
||||
if options
|
||||
.min_size
|
||||
.zip(options.max_size)
|
||||
.is_some_and(|(min, max)| min > max)
|
||||
{
|
||||
bail!("minimum size exceeds maximum size");
|
||||
}
|
||||
if options
|
||||
.modified_after_unix_ms
|
||||
.zip(options.modified_before_unix_ms)
|
||||
.is_some_and(|(min, max)| min > max)
|
||||
{
|
||||
bail!("modified-after exceeds modified-before");
|
||||
}
|
||||
let folder = options
|
||||
.folder
|
||||
.as_ref()
|
||||
.map(fs::canonicalize)
|
||||
.transpose()?
|
||||
.map(|path| path.to_string_lossy().into_owned());
|
||||
let folder_prefix = folder.as_ref().map(|value| {
|
||||
format!(
|
||||
"{}{}",
|
||||
value.trim_end_matches(['/', '\\']),
|
||||
std::path::MAIN_SEPARATOR
|
||||
)
|
||||
});
|
||||
let extension = options
|
||||
.extension
|
||||
.as_ref()
|
||||
.map(|value| value.trim_start_matches('.').to_ascii_lowercase());
|
||||
let min_size = options
|
||||
.min_size
|
||||
.map(i64::try_from)
|
||||
.transpose()
|
||||
.context("minimum size exceeds supported range")?;
|
||||
let max_size = options
|
||||
.max_size
|
||||
.map(i64::try_from)
|
||||
.transpose()
|
||||
.context("maximum size exceeds supported range")?;
|
||||
let query = query.trim();
|
||||
if query.is_empty() || limit == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
if query.len() > 4096 {
|
||||
bail!("query exceeds 4096 bytes");
|
||||
}
|
||||
let fts = query
|
||||
.split_whitespace()
|
||||
.map(|term| format!("\"{}\"", term.replace('"', "\"\"")))
|
||||
.collect::<Vec<_>>()
|
||||
.join(" AND ");
|
||||
let mut statement = self.connection.prepare(
|
||||
"WITH hits AS (
|
||||
SELECT rowid AS id, bm25(search, 8.0, 2.0, 1.0) AS rank,
|
||||
snippet(search, 2, '', '', ' … ', 24) AS excerpt
|
||||
FROM search WHERE search MATCH ?1)
|
||||
SELECT f.path, f.filename, COALESCE(h.excerpt, ''), COALESCE(h.rank, 0),
|
||||
CASE WHEN lower(f.filename)=lower(?2) THEN 0
|
||||
WHEN instr(lower(f.filename), lower(?2)) > 0 THEN 1 ELSE 2 END AS tier
|
||||
FROM files f LEFT JOIN hits h ON h.id=f.id
|
||||
WHERE (h.id IS NOT NULL OR instr(lower(f.filename), lower(?2)) > 0)
|
||||
AND (?4 IS NULL OR substr(f.path,1,length(?4))=?4)
|
||||
AND (?5 IS NULL OR f.extension=?5)
|
||||
AND (?6 IS NULL OR f.size>=?6) AND (?7 IS NULL OR f.size<=?7)
|
||||
AND (?8 IS NULL OR f.mtime/1000000>=?8) AND (?9 IS NULL OR f.mtime/1000000<=?9)
|
||||
ORDER BY tier, h.rank, f.path LIMIT ?3",
|
||||
)?;
|
||||
let results = statement
|
||||
.query_map(
|
||||
params![
|
||||
fts,
|
||||
query,
|
||||
limit.min(MAX_RESULTS) as i64,
|
||||
folder_prefix,
|
||||
extension,
|
||||
min_size,
|
||||
max_size,
|
||||
options.modified_after_unix_ms,
|
||||
options.modified_before_unix_ms
|
||||
],
|
||||
|row| {
|
||||
let tier: i64 = row.get(4)?;
|
||||
Ok(SearchResult {
|
||||
path: row.get(0)?,
|
||||
filename: row.get(1)?,
|
||||
snippet: row.get(2)?,
|
||||
score: -row.get::<_, f64>(3)?,
|
||||
match_kind: match tier {
|
||||
0 => "exact_filename",
|
||||
1 => "filename",
|
||||
_ => "content",
|
||||
}
|
||||
.to_owned(),
|
||||
})
|
||||
},
|
||||
)?
|
||||
.collect::<rusqlite::Result<Vec<_>>>()?;
|
||||
Ok(results)
|
||||
}
|
||||
|
||||
pub fn stats(&self) -> Result<IndexStats> {
|
||||
Ok(self.connection.query_row(
|
||||
"SELECT count(*), count(NULLIF(content, '')) FROM files",
|
||||
[],
|
||||
|row| {
|
||||
Ok(IndexStats {
|
||||
files: row.get(0)?,
|
||||
content_files: row.get(1)?,
|
||||
})
|
||||
},
|
||||
)?)
|
||||
}
|
||||
}
|
||||
|
||||
fn extract_text(path: &Path) -> Result<Option<String>> {
|
||||
let extension = path
|
||||
.extension()
|
||||
.and_then(|value| value.to_str())
|
||||
.unwrap_or("")
|
||||
.to_ascii_lowercase();
|
||||
if !matches!(
|
||||
extension.as_str(),
|
||||
"txt"
|
||||
| "md"
|
||||
| "rs"
|
||||
| "py"
|
||||
| "js"
|
||||
| "ts"
|
||||
| "tsx"
|
||||
| "jsx"
|
||||
| "json"
|
||||
| "toml"
|
||||
| "yaml"
|
||||
| "yml"
|
||||
| "csv"
|
||||
| "html"
|
||||
| "css"
|
||||
| "xml"
|
||||
| "log"
|
||||
| "sh"
|
||||
) {
|
||||
return Ok(None);
|
||||
}
|
||||
let file = fs::File::open(path)?;
|
||||
if file.metadata()?.len() > MAX_TEXT_BYTES {
|
||||
return Ok(None);
|
||||
}
|
||||
let mut bytes = Vec::new();
|
||||
file.take(MAX_TEXT_BYTES + 1).read_to_end(&mut bytes)?;
|
||||
if bytes.len() as u64 > MAX_TEXT_BYTES || bytes.contains(&0) {
|
||||
return Ok(None);
|
||||
}
|
||||
Ok(String::from_utf8(bytes).ok())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
#[test]
|
||||
fn persistence_updates_deletes_and_exact_priority() -> Result<()> {
|
||||
let folder = tempfile::tempdir()?;
|
||||
let root = folder.path().join("docs");
|
||||
fs::create_dir(&root)?;
|
||||
let db = folder.path().join("index.db");
|
||||
fs::write(root.join("contract.txt"), "supplier cancellation agreement")?;
|
||||
fs::write(root.join("other.txt"), "contract contract contract")?;
|
||||
let mut index = SearchIndex::open(&db)?;
|
||||
assert_eq!(index.index_folder(&root)?.indexed, 2);
|
||||
assert_eq!(
|
||||
index.search("contract.txt", 10)?[0].match_kind,
|
||||
"exact_filename"
|
||||
);
|
||||
assert_eq!(index.search("cancellation", 10)?.len(), 1);
|
||||
drop(index);
|
||||
let mut index = SearchIndex::open(&db)?;
|
||||
assert_eq!(index.stats()?.files, 2);
|
||||
fs::write(root.join("contract.txt"), "replacement document")?;
|
||||
fs::rename(root.join("other.txt"), root.join("renamed.txt"))?;
|
||||
let report = index.index_folder(&root)?;
|
||||
assert_eq!(report.removed, 1);
|
||||
assert!(index.search("cancellation", 10)?.is_empty());
|
||||
assert!(index.search("other.txt", 10)?.is_empty());
|
||||
assert_eq!(index.search("renamed.txt", 10)?.len(), 1);
|
||||
fs::remove_file(root.join("contract.txt"))?;
|
||||
index.index_folder(&root)?;
|
||||
assert_eq!(index.stats()?.files, 1);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn safe_queries_limits_and_unsupported_files() -> Result<()> {
|
||||
let root = tempfile::tempdir()?;
|
||||
fs::write(
|
||||
root.path().join("notes.md"),
|
||||
"German Kündigung and English cancellation",
|
||||
)?;
|
||||
fs::write(root.path().join("photo.jpg"), [0, 1, 2])?;
|
||||
let mut index = SearchIndex::open(":memory:")?;
|
||||
assert_eq!(index.index_folder(root.path())?.skipped, 1);
|
||||
assert_eq!(index.search("Kündigung", 1)?.len(), 1);
|
||||
assert!(index.search("\" OR *", 10)?.is_empty());
|
||||
assert!(index.search("", 10)?.is_empty());
|
||||
assert!(index.search("notes", 0)?.is_empty());
|
||||
assert_eq!(index.search("photo.jpg", 10)?.len(), 1);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invalid_root_preserves_index_and_unknown_schema_is_rejected() -> Result<()> {
|
||||
let folder = tempfile::tempdir()?;
|
||||
fs::write(folder.path().join("notes.txt"), "persistent evidence")?;
|
||||
let mut index = SearchIndex::open(":memory:")?;
|
||||
index.index_folder(folder.path())?;
|
||||
assert!(index.index_folder(folder.path().join("missing")).is_err());
|
||||
assert_eq!(index.search("evidence", 10)?.len(), 1);
|
||||
let db = folder.path().join("future.db");
|
||||
let connection = Connection::open(&db)?;
|
||||
connection.execute_batch("PRAGMA user_version=999")?;
|
||||
drop(connection);
|
||||
assert!(SearchIndex::open(db).is_err());
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user