Add Phase 1 frxd implementation with conformance test suite
This commit is contained in:
George Coles
+359
@@ -0,0 +1,359 @@
|
||||
use std::fs;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Mutex;
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use tantivy::collector::{Count, TopDocs};
|
||||
use tantivy::directory::MmapDirectory;
|
||||
use tantivy::query::{AllQuery, BooleanQuery, Occur, Query, QueryParser, TermQuery};
|
||||
use tantivy::schema::{Field, IndexRecordOption, STORED, STRING, Schema, TEXT, Value};
|
||||
use tantivy::{Index, IndexReader, IndexWriter, TantivyDocument, Term, doc};
|
||||
|
||||
use crate::extract::{extract_file, summary_of, supported};
|
||||
use crate::message::{EXPOSURE_FULL, ResponseItem};
|
||||
|
||||
const WRITER_BUDGET: usize = 50_000_000;
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct Collection {
|
||||
pub name: String,
|
||||
pub path: String,
|
||||
pub shared: bool,
|
||||
pub exposure: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Default, Serialize, Deserialize)]
|
||||
struct CollectionsFile {
|
||||
#[serde(default, rename = "collection")]
|
||||
collections: Vec<Collection>,
|
||||
}
|
||||
|
||||
pub fn load_collections(manifest: &Path) -> Result<Vec<Collection>> {
|
||||
if !manifest.exists() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let raw = fs::read_to_string(manifest).context("reading collections manifest")?;
|
||||
let parsed: CollectionsFile = toml::from_str(&raw).context("parsing collections manifest")?;
|
||||
Ok(parsed.collections)
|
||||
}
|
||||
|
||||
pub fn save_collections(manifest: &Path, collections: &[Collection]) -> Result<()> {
|
||||
if let Some(parent) = manifest.parent() {
|
||||
fs::create_dir_all(parent)?;
|
||||
}
|
||||
let file = CollectionsFile {
|
||||
collections: collections.to_vec(),
|
||||
};
|
||||
fs::write(manifest, toml::to_string_pretty(&file)?)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct SearchHit {
|
||||
pub url: String,
|
||||
pub title: String,
|
||||
pub summary: String,
|
||||
pub published: String,
|
||||
pub exposure: String,
|
||||
pub collection: String,
|
||||
pub shared: bool,
|
||||
pub content: Option<String>,
|
||||
}
|
||||
|
||||
pub fn response_items(hits: &[SearchHit]) -> Vec<ResponseItem> {
|
||||
hits.iter()
|
||||
.map(|hit| ResponseItem {
|
||||
url: hit.url.clone(),
|
||||
title: hit.title.clone(),
|
||||
summary: hit.summary.clone(),
|
||||
published: hit.published.clone(),
|
||||
exposure: hit.exposure.clone(),
|
||||
content: if hit.exposure == EXPOSURE_FULL {
|
||||
hit.content.clone()
|
||||
} else {
|
||||
None
|
||||
},
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
struct Fields {
|
||||
url: Field,
|
||||
title: Field,
|
||||
body: Field,
|
||||
summary: Field,
|
||||
published: Field,
|
||||
exposure: Field,
|
||||
collection: Field,
|
||||
shared: Field,
|
||||
}
|
||||
|
||||
pub struct LocalIndex {
|
||||
index: Index,
|
||||
reader: IndexReader,
|
||||
writer: Mutex<IndexWriter>,
|
||||
fields: Fields,
|
||||
}
|
||||
|
||||
impl LocalIndex {
|
||||
pub fn open(dir: &Path) -> Result<Self> {
|
||||
fs::create_dir_all(dir)?;
|
||||
let schema = build_schema();
|
||||
let index = Index::open_or_create(MmapDirectory::open(dir)?, schema)
|
||||
.context("opening tantivy index")?;
|
||||
let schema = index.schema();
|
||||
let fields = Fields {
|
||||
url: schema.get_field("url")?,
|
||||
title: schema.get_field("title")?,
|
||||
body: schema.get_field("body")?,
|
||||
summary: schema.get_field("summary")?,
|
||||
published: schema.get_field("published")?,
|
||||
exposure: schema.get_field("exposure")?,
|
||||
collection: schema.get_field("collection")?,
|
||||
shared: schema.get_field("shared")?,
|
||||
};
|
||||
let reader = index.reader()?;
|
||||
let writer = index.writer::<TantivyDocument>(WRITER_BUDGET)?;
|
||||
Ok(Self {
|
||||
index,
|
||||
reader,
|
||||
writer: Mutex::new(writer),
|
||||
fields,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn add_collection(&self, collection: &Collection) -> Result<usize> {
|
||||
let root = PathBuf::from(&collection.path);
|
||||
{
|
||||
let writer = self.writer.lock().expect("writer lock");
|
||||
writer.delete_term(Term::from_field_text(
|
||||
self.fields.collection,
|
||||
&collection.name,
|
||||
));
|
||||
}
|
||||
let mut files = Vec::new();
|
||||
walk(&root, &mut files);
|
||||
let mut added = 0;
|
||||
for file in files {
|
||||
if self.add_file(collection, &file)? {
|
||||
added += 1;
|
||||
}
|
||||
}
|
||||
self.commit()?;
|
||||
Ok(added)
|
||||
}
|
||||
|
||||
pub fn add_file(&self, collection: &Collection, path: &Path) -> Result<bool> {
|
||||
let Some(extracted) = extract_file(path) else {
|
||||
return Ok(false);
|
||||
};
|
||||
let canonical = path.canonicalize().unwrap_or_else(|_| path.to_path_buf());
|
||||
let url = format!("file://{}", canonical.display());
|
||||
let summary = summary_of(&extracted.body);
|
||||
let shared = if collection.shared { "true" } else { "false" };
|
||||
let writer = self.writer.lock().expect("writer lock");
|
||||
writer.delete_term(Term::from_field_text(self.fields.url, &url));
|
||||
writer.add_document(doc!(
|
||||
self.fields.url => url.as_str(),
|
||||
self.fields.title => extracted.title.as_str(),
|
||||
self.fields.body => extracted.body.as_str(),
|
||||
self.fields.summary => summary.as_str(),
|
||||
self.fields.published => extracted.published.as_str(),
|
||||
self.fields.exposure => collection.exposure.as_str(),
|
||||
self.fields.collection => collection.name.as_str(),
|
||||
self.fields.shared => shared,
|
||||
))?;
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
pub fn commit(&self) -> Result<()> {
|
||||
self.writer.lock().expect("writer lock").commit()?;
|
||||
self.reader.reload()?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn search(
|
||||
&self,
|
||||
text: &str,
|
||||
limit: usize,
|
||||
only_shared: bool,
|
||||
) -> Result<(Vec<SearchHit>, u64)> {
|
||||
let limit = limit.max(1);
|
||||
self.reader.reload()?;
|
||||
let searcher = self.reader.searcher();
|
||||
let user_query: Box<dyn Query> = if text.trim().is_empty() {
|
||||
Box::new(AllQuery)
|
||||
} else {
|
||||
let parser =
|
||||
QueryParser::for_index(&self.index, vec![self.fields.title, self.fields.body]);
|
||||
let (query, _errors) = parser.parse_query_lenient(text);
|
||||
query
|
||||
};
|
||||
let query: Box<dyn Query> = if only_shared {
|
||||
let shared_term = TermQuery::new(
|
||||
Term::from_field_text(self.fields.shared, "true"),
|
||||
IndexRecordOption::Basic,
|
||||
);
|
||||
Box::new(BooleanQuery::new(vec![
|
||||
(Occur::Must, user_query),
|
||||
(Occur::Must, Box::new(shared_term)),
|
||||
]))
|
||||
} else {
|
||||
user_query
|
||||
};
|
||||
let total = searcher.search(&*query, &Count)? as u64;
|
||||
let top = searcher.search(&*query, &TopDocs::with_limit(limit).order_by_score())?;
|
||||
let mut hits = Vec::with_capacity(top.len());
|
||||
for (_score, address) in top {
|
||||
let document: TantivyDocument = searcher.doc(address)?;
|
||||
hits.push(SearchHit {
|
||||
url: text_value(&document, self.fields.url),
|
||||
title: text_value(&document, self.fields.title),
|
||||
summary: text_value(&document, self.fields.summary),
|
||||
published: text_value(&document, self.fields.published),
|
||||
exposure: text_value(&document, self.fields.exposure),
|
||||
collection: text_value(&document, self.fields.collection),
|
||||
shared: text_value(&document, self.fields.shared) == "true",
|
||||
content: Some(text_value(&document, self.fields.body)),
|
||||
});
|
||||
}
|
||||
Ok((hits, total))
|
||||
}
|
||||
|
||||
pub fn doc_count(&self) -> u64 {
|
||||
self.reader.searcher().num_docs()
|
||||
}
|
||||
}
|
||||
|
||||
fn text_value(document: &TantivyDocument, field: Field) -> String {
|
||||
document
|
||||
.get_first(field)
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or("")
|
||||
.to_string()
|
||||
}
|
||||
|
||||
fn build_schema() -> Schema {
|
||||
let mut builder = Schema::builder();
|
||||
builder.add_text_field("url", STRING | STORED);
|
||||
builder.add_text_field("title", TEXT | STORED);
|
||||
builder.add_text_field("body", TEXT | STORED);
|
||||
builder.add_text_field("summary", STORED);
|
||||
builder.add_text_field("published", STORED);
|
||||
builder.add_text_field("exposure", STRING | STORED);
|
||||
builder.add_text_field("collection", STRING | STORED);
|
||||
builder.add_text_field("shared", STRING | STORED);
|
||||
builder.build()
|
||||
}
|
||||
|
||||
fn walk(root: &Path, out: &mut Vec<PathBuf>) {
|
||||
let Ok(entries) = fs::read_dir(root) else {
|
||||
return;
|
||||
};
|
||||
for entry in entries.flatten() {
|
||||
let Ok(file_type) = entry.file_type() else {
|
||||
continue;
|
||||
};
|
||||
if file_type.is_symlink() {
|
||||
continue;
|
||||
}
|
||||
let name = entry.file_name();
|
||||
let name = name.to_string_lossy();
|
||||
if file_type.is_dir() {
|
||||
if name.starts_with('.') {
|
||||
continue;
|
||||
}
|
||||
walk(&entry.path(), out);
|
||||
} else if file_type.is_file() && supported(&entry.path()) {
|
||||
out.push(entry.path());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::message::EXPOSURE_METADATA;
|
||||
|
||||
fn collection(name: &str, root: &Path, shared: bool, exposure: &str) -> Collection {
|
||||
Collection {
|
||||
name: name.to_string(),
|
||||
path: root.display().to_string(),
|
||||
shared,
|
||||
exposure: exposure.to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn shared_filter_and_exposure_are_enforced() {
|
||||
let temp = tempfile::tempdir().unwrap();
|
||||
let shared_dir = temp.path().join("shared");
|
||||
let private_dir = temp.path().join("private");
|
||||
fs::create_dir_all(&shared_dir).unwrap();
|
||||
fs::create_dir_all(&private_dir).unwrap();
|
||||
fs::write(shared_dir.join("a.txt"), "rust ownership rules").unwrap();
|
||||
fs::write(private_dir.join("b.txt"), "rust secret notes").unwrap();
|
||||
|
||||
let index = LocalIndex::open(&temp.path().join("index")).unwrap();
|
||||
index
|
||||
.add_collection(&collection("shared", &shared_dir, true, EXPOSURE_FULL))
|
||||
.unwrap();
|
||||
index
|
||||
.add_collection(&collection(
|
||||
"private",
|
||||
&private_dir,
|
||||
false,
|
||||
EXPOSURE_METADATA,
|
||||
))
|
||||
.unwrap();
|
||||
|
||||
let (hits, total) = index.search("rust", 10, false).unwrap();
|
||||
assert_eq!(total, 2);
|
||||
assert_eq!(hits.len(), 2);
|
||||
|
||||
let (shared_hits, shared_total) = index.search("rust", 10, true).unwrap();
|
||||
assert_eq!(shared_total, 1);
|
||||
assert_eq!(shared_hits.len(), 1);
|
||||
assert!(shared_hits[0].url.contains("shared"));
|
||||
|
||||
let items = response_items(&shared_hits);
|
||||
assert!(items[0].content.is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reindex_removes_deleted_files() {
|
||||
let temp = tempfile::tempdir().unwrap();
|
||||
let dir = temp.path().join("corpus");
|
||||
fs::create_dir_all(&dir).unwrap();
|
||||
fs::write(dir.join("a.txt"), "rust one").unwrap();
|
||||
fs::write(dir.join("b.txt"), "rust two").unwrap();
|
||||
let index = LocalIndex::open(&temp.path().join("index")).unwrap();
|
||||
index
|
||||
.add_collection(&collection("test", &dir, true, EXPOSURE_FULL))
|
||||
.unwrap();
|
||||
assert_eq!(index.search("rust", 10, false).unwrap().1, 2);
|
||||
|
||||
fs::remove_file(dir.join("b.txt")).unwrap();
|
||||
index
|
||||
.add_collection(&collection("test", &dir, true, EXPOSURE_FULL))
|
||||
.unwrap();
|
||||
assert_eq!(index.search("rust", 10, false).unwrap().1, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn metadata_exposure_withholds_content() {
|
||||
let temp = tempfile::tempdir().unwrap();
|
||||
let dir = temp.path().join("corpus");
|
||||
fs::create_dir_all(&dir).unwrap();
|
||||
fs::write(dir.join("a.txt"), "metadata only").unwrap();
|
||||
let index = LocalIndex::open(&temp.path().join("index")).unwrap();
|
||||
index
|
||||
.add_collection(&collection("test", &dir, true, EXPOSURE_METADATA))
|
||||
.unwrap();
|
||||
let (hits, _) = index.search("metadata", 10, true).unwrap();
|
||||
let items = response_items(&hits);
|
||||
assert!(items[0].content.is_none());
|
||||
assert_eq!(items[0].exposure, EXPOSURE_METADATA);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user