use std::collections::{HashMap, HashSet};
use std::path::{Path, PathBuf};
use rusqlite::{Connection, OptionalExtension, params};
use serde::{Deserialize, Serialize};
use crate::embedder::ModelIdentity;
use crate::embedding::{from_bytes, to_bytes};
use crate::lexical::{FIELDS, LexicalDoc, LexicalQuery, order_free_text};
use crate::{
Error, Modality, ResourceId, ResourceLocator, ResourceMetadata, ResourceVersion, Result,
};
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct CacheProfile {
pub model: ModelIdentity,
pub segmenter_version: String,
}
impl CacheProfile {
pub fn id(&self) -> String {
let json = serde_json::to_vec(self).expect("profile serialises");
blake3::hash(&json).to_hex()[..24].to_string()
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct CachedSegment {
pub locator: ResourceLocator,
pub heading: Option<String>,
pub input_hash: String,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CachedResource {
pub metadata: ResourceMetadata,
pub version: ResourceVersion,
pub descriptor_hash: Option<String>,
pub segments: Vec<CachedSegment>,
}
#[derive(Debug, Clone)]
pub struct SegmentWrite<'a> {
pub segment: &'a CachedSegment,
pub lexical_body: &'a str,
}
#[derive(Debug, Clone)]
pub struct ResourceWrite<'a> {
pub metadata: &'a ResourceMetadata,
pub version: &'a ResourceVersion,
pub descriptor_hash: Option<&'a str>,
pub segments: Vec<SegmentWrite<'a>>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CachedResourceState {
pub metadata: ResourceMetadata,
pub version: ResourceVersion,
pub descriptor_hash: Option<String>,
}
#[derive(Debug, Clone)]
pub struct StoredVector {
pub resource: ResourceId,
pub ordinal: usize,
pub segment: CachedSegment,
pub vector: Vec<f32>,
}
#[derive(Debug, Clone, PartialEq)]
pub struct LexicalHit {
pub resource: ResourceId,
pub ordinal: usize,
pub score: f32,
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct ProfileStats {
pub profile_id: String,
pub profile: Option<CacheProfile>,
pub resources: u64,
pub segments: u64,
pub embeddings: u64,
#[serde(default)]
pub image_codes: u64,
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct CacheStats {
pub location: Option<PathBuf>,
pub size_bytes: Option<u64>,
#[serde(skip)]
pub scrub_pending: bool,
pub profiles: Vec<ProfileStats>,
#[serde(default)]
pub backup_excluded: Option<bool>,
#[serde(default)]
pub in_synced_folder: bool,
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct CacheCleanup {
pub profiles_scrubbed: usize,
pub pending: bool,
}
impl CacheCleanup {
pub fn is_empty(&self) -> bool {
self.profiles_scrubbed == 0 && !self.pending
}
}
pub trait EmbeddingCache: Send {
fn scrub(&mut self) -> Result<CacheCleanup> {
Ok(CacheCleanup::default())
}
fn scrub_if_idle(&mut self, _quiet: std::time::Duration) -> Result<Option<CacheCleanup>> {
Ok(None)
}
fn scrub_delay(&self, quiet: std::time::Duration) -> std::time::Duration {
quiet
}
fn register_profile(&mut self, id: &str, profile: &CacheProfile) -> Result<()>;
fn resource_index(&self, profile: &str) -> Result<HashMap<ResourceId, CachedResourceState>>;
fn resource(&self, profile: &str, path: &ResourceId) -> Result<Option<CachedResource>>;
fn put_resource(
&mut self,
profile: &str,
path: &ResourceId,
resource: &ResourceWrite<'_>,
) -> Result<()>;
fn touch_resource(
&mut self,
profile: &str,
path: &ResourceId,
metadata: &ResourceMetadata,
version: &ResourceVersion,
) -> Result<()>;
fn remove_resources(&mut self, profile: &str, paths: &[ResourceId]) -> Result<()>;
fn known_inputs(&self, profile: &str, hashes: &[String]) -> Result<HashSet<String>>;
fn put_embeddings(&mut self, profile: &str, items: &[(String, Vec<f32>)]) -> Result<()>;
fn supports_image_codes(&self) -> bool {
false
}
fn sync_image_codes(&mut self, _profile: &str, _paths: &HashSet<ResourceId>) -> Result<()> {
Err(Error::Config(
"q4 image scan requires a persistent SQLite cache; --no-cache is unsupported".into(),
))
}
fn image_scan_rows(
&self,
_profile: &str,
_paths: &HashSet<ResourceId>,
) -> Result<Vec<ScanVector>> {
Err(Error::Config(
"q4 image scan is unsupported by this cache".into(),
))
}
fn selected_image_vectors(
&self,
_profile: &str,
_hashes: &HashSet<String>,
) -> Result<HashMap<String, Vec<f32>>> {
Err(Error::Config(
"q4 image scan is unsupported by this cache".into(),
))
}
fn vectors_for(&self, profile: &str, paths: &HashSet<ResourceId>) -> Result<Vec<StoredVector>>;
fn lexical_search(
&self,
profile: &str,
query: &LexicalQuery,
paths: &HashSet<ResourceId>,
limit: usize,
) -> Result<Vec<LexicalHit>>;
fn lexical_matched_terms(
&self,
profile: &str,
query: &LexicalQuery,
keys: &[(ResourceId, usize)],
) -> Result<Vec<Vec<String>>>;
fn prune_orphans(&mut self, profile: &str) -> Result<u64>;
fn stats(&self) -> Result<CacheStats>;
fn clear(&mut self) -> Result<()>;
}
#[derive(Debug, Default)]
pub struct MemoryCache {
profiles: HashMap<String, CacheProfile>,
resources: HashMap<String, HashMap<ResourceId, CachedResource>>,
vectors: HashMap<String, HashMap<String, Vec<f32>>>,
lexical: HashMap<String, HashMap<ResourceId, Vec<LexicalDoc>>>,
}
impl MemoryCache {
pub fn new() -> Self {
Self::default()
}
}
impl EmbeddingCache for MemoryCache {
fn register_profile(&mut self, id: &str, profile: &CacheProfile) -> Result<()> {
self.profiles.insert(id.to_string(), profile.clone());
Ok(())
}
fn resource_index(&self, profile: &str) -> Result<HashMap<ResourceId, CachedResourceState>> {
Ok(self
.resources
.get(profile)
.into_iter()
.flatten()
.map(|(path, r)| {
(
path.clone(),
CachedResourceState {
metadata: r.metadata.clone(),
version: r.version.clone(),
descriptor_hash: r.descriptor_hash.clone(),
},
)
})
.collect())
}
fn resource(&self, profile: &str, path: &ResourceId) -> Result<Option<CachedResource>> {
Ok(self
.resources
.get(profile)
.and_then(|b| b.get(path))
.cloned())
}
fn put_resource(
&mut self,
profile: &str,
path: &ResourceId,
resource: &ResourceWrite<'_>,
) -> Result<()> {
let docs = lexical_docs(path, resource);
let stored = CachedResource {
metadata: resource.metadata.clone(),
version: resource.version.clone(),
descriptor_hash: resource.descriptor_hash.map(str::to_owned),
segments: resource
.segments
.iter()
.map(|s| s.segment.clone())
.collect(),
};
self.resources
.entry(profile.to_string())
.or_default()
.insert(path.clone(), stored);
self.lexical
.entry(profile.to_string())
.or_default()
.insert(path.clone(), docs);
Ok(())
}
fn touch_resource(
&mut self,
profile: &str,
path: &ResourceId,
metadata: &ResourceMetadata,
version: &ResourceVersion,
) -> Result<()> {
if let Some(r) = self
.resources
.get_mut(profile)
.and_then(|b| b.get_mut(path))
{
r.metadata = metadata.clone();
r.version = version.clone();
}
Ok(())
}
fn remove_resources(&mut self, profile: &str, paths: &[ResourceId]) -> Result<()> {
if let Some(bucket) = self.resources.get_mut(profile) {
for p in paths {
bucket.remove(p);
}
}
if let Some(bucket) = self.lexical.get_mut(profile) {
for p in paths {
bucket.remove(p);
}
}
Ok(())
}
fn known_inputs(&self, profile: &str, hashes: &[String]) -> Result<HashSet<String>> {
let Some(vectors) = self.vectors.get(profile) else {
return Ok(HashSet::new());
};
Ok(hashes
.iter()
.filter(|h| vectors.contains_key(*h))
.cloned()
.collect())
}
fn put_embeddings(&mut self, profile: &str, items: &[(String, Vec<f32>)]) -> Result<()> {
let vectors = self.vectors.entry(profile.to_string()).or_default();
for (h, v) in items {
vectors.insert(h.clone(), v.clone());
}
Ok(())
}
fn vectors_for(&self, profile: &str, paths: &HashSet<ResourceId>) -> Result<Vec<StoredVector>> {
let (Some(resources), Some(vectors)) =
(self.resources.get(profile), self.vectors.get(profile))
else {
return Ok(Vec::new());
};
let mut out = Vec::new();
for (path, r) in resources {
if !paths.contains(path) {
continue;
}
for (ordinal, s) in r.segments.iter().enumerate() {
if let Some(v) = vectors.get(&s.input_hash) {
out.push(StoredVector {
resource: path.clone(),
ordinal,
segment: s.clone(),
vector: v.clone(),
});
}
}
}
Ok(out)
}
fn lexical_search(
&self,
profile: &str,
query: &LexicalQuery,
paths: &HashSet<ResourceId>,
limit: usize,
) -> Result<Vec<LexicalHit>> {
let Some(bucket) = self.lexical.get(profile) else {
return Ok(Vec::new());
};
let docs = bucket
.iter()
.flat_map(|(path, docs)| docs.iter().enumerate().map(move |(i, d)| ((path, i), d)));
let mut scored = crate::lexical::bm25(docs, query);
scored.sort_by(|a, b| {
b.1.total_cmp(&a.1)
.then_with(|| resource_key(a.0.0).cmp(&resource_key(b.0.0)))
.then_with(|| a.0.1.cmp(&b.0.1))
});
Ok(scored
.into_iter()
.filter(|((path, _), _)| paths.contains(*path))
.take(limit)
.map(|((path, ordinal), score)| LexicalHit {
resource: path.clone(),
ordinal,
score: score as f32,
})
.collect())
}
fn lexical_matched_terms(
&self,
profile: &str,
query: &LexicalQuery,
keys: &[(ResourceId, usize)],
) -> Result<Vec<Vec<String>>> {
Ok(keys
.iter()
.map(|(path, ordinal)| {
self.lexical
.get(profile)
.and_then(|bucket| bucket.get(path))
.and_then(|docs| docs.get(*ordinal))
.map(|doc| doc.matched_terms(query))
.unwrap_or_default()
})
.collect())
}
fn prune_orphans(&mut self, profile: &str) -> Result<u64> {
let Some(vectors) = self.vectors.get_mut(profile) else {
return Ok(0);
};
let live: HashSet<&str> = self
.resources
.get(profile)
.into_iter()
.flat_map(|b| b.values())
.flat_map(|r| r.segments.iter().map(|s| s.input_hash.as_str()))
.collect();
let before = vectors.len();
vectors.retain(|h, _| live.contains(h.as_str()));
Ok((before - vectors.len()) as u64)
}
fn stats(&self) -> Result<CacheStats> {
let profiles = self
.profiles
.iter()
.map(|(id, p)| {
let resources = self.resources.get(id);
ProfileStats {
profile_id: id.clone(),
profile: Some(p.clone()),
resources: resources.map_or(0, |b| b.len() as u64),
segments: resources
.into_iter()
.flat_map(|b| b.values())
.map(|r| r.segments.len() as u64)
.sum(),
embeddings: self.vectors.get(id).map_or(0, |b| b.len() as u64),
image_codes: 0,
}
})
.collect();
Ok(CacheStats {
location: None,
size_bytes: None,
backup_excluded: None,
in_synced_folder: false,
scrub_pending: false,
profiles,
})
}
fn clear(&mut self) -> Result<()> {
*self = Self::default();
Ok(())
}
}
fn lexical_docs(path: &ResourceId, resource: &ResourceWrite<'_>) -> Vec<LexicalDoc> {
let (name, full) = path.labels();
resource
.segments
.iter()
.map(|s| LexicalDoc::new(&name, &full, s.segment.heading.as_deref(), s.lexical_body))
.collect()
}
pub(crate) mod image_codes;
pub use image_codes::{IMAGE_CODEC_ID, ScanData, ScanVector};
mod import;
pub use import::{CacheImportMapping, CacheImportReport};
pub const SCHEMA_VERSION: &str = "7";
fn fts_table(profile: &str) -> String {
let hex: String = profile
.as_bytes()
.iter()
.map(|b| format!("{b:02x}"))
.collect();
format!("\"segments_fts_p_{hex}\"")
}
fn create_profile_fts(tx: &rusqlite::Transaction<'_>, profile: &str) -> Result<()> {
tx.execute_batch(&format!(
"CREATE VIRTUAL TABLE IF NOT EXISTS {} USING fts5(
name, path, heading, body, detail=full, columnsize=1,
tokenize='unicode61 remove_diacritics 2');",
fts_table(profile)
))?;
Ok(())
}
fn migrate_shared_fts(tx: &rusqlite::Transaction<'_>, profiles: &[String]) -> Result<()> {
let missing: bool = tx.query_row(
"SELECT EXISTS(SELECT 1 FROM segments s JOIN resources r
ON r.profile=s.profile AND r.path=s.path
WHERE s.rowid NOT IN (SELECT rowid FROM segments_fts)
AND NOT (r.content_hash='' AND r.modified_ns IS NULL
AND (r.version IS NULL OR r.version='null') AND r.descriptor_hash IS NULL))",
[],
|r| r.get(0),
)?;
if missing {
return Err(Error::Config(
"cache migration: live segment is missing its keyword row".into(),
));
}
tx.execute_batch(
"CREATE VIRTUAL TABLE temp.migration_vocab USING fts5vocab(main, segments_fts, instance);",
)?;
let mut docs: HashMap<i64, (String, [Vec<String>; 4])> = tx
.prepare("SELECT s.rowid,s.profile FROM segments s JOIN segments_fts f ON f.rowid=s.rowid")?
.query_map([], |r| {
Ok((r.get(0)?, (r.get(1)?, std::array::from_fn(|_| Vec::new()))))
})?
.collect::<rusqlite::Result<_>>()?;
{
let mut stmt = tx.prepare("SELECT term,doc,col FROM temp.migration_vocab")?;
let mut rows = stmt.query([])?;
while let Some(row) = rows.next()? {
let doc = row.get::<_, i64>(1)?;
let Some((_, fields)) = docs.get_mut(&doc) else {
continue;
};
let col: String = row.get(2)?;
let index = ["name", "path", "heading", "body"]
.iter()
.position(|f| *f == col)
.ok_or_else(|| Error::Config("cache migration: invalid keyword field".into()))?;
fields[index].push(row.get(0)?);
}
}
for (rowid, (profile, fields)) in &mut docs {
let fields = fields.each_mut().map(|terms| {
terms.sort();
terms.join(" ")
});
tx.prepare_cached(&format!(
"INSERT INTO {}(rowid,name,path,heading,body) VALUES(?1,?2,?3,?4,?5)",
fts_table(profile)
))?
.execute(params![rowid, fields[0], fields[1], fields[2], fields[3]])?;
}
for profile in profiles {
let table = fts_table(profile);
tx.execute_batch(&format!(
"CREATE VIRTUAL TABLE temp.migration_new USING fts5vocab(main, {table}, instance);"
))?;
let different: bool = tx.query_row(
"WITH old AS (SELECT term,doc,col,offset FROM temp.migration_vocab
JOIN segments s ON s.rowid=doc WHERE s.profile=?1),
missing AS (SELECT * FROM old EXCEPT SELECT * FROM temp.migration_new),
added AS (SELECT * FROM temp.migration_new EXCEPT SELECT * FROM old)
SELECT EXISTS(SELECT 1 FROM missing UNION ALL SELECT 1 FROM added)",
[profile],
|r| r.get(0),
)?;
let count: i64 =
tx.query_row(&format!("SELECT count(*) FROM {table}"), [], |r| r.get(0))?;
if different || count != docs.values().filter(|(p, _)| p == profile).count() as i64 {
return Err(Error::Config(
"cache migration: keyword round trip failed".into(),
));
}
tx.execute_batch("DROP TABLE temp.migration_new;")?;
}
tx.execute_batch("DROP TABLE temp.migration_vocab; DROP TABLE segments_fts;")?;
Ok(())
}
pub struct SqliteCache {
conn: Connection,
path: PathBuf,
guards: super::cache_guards::GuardStatus,
last_change: std::time::Instant,
last_attempt: Option<std::time::Instant>,
observed_generation: Option<String>,
#[cfg(test)]
after_vector_read: Option<Box<dyn FnMut() + Send>>,
}
impl SqliteCache {
pub fn default_path() -> Result<PathBuf> {
let base = dirs::cache_dir()
.ok_or_else(|| Error::Config("no cache directory on this platform".into()))?;
Ok(base.join("seyk").join("cache.sqlite3"))
}
pub fn open_default() -> Result<Self> {
Self::open(Self::default_path()?)
}
pub fn open(path: impl Into<PathBuf>) -> Result<Self> {
Self::open_with_guards(path.into(), super::cache_guards::prepare)
}
fn open_with_guards(
path: PathBuf,
prepare: impl FnOnce(&Path, bool) -> super::cache_guards::GuardStatus,
) -> Result<Self> {
Self::open_internal(path, prepare, true)
}
fn open_internal(
path: PathBuf,
prepare: impl FnOnce(&Path, bool) -> super::cache_guards::GuardStatus,
initialize: bool,
) -> Result<Self> {
let dir = super::cache_guards::parent_dir(&path);
let created = match std::fs::create_dir(dir) {
Ok(()) => true,
Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
std::fs::create_dir_all(super::cache_guards::parent_dir(dir))
.map_err(|e| Error::io(dir, e))?;
match std::fs::create_dir(dir) {
Ok(()) => true,
Err(e) if e.kind() == std::io::ErrorKind::AlreadyExists && dir.is_dir() => {
false
}
Err(e) => return Err(Error::io(dir, e)),
}
}
Err(e) if e.kind() == std::io::ErrorKind::AlreadyExists && dir.is_dir() => false,
Err(e) => return Err(Error::io(dir, e)),
};
let owned = created || super::cache_guards::is_default_dir(dir);
if owned {
restrict_permissions(dir, 0o700);
}
let guards = prepare(dir, owned);
let conn = Connection::open(&path)?;
restrict_permissions(&path, 0o600);
conn.pragma_update(None, "journal_mode", "WAL")?;
conn.pragma_update(None, "synchronous", "NORMAL")?;
conn.pragma_update(None, "foreign_keys", "ON")?;
conn.pragma_update(None, "secure_delete", "ON")?;
conn.busy_timeout(std::time::Duration::from_secs(10))?;
let mut cache = Self {
conn,
path,
guards,
last_change: std::time::Instant::now(),
last_attempt: None,
observed_generation: None,
#[cfg(test)]
after_vector_read: None,
};
if !initialize {
return Ok(cache);
}
cache.migrate()?;
cache.observed_generation = cache
.conn
.query_row(
"SELECT value FROM meta WHERE key='wal_scrub_pending'",
[],
|r| r.get(0),
)
.optional()?;
cache.last_change = std::time::Instant::now();
Ok(cache)
}
pub fn path(&self) -> &Path {
&self.path
}
fn cleanup_pending(&self) -> Result<bool> {
Ok(self.conn.query_row(
"SELECT EXISTS(SELECT 1 FROM meta WHERE key IN ('wal_scrub_pending', 'compact_pending') OR key GLOB 'fts_scrub_pending:*')",
[], |r| r.get(0),
)?)
}
fn run_scrub(&mut self) -> Result<CacheCleanup> {
self.last_attempt = Some(std::time::Instant::now());
if !self.cleanup_pending()? {
return Ok(CacheCleanup::default());
}
self.conn.busy_timeout(std::time::Duration::ZERO)?;
let mut report = CacheCleanup::default();
let result = self.scrub_pending_inner(&mut report);
self.conn.busy_timeout(std::time::Duration::from_secs(10))?;
match result {
Ok(()) => Ok(report),
Err(Error::Cache(error)) if sqlite_busy(&error) => {
report.pending = true;
Ok(report)
}
Err(error) => Err(error),
}
}
fn scrub_pending_inner(&mut self, report: &mut CacheCleanup) -> Result<()> {
let tx = self
.conn
.transaction_with_behavior(rusqlite::TransactionBehavior::Immediate)?;
let profiles = tx
.prepare("SELECT key,value FROM meta WHERE key GLOB 'fts_scrub_pending:*'")?
.query_map([], |r| Ok((r.get::<_, String>(0)?, r.get::<_, String>(1)?)))?
.collect::<rusqlite::Result<Vec<_>>>()?;
for (key, _) in &profiles {
optimize_fts(&tx, key.strip_prefix("fts_scrub_pending:").unwrap())?;
}
let generation: Option<String> = tx
.query_row(
"SELECT value FROM meta WHERE key='wal_scrub_pending'",
[],
|r| r.get(0),
)
.optional()?;
tx.commit()?;
report.profiles_scrubbed = profiles.len();
let Some(generation) = generation else {
report.pending = self.cleanup_pending()?;
return Ok(());
};
let compact: bool = self.conn.query_row(
"SELECT EXISTS(SELECT 1 FROM meta WHERE key='compact_pending')",
[],
|r| r.get(0),
)?;
if compact {
self.conn.flush_prepared_statement_cache();
self.conn.execute_batch("VACUUM")?;
}
let busy: i64 = self
.conn
.query_row("PRAGMA wal_checkpoint(TRUNCATE)", [], |r| r.get(0))?;
if busy != 0 {
report.pending = true;
return Ok(());
}
let tx = self
.conn
.transaction_with_behavior(rusqlite::TransactionBehavior::Immediate)?;
for (key, value) in &profiles {
tx.execute(
"DELETE FROM meta WHERE key=?1 AND value=?2",
params![key, value],
)?;
}
let finished = tx.execute(
"DELETE FROM meta WHERE key='wal_scrub_pending' AND value=?1",
[generation],
)?;
if finished != 0 {
tx.execute("DELETE FROM meta WHERE key='compact_pending'", [])?;
}
tx.commit()?;
report.pending = self.cleanup_pending()?;
Ok(())
}
fn record_write(&mut self, changed: bool, generation: Option<String>) {
if changed {
self.last_change = std::time::Instant::now();
self.observed_generation = generation;
}
}
fn migrate(&mut self) -> Result<()> {
let tx = self
.conn
.transaction_with_behavior(rusqlite::TransactionBehavior::Immediate)?;
Self::initialize_schema(&tx)?;
tx.commit()?;
Ok(())
}
fn initialize_schema(tx: &rusqlite::Transaction<'_>) -> Result<()> {
let has_meta: bool = tx.query_row(
"SELECT EXISTS(SELECT 1 FROM sqlite_master WHERE name = 'meta')",
[],
|r| r.get(0),
)?;
let version: Option<String> = if has_meta {
tx.query_row(
"SELECT value FROM meta WHERE key = 'schema_version'",
[],
|r| r.get(0),
)
.optional()?
} else {
None
};
if version.as_deref() == Some(SCHEMA_VERSION) {
return Ok(());
}
if version.as_deref() == Some("6") {
tx.execute_batch(image_codes::CREATE)?;
tx.execute(
"UPDATE meta SET value=?1 WHERE key='schema_version'",
[SCHEMA_VERSION],
)?;
return Ok(());
}
if let Some(version) = &version {
if !matches!(version.as_str(), "1" | "2" | "3" | "4" | "5") {
return Err(Error::Config(format!(
"cache schema {version} is newer than this binary supports; use a newer seyk or a separate --cache-path"
)));
}
} else {
let nonempty: bool = tx.query_row(
"SELECT EXISTS(SELECT 1 FROM sqlite_master WHERE name NOT LIKE 'sqlite_%')",
[],
|r| r.get(0),
)?;
if nonempty {
return Err(Error::Config(
"nonempty cache is missing schema_version metadata".into(),
));
}
}
tx.execute_batch(
"CREATE TABLE IF NOT EXISTS meta (key TEXT PRIMARY KEY, value TEXT NOT NULL);",
)?;
tx.execute_batch(
"CREATE TABLE IF NOT EXISTS profiles (
id TEXT PRIMARY KEY,
profile TEXT NOT NULL,
dimension INTEGER NOT NULL
);
CREATE TABLE IF NOT EXISTS resources (
profile TEXT NOT NULL REFERENCES profiles(id) ON DELETE CASCADE,
path BLOB NOT NULL,
modality TEXT NOT NULL,
size_bytes INTEGER NOT NULL,
modified_ns INTEGER,
content_hash TEXT NOT NULL,
version TEXT,
PRIMARY KEY (profile, path)
);
CREATE TABLE IF NOT EXISTS segments (
id INTEGER PRIMARY KEY,
profile TEXT NOT NULL,
path BLOB NOT NULL,
ordinal INTEGER NOT NULL,
locator TEXT NOT NULL,
heading TEXT,
input_hash TEXT NOT NULL,
UNIQUE (profile, path, ordinal),
FOREIGN KEY (profile, path) REFERENCES resources(profile, path) ON DELETE CASCADE
);
CREATE INDEX IF NOT EXISTS segments_by_input ON segments(profile, input_hash);
CREATE TABLE IF NOT EXISTS embeddings (
profile TEXT NOT NULL REFERENCES profiles(id) ON DELETE CASCADE,
input_hash TEXT NOT NULL,
vector BLOB NOT NULL,
PRIMARY KEY (profile, input_hash)
);
",
)?;
tx.execute_batch(image_codes::CREATE)?;
let has_media_type: bool = tx.query_row(
"SELECT EXISTS(SELECT 1 FROM pragma_table_info('resources') WHERE name = 'media_type')",
[],
|r| r.get(0),
)?;
if !has_media_type {
tx.execute_batch("ALTER TABLE resources ADD COLUMN media_type TEXT;")?;
}
let has_version: bool = tx.query_row(
"SELECT EXISTS(SELECT 1 FROM pragma_table_info('resources') WHERE name = 'version')",
[],
|r| r.get(0),
)?;
if !has_version {
tx.execute_batch("ALTER TABLE resources ADD COLUMN version TEXT;")?;
}
let has_descriptor: bool = tx.query_row(
"SELECT EXISTS(SELECT 1 FROM pragma_table_info('resources') WHERE name = 'descriptor_hash')",
[], |r| r.get(0),
)?;
if !has_descriptor {
tx.execute_batch("ALTER TABLE resources ADD COLUMN descriptor_hash TEXT;")?;
}
if matches!(version.as_deref(), Some("1" | "2" | "3")) {
tx.execute_batch(
r#"UPDATE resources SET version = CASE
WHEN modified_ns IS NOT NULL AND content_hash != ''
THEN '"fs:' || size_bytes || ':' || modified_ns || '"'
ELSE 'null' END;"#,
)?;
}
if version.is_some() {
tx.execute_batch(
"CREATE TABLE segments_v6 (
id INTEGER PRIMARY KEY,
profile TEXT NOT NULL, path BLOB NOT NULL, ordinal INTEGER NOT NULL,
locator TEXT NOT NULL, heading TEXT, input_hash TEXT NOT NULL,
UNIQUE(profile, path, ordinal),
FOREIGN KEY(profile, path) REFERENCES resources(profile, path) ON DELETE CASCADE
);
INSERT INTO segments_v6 SELECT rowid, profile, path, ordinal, locator, heading, input_hash FROM segments;
DROP TABLE segments;
ALTER TABLE segments_v6 RENAME TO segments;
CREATE INDEX segments_by_input ON segments(profile, input_hash);",
)?;
}
let profiles = tx
.prepare("SELECT id FROM profiles")?
.query_map([], |r| r.get::<_, String>(0))?
.collect::<rusqlite::Result<Vec<_>>>()?;
for profile in &profiles {
create_profile_fts(tx, profile)?;
}
if matches!(version.as_deref(), Some("1" | "2")) {
tx.execute_batch(
"DROP TABLE IF EXISTS segments_fts;
UPDATE resources SET modified_ns = NULL, content_hash = '', version = 'null', descriptor_hash = NULL;",
)?;
} else if version.is_some() {
migrate_shared_fts(tx, &profiles)?;
}
tx.execute(
"INSERT OR REPLACE INTO meta (key, value) VALUES ('schema_version', ?1)",
[SCHEMA_VERSION],
)?;
if version.is_some() {
tx.execute_batch("INSERT OR IGNORE INTO meta VALUES('compact_pending','1');")?;
mark_cleanup(tx, None)?;
}
Ok(())
}
}
pub(crate) fn resource_key(id: &ResourceId) -> Vec<u8> {
match id.to_file_path() {
Some(path) => path_key(&path),
None => std::iter::once(0).chain(id.as_str().bytes()).collect(),
}
}
fn key_resource(bytes: Vec<u8>) -> rusqlite::Result<ResourceId> {
let result = if let Some(uri) = bytes.strip_prefix(&[0]) {
std::str::from_utf8(uri)
.map_err(|e| Error::Config(e.to_string()))
.and_then(str::parse)
} else {
ResourceId::from_path(key_path(bytes))
};
result.map_err(|e| {
rusqlite::Error::FromSqlConversionFailure(0, rusqlite::types::Type::Blob, Box::new(e))
})
}
fn decode_version(value: Option<String>) -> rusqlite::Result<ResourceVersion> {
value
.map(|v| {
serde_json::from_str(&v).map_err(|e| {
rusqlite::Error::FromSqlConversionFailure(
0,
rusqlite::types::Type::Text,
Box::new(e),
)
})
})
.transpose()
.map(Option::unwrap_or_default)
}
#[cfg(unix)]
pub(crate) fn path_key(p: &Path) -> Vec<u8> {
use std::os::unix::ffi::OsStrExt;
p.as_os_str().as_bytes().to_vec()
}
#[cfg(unix)]
fn key_path(b: Vec<u8>) -> PathBuf {
use std::os::unix::ffi::OsStringExt;
PathBuf::from(std::ffi::OsString::from_vec(b))
}
#[cfg(windows)]
pub(crate) fn path_key(p: &Path) -> Vec<u8> {
use std::path::{Component, Prefix};
let raw = p.to_string_lossy();
match p.components().next() {
Some(Component::Prefix(prefix)) if matches!(prefix.kind(), Prefix::Disk(_)) => {
format!(r"\\?\{raw}").into_bytes()
}
_ => raw.into_owned().into_bytes(),
}
}
#[cfg(not(any(unix, windows)))]
pub(crate) fn path_key(p: &Path) -> Vec<u8> {
p.to_string_lossy().into_owned().into_bytes()
}
#[cfg(not(unix))]
fn key_path(b: Vec<u8>) -> PathBuf {
PathBuf::from(String::from_utf8_lossy(&b).into_owned())
}
#[cfg(unix)]
fn restrict_permissions(p: &Path, mode: u32) {
use std::os::unix::fs::PermissionsExt;
let _ = std::fs::set_permissions(p, std::fs::Permissions::from_mode(mode));
}
#[cfg(not(unix))]
fn restrict_permissions(_: &Path, _: u32) {}
fn sqlite_busy(error: &rusqlite::Error) -> bool {
matches!(error, rusqlite::Error::SqliteFailure(code, _)
if matches!(code.code, rusqlite::ErrorCode::DatabaseBusy | rusqlite::ErrorCode::DatabaseLocked))
}
fn mark_cleanup(tx: &rusqlite::Transaction<'_>, profile: Option<&str>) -> Result<String> {
let generation: String = tx.query_row("SELECT lower(hex(randomblob(16)))", [], |r| r.get(0))?;
tx.execute(
"INSERT INTO meta VALUES('wal_scrub_pending',?1)
ON CONFLICT(key) DO UPDATE SET value=excluded.value",
[&generation],
)?;
if let Some(profile) = profile {
tx.execute(
"INSERT INTO meta VALUES(?1,?2) ON CONFLICT(key) DO UPDATE SET value=excluded.value",
params![format!("fts_scrub_pending:{profile}"), generation],
)?;
}
Ok(generation)
}
fn refresh_pending_activity(tx: &rusqlite::Transaction<'_>) -> Result<Option<String>> {
Ok(tx.query_row(
"UPDATE meta SET value=lower(hex(randomblob(16))) WHERE key='wal_scrub_pending' RETURNING value",
[], |r| r.get(0),
).optional()?)
}
fn delete_fts_rows(tx: &rusqlite::Transaction, profile: &str, key: &[u8]) -> Result<bool> {
let deleted = tx
.prepare_cached(&format!(
"DELETE FROM {} WHERE rowid IN
(SELECT rowid FROM segments WHERE profile = ?1 AND path = ?2)",
fts_table(profile)
))?
.execute(params![profile, key])?;
Ok(deleted != 0)
}
fn optimize_fts(tx: &rusqlite::Transaction, profile: &str) -> Result<()> {
let table = fts_table(profile);
tx.execute_batch(&format!("INSERT INTO {table}({table}) VALUES('optimize');"))?;
Ok(())
}
fn parse_modality(s: &str) -> Modality {
serde_json::from_value(serde_json::Value::String(s.to_string())).unwrap_or(Modality::Text)
}
fn json_err(e: serde_json::Error) -> Error {
Error::Config(format!("corrupt cache entry: {e}"))
}
impl EmbeddingCache for SqliteCache {
fn scrub(&mut self) -> Result<CacheCleanup> {
self.run_scrub()
}
fn scrub_if_idle(&mut self, quiet: std::time::Duration) -> Result<Option<CacheCleanup>> {
let now = std::time::Instant::now();
let generation: Option<String> = self
.conn
.query_row(
"SELECT value FROM meta WHERE key='wal_scrub_pending'",
[],
|r| r.get(0),
)
.optional()?;
if generation != self.observed_generation {
self.observed_generation = generation;
self.last_change = now;
}
if now.duration_since(self.last_change) < quiet
|| self
.last_attempt
.is_some_and(|attempt| now.duration_since(attempt) < quiet)
|| !self.cleanup_pending()?
{
return Ok(None);
}
self.run_scrub().map(Some)
}
fn scrub_delay(&self, quiet: std::time::Duration) -> std::time::Duration {
let remaining = quiet.saturating_sub(self.last_change.elapsed()).max(
self.last_attempt
.map_or(std::time::Duration::ZERO, |attempt| {
quiet.saturating_sub(attempt.elapsed())
}),
);
if remaining.is_zero() {
quiet
} else {
remaining
}
}
fn register_profile(&mut self, id: &str, profile: &CacheProfile) -> Result<()> {
let tx = self
.conn
.transaction_with_behavior(rusqlite::TransactionBehavior::Immediate)?;
let inserted = tx.execute(
"INSERT OR IGNORE INTO profiles (id, profile, dimension) VALUES (?1, ?2, ?3)",
params![
id,
serde_json::to_string(profile).map_err(json_err)?,
profile.model.dimension.get() as i64
],
)?;
if inserted == 0 {
let table_exists: bool = tx.query_row(
"SELECT EXISTS(SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = ?1)",
[fts_table(id).trim_matches('"')],
|r| r.get(0),
)?;
if !table_exists {
return Err(Error::Config(
"corrupt cache: registered profile is missing its keyword table".into(),
));
}
} else {
create_profile_fts(&tx, id)?;
}
tx.commit()?;
Ok(())
}
fn resource_index(&self, profile: &str) -> Result<HashMap<ResourceId, CachedResourceState>> {
let mut stmt = self.conn.prepare_cached(
"SELECT path, modality, size_bytes, modified_ns, content_hash, version, descriptor_hash, media_type FROM resources WHERE profile = ?1",
)?;
let rows = stmt.query_map([profile], |r| {
let id = key_resource(r.get(0)?)?;
let metadata = ResourceMetadata {
media_type: r.get(7)?,
modality: parse_modality(&r.get::<_, String>(1)?),
size_bytes: r.get::<_, i64>(2)? as u64,
modified_ns: r.get(3)?,
content_hash: r.get(4)?,
};
let version = decode_version(r.get(5)?)?;
Ok((
id,
CachedResourceState {
metadata,
version,
descriptor_hash: r.get(6)?,
},
))
})?;
Ok(rows.collect::<rusqlite::Result<_>>()?)
}
fn resource(&self, profile: &str, path: &ResourceId) -> Result<Option<CachedResource>> {
let key = resource_key(path);
let metadata = self
.conn
.query_row(
"SELECT modality, size_bytes, modified_ns, content_hash, version, descriptor_hash, media_type FROM resources WHERE profile = ?1 AND path = ?2",
params![profile, key],
|r| {
Ok((ResourceMetadata {
media_type: r.get(6)?,
modality: parse_modality(&r.get::<_, String>(0)?),
size_bytes: r.get::<_, i64>(1)? as u64,
modified_ns: r.get(2)?,
content_hash: r.get(3)?,
}, decode_version(r.get(4)?)?, r.get::<_, Option<String>>(5)?))
},
)
.optional()?;
let Some((metadata, version, descriptor_hash)) = metadata else {
return Ok(None);
};
let mut stmt = self.conn.prepare_cached(
"SELECT locator, heading, input_hash FROM segments WHERE profile = ?1 AND path = ?2 ORDER BY ordinal",
)?;
let segments = stmt
.query_map(params![profile, key], |r| {
Ok((
r.get::<_, String>(0)?,
r.get::<_, Option<String>>(1)?,
r.get::<_, String>(2)?,
))
})?
.map(|row| {
let (loc, heading, input_hash) = row?;
Ok(CachedSegment {
locator: serde_json::from_str(&loc).map_err(json_err)?,
heading,
input_hash,
})
})
.collect::<Result<Vec<_>>>()?;
Ok(Some(CachedResource {
metadata,
version,
descriptor_hash,
segments,
}))
}
fn put_resource(
&mut self,
profile: &str,
path: &ResourceId,
resource: &ResourceWrite<'_>,
) -> Result<()> {
let key = resource_key(path);
let tx = self
.conn
.transaction_with_behavior(rusqlite::TransactionBehavior::Immediate)?;
let deleted = delete_fts_rows(&tx, profile, &key)?;
let replaced = tx.execute(
"DELETE FROM resources WHERE profile = ?1 AND path = ?2",
params![profile, key],
)?;
let m = &resource.metadata;
tx.execute(
"INSERT INTO resources (profile, path, modality, size_bytes, modified_ns, content_hash, version, descriptor_hash, media_type)
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9)",
params![
profile,
key,
m.modality.to_string(),
m.size_bytes as i64,
m.modified_ns,
m.content_hash,
serde_json::to_string(&resource.version).map_err(json_err)?,
resource.descriptor_hash,
m.media_type
],
)?;
{
let mut stmt = tx.prepare_cached(
"INSERT INTO segments (profile, path, ordinal, locator, heading, input_hash)
VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
)?;
let mut fts = tx.prepare_cached(&format!(
"INSERT INTO {} (rowid, name, path, heading, body) VALUES (?1, ?2, ?3, ?4, ?5)",
fts_table(profile)
))?;
let (name, full) = path.labels();
let name = order_free_text(&name);
let full = order_free_text(&full);
for (i, s) in resource.segments.iter().enumerate() {
stmt.execute(params![
profile,
key,
i as i64,
serde_json::to_string(&s.segment.locator).map_err(json_err)?,
s.segment.heading,
s.segment.input_hash
])?;
fts.execute(params![
tx.last_insert_rowid(),
name,
full,
order_free_text(s.segment.heading.as_deref().unwrap_or("")),
order_free_text(s.lexical_body)
])?;
}
}
let generation = if deleted || replaced != 0 {
Some(mark_cleanup(&tx, deleted.then_some(profile))?)
} else {
refresh_pending_activity(&tx)?
};
tx.commit()?;
self.record_write(true, generation);
Ok(())
}
fn touch_resource(
&mut self,
profile: &str,
path: &ResourceId,
metadata: &ResourceMetadata,
version: &ResourceVersion,
) -> Result<()> {
let tx = self
.conn
.transaction_with_behavior(rusqlite::TransactionBehavior::Immediate)?;
let changed = tx.execute(
"UPDATE resources SET modified_ns = ?3, size_bytes = ?4, version = ?5, content_hash = ?6, modality = ?7, media_type = ?8 WHERE profile = ?1 AND path = ?2
AND (modified_ns IS NOT ?3 OR size_bytes IS NOT ?4 OR version IS NOT ?5 OR content_hash IS NOT ?6 OR modality IS NOT ?7 OR media_type IS NOT ?8)",
params![profile, resource_key(path), metadata.modified_ns, metadata.size_bytes as i64, serde_json::to_string(version).map_err(json_err)?, metadata.content_hash, metadata.modality.to_string(), metadata.media_type],
)?;
let generation = if changed != 0 {
Some(mark_cleanup(&tx, None)?)
} else {
None
};
tx.commit()?;
self.record_write(changed != 0, generation);
Ok(())
}
fn remove_resources(&mut self, profile: &str, paths: &[ResourceId]) -> Result<()> {
let tx = self
.conn
.transaction_with_behavior(rusqlite::TransactionBehavior::Immediate)?;
if !tx.query_row(
"SELECT EXISTS(SELECT 1 FROM profiles WHERE id = ?1)",
[profile],
|r| r.get::<_, bool>(0),
)? {
return Ok(());
}
let mut deleted = false;
let mut removed = 0;
{
let mut stmt =
tx.prepare_cached("DELETE FROM resources WHERE profile = ?1 AND path = ?2")?;
for p in paths {
let key = resource_key(p);
deleted |= delete_fts_rows(&tx, profile, &key)?;
removed += stmt.execute(params![profile, key])?;
}
}
let generation = if deleted || removed != 0 {
Some(mark_cleanup(&tx, deleted.then_some(profile))?)
} else {
None
};
tx.commit()?;
self.record_write(deleted || removed != 0, generation);
Ok(())
}
fn known_inputs(&self, profile: &str, hashes: &[String]) -> Result<HashSet<String>> {
let mut stmt = self
.conn
.prepare_cached("SELECT 1 FROM embeddings WHERE profile = ?1 AND input_hash = ?2")?;
let mut out = HashSet::new();
for h in hashes {
if stmt.exists(params![profile, h])? {
out.insert(h.clone());
}
}
Ok(out)
}
fn put_embeddings(&mut self, profile: &str, items: &[(String, Vec<f32>)]) -> Result<()> {
let tx = self
.conn
.transaction_with_behavior(rusqlite::TransactionBehavior::Immediate)?;
let mut overwritten = false;
let mut changed = false;
{
let mut read = tx.prepare_cached(
"SELECT vector FROM embeddings WHERE profile=?1 AND input_hash=?2",
)?;
let mut stmt = tx.prepare_cached(
"INSERT OR REPLACE INTO embeddings (profile,input_hash,vector) VALUES(?1,?2,?3)",
)?;
for (hash, vector) in items {
let bytes = to_bytes(vector);
let old: Option<Vec<u8>> = read
.query_row(params![profile, hash], |r| r.get(0))
.optional()?;
#[cfg(test)]
if let Some(observer) = &mut self.after_vector_read {
observer();
}
if old.as_ref() == Some(&bytes) {
continue;
}
overwritten |= old.is_some();
changed = true;
stmt.execute(params![profile, hash, bytes])?;
}
}
let generation = if overwritten {
Some(mark_cleanup(&tx, None)?)
} else if changed {
refresh_pending_activity(&tx)?
} else {
None
};
tx.commit()?;
self.record_write(changed, generation);
Ok(())
}
fn supports_image_codes(&self) -> bool {
true
}
fn sync_image_codes(&mut self, profile: &str, paths: &HashSet<ResourceId>) -> Result<()> {
self.backfill_codes(profile, paths)
}
fn image_scan_rows(
&self,
profile: &str,
paths: &HashSet<ResourceId>,
) -> Result<Vec<ScanVector>> {
self.scan_rows(profile, paths)
}
fn selected_image_vectors(
&self,
profile: &str,
hashes: &HashSet<String>,
) -> Result<HashMap<String, Vec<f32>>> {
self.selected_vectors(profile, hashes)
}
fn vectors_for(&self, profile: &str, paths: &HashSet<ResourceId>) -> Result<Vec<StoredVector>> {
let mut stmt = self.conn.prepare_cached(
"SELECT s.path, s.locator, s.heading, s.input_hash, e.vector, s.ordinal
FROM segments s JOIN embeddings e ON e.profile = s.profile AND e.input_hash = s.input_hash
WHERE s.profile = ?1
ORDER BY s.path, s.ordinal",
)?;
let mut rows = stmt.query([profile])?;
let mut out = Vec::new();
while let Some(r) = rows.next()? {
let path = key_resource(r.get(0)?)?;
if !paths.contains(&path) {
continue;
}
let locator: String = r.get(1)?;
out.push(StoredVector {
resource: path,
ordinal: r.get::<_, i64>(5)? as usize,
segment: CachedSegment {
locator: serde_json::from_str(&locator).map_err(json_err)?,
heading: r.get(2)?,
input_hash: r.get(3)?,
},
vector: from_bytes(&r.get::<_, Vec<u8>>(4)?).map_err(|error| {
rusqlite::Error::FromSqlConversionFailure(
4,
rusqlite::types::Type::Blob,
Box::new(error),
)
})?,
});
}
Ok(out)
}
fn lexical_search(
&self,
profile: &str,
query: &LexicalQuery,
paths: &HashSet<ResourceId>,
limit: usize,
) -> Result<Vec<LexicalHit>> {
let Some(expr) = query.fts5() else {
return Ok(Vec::new());
};
if !self.conn.query_row(
"SELECT EXISTS(SELECT 1 FROM profiles WHERE id = ?1)",
[profile],
|r| r.get::<_, bool>(0),
)? {
return Ok(Vec::new());
}
let table = fts_table(profile);
let weights = FIELDS.map(|(_, w)| w.to_string()).join(", ");
let sql = format!(
"SELECT s.path, s.ordinal, -bm25({table}, {weights})
FROM {table} JOIN segments s ON s.rowid = {table}.rowid
WHERE {table} MATCH ?1 AND s.profile = ?2
ORDER BY bm25({table}, {weights}), s.path, s.ordinal"
);
let mut stmt = self.conn.prepare_cached(&sql)?;
let mut rows = stmt.query(params![expr, profile])?;
let mut out = Vec::new();
while out.len() < limit
&& let Some(r) = rows.next()?
{
let path = key_resource(r.get(0)?)?;
if paths.contains(&path) {
out.push(LexicalHit {
resource: path,
ordinal: r.get::<_, i64>(1)? as usize,
score: r.get::<_, f64>(2)? as f32,
});
}
}
Ok(out)
}
fn lexical_matched_terms(
&self,
profile: &str,
query: &LexicalQuery,
keys: &[(ResourceId, usize)],
) -> Result<Vec<Vec<String>>> {
let mut terms = vec![Vec::new(); keys.len()];
if keys.is_empty() || query.groups.is_empty() {
return Ok(terms);
}
let mut positions = HashMap::new();
for (batch, chunk) in keys.chunks(400).enumerate() {
let values = (0..chunk.len())
.map(|i| format!("(?, ?, {})", batch * 400 + i))
.collect::<Vec<_>>()
.join(",");
let sql = format!(
"WITH requested(path, ordinal, position) AS (VALUES {values})
SELECT s.rowid, r.position FROM requested r
CROSS JOIN segments s ON s.path = r.path AND s.ordinal = r.ordinal
WHERE s.profile = ?"
);
let mut params = Vec::with_capacity(chunk.len() * 2 + 1);
for (path, ordinal) in chunk {
params.push(rusqlite::types::Value::Blob(resource_key(path)));
params.push(rusqlite::types::Value::Integer(*ordinal as i64));
}
params.push(rusqlite::types::Value::Text(profile.to_owned()));
let mut stmt = self.conn.prepare(&sql)?;
let mut rows = stmt.query(rusqlite::params_from_iter(params))?;
while let Some(row) = rows.next()? {
positions
.entry(row.get::<_, i64>(0)?)
.or_insert_with(Vec::new)
.push(row.get::<_, i64>(1)? as usize);
}
}
if positions.is_empty() {
return Ok(terms);
}
let rowids = serde_json::to_string(&positions.keys().collect::<Vec<_>>())
.expect("integer rowids serialize to JSON");
let table = fts_table(profile);
let mut stmt = self.conn.prepare_cached(&format!(
"SELECT rowid FROM {table} WHERE {table} MATCH ?1
AND +rowid IN (SELECT value FROM json_each(?2))"
))?;
for group in &query.groups {
let expr = LexicalQuery {
groups: vec![group.clone()],
}
.fts5()
.unwrap();
let mut rows = stmt.query(params![expr, rowids])?;
while let Some(row) = rows.next()? {
for &position in &positions[&row.get::<_, i64>(0)?] {
let matched = &mut terms[position];
for term in group {
if !matched.contains(term) {
matched.push(term.clone());
}
}
}
}
}
Ok(terms)
}
fn prune_orphans(&mut self, profile: &str) -> Result<u64> {
let tx = self
.conn
.transaction_with_behavior(rusqlite::TransactionBehavior::Immediate)?;
let n = tx.execute(
"DELETE FROM embeddings WHERE profile = ?1 AND NOT EXISTS (
SELECT 1 FROM segments s WHERE s.profile = embeddings.profile AND s.input_hash = embeddings.input_hash)",
[profile],
)?;
let generation = if n != 0 {
Some(mark_cleanup(&tx, None)?)
} else {
None
};
tx.commit()?;
self.record_write(n != 0, generation);
Ok(n as u64)
}
fn stats(&self) -> Result<CacheStats> {
let mut stmt = self.conn.prepare(
"SELECT p.id, p.profile,
(SELECT COUNT(*) FROM resources r WHERE r.profile = p.id),
(SELECT COUNT(*) FROM segments s WHERE s.profile = p.id),
(SELECT COUNT(*) FROM embeddings e WHERE e.profile = p.id),
(SELECT COUNT(*) FROM image_codes c WHERE c.profile = p.id)
FROM profiles p ORDER BY p.id",
)?;
let profiles = stmt
.query_map([], |r| {
Ok(ProfileStats {
profile_id: r.get(0)?,
profile: serde_json::from_str(&r.get::<_, String>(1)?).ok(),
resources: r.get::<_, i64>(2)? as u64,
segments: r.get::<_, i64>(3)? as u64,
embeddings: r.get::<_, i64>(4)? as u64,
image_codes: r.get::<_, i64>(5)? as u64,
})
})?
.collect::<rusqlite::Result<Vec<_>>>()?;
let size_bytes = ["", "-wal", "-shm"]
.iter()
.filter_map(|s| std::fs::metadata(format!("{}{s}", self.path.display())).ok())
.map(|m| m.len())
.sum();
Ok(CacheStats {
location: Some(self.path.clone()),
size_bytes: Some(size_bytes),
backup_excluded: self.guards.backup_excluded,
in_synced_folder: self.guards.in_synced_folder,
scrub_pending: self.cleanup_pending()?,
profiles,
})
}
fn clear(&mut self) -> Result<()> {
self.conn.flush_prepared_statement_cache();
let tx = self
.conn
.transaction_with_behavior(rusqlite::TransactionBehavior::Immediate)?;
let tables = tx.prepare("SELECT name FROM sqlite_master WHERE name GLOB 'segments_fts_p_*' AND sql LIKE 'CREATE VIRTUAL TABLE%'")?
.query_map([], |r| r.get::<_, String>(0))?
.collect::<rusqlite::Result<Vec<_>>>()?;
for table in tables {
tx.execute_batch(&format!("DROP TABLE \"{}\"", table.replace('"', "\"\"")))?;
}
tx.execute_batch("DELETE FROM embeddings; DELETE FROM segments; DELETE FROM resources; DELETE FROM profiles;")?;
tx.execute("DELETE FROM meta WHERE key GLOB 'fts_scrub_pending:*'", [])?;
tx.execute_batch("INSERT OR IGNORE INTO meta VALUES('compact_pending','1');")?;
mark_cleanup(&tx, None)?;
tx.commit()?;
if self.run_scrub()?.pending {
return Err(Error::CacheCleanupPending);
}
let busy: i64 = self
.conn
.query_row("PRAGMA wal_checkpoint(TRUNCATE)", [], |r| r.get(0))?;
if busy != 0 {
let tx = self
.conn
.transaction_with_behavior(rusqlite::TransactionBehavior::Immediate)?;
mark_cleanup(&tx, None)?;
tx.commit()?;
return Err(Error::CacheCleanupPending);
}
Ok(())
}
}
#[cfg(all(test, windows))]
mod windows_tests {
use super::*;
#[test]
fn file_keys_keep_the_legacy_canonical_drive_prefix() {
let id = ResourceId::from_path(r"C:\notes\chapter.md").unwrap();
let legacy = br"\\?\C:\notes\chapter.md".to_vec();
assert_eq!(resource_key(&id), legacy);
assert_eq!(key_resource(legacy).unwrap(), id);
}
}
#[cfg(test)]
mod guard_tests {
use super::*;
use std::cell::{Cell, RefCell};
#[test]
fn stats_reuses_one_verification_even_after_directory_changes() {
let temp = tempfile::tempdir().unwrap();
let calls = Cell::new(0);
let cache =
SqliteCache::open_with_guards(temp.path().join("cache.sqlite3"), |dir, owned| {
assert!(!owned);
super::super::cache_guards::prepare_with(
dir,
owned,
|_, _| panic!("foreign directory publication"),
|_| {
calls.set(calls.get() + 1);
Some(false)
},
|_| {},
)
})
.unwrap();
std::fs::create_dir(temp.path().join(".stfolder")).unwrap();
for _ in 0..4 {
let stats = cache.stats().unwrap();
assert_eq!(stats.backup_excluded, Some(false));
assert!(!stats.in_synced_folder);
}
assert_eq!(calls.get(), 1);
assert!(!temp.path().join("CACHEDIR.TAG").exists());
}
#[test]
fn unsupported_tag_publication_warns_and_cache_still_opens() {
let temp = tempfile::tempdir().unwrap();
let warnings = RefCell::new(Vec::new());
let path = temp.path().join("dedicated/cache.sqlite3");
let cache = SqliteCache::open_with_guards(path, |dir, owned| {
assert!(owned);
super::super::cache_guards::prepare_with(
dir,
owned,
|_, _| {
Err(std::io::Error::new(
std::io::ErrorKind::Unsupported,
"hard links unsupported",
))
},
|_| None,
|warning| warnings.borrow_mut().push(warning),
)
})
.unwrap();
assert!(cache.stats().is_ok());
assert!(
warnings
.borrow()
.iter()
.any(|w| w.contains("hard links unsupported"))
);
assert!(!cache.path().parent().unwrap().join("CACHEDIR.TAG").exists());
assert!(
std::fs::read_dir(cache.path().parent().unwrap())
.unwrap()
.flatten()
.all(|e| !e
.file_name()
.to_string_lossy()
.starts_with(".CACHEDIR.TAG-"))
);
}
}
#[cfg(test)]
mod scrub_tests {
use super::*;
use std::time::{Duration, Instant};
#[test]
fn sync_serializes_external_writer_between_vector_read_and_insert() {
use std::sync::{
Arc, Mutex,
atomic::{AtomicUsize, Ordering},
};
struct SyntheticEmbedder(crate::ModelIdentity);
impl crate::Embedder for SyntheticEmbedder {
fn identity(&self) -> &crate::ModelIdentity {
&self.0
}
fn supports(&self, modality: Modality) -> bool {
modality == Modality::Text
}
fn embed_batch(
&self,
requests: &[seyk_embed::EmbeddingRequest<'_>],
) -> seyk_embed::Result<Vec<crate::Embedding>> {
requests
.iter()
.map(|_| crate::Embedding::from_raw(vec![1.0; 128], crate::Dimension::D128))
.collect()
}
}
let temp = tempfile::tempdir().unwrap();
let path = temp.path().join("cache.sqlite3");
let file = temp.path().join("synthetic.txt");
std::fs::write(&file, "synthetic alpha garden words").unwrap();
let mut cache = SqliteCache::open(&path).unwrap();
let writer = Arc::new(Mutex::new(Connection::open(&path).unwrap()));
writer.lock().unwrap().busy_timeout(Duration::ZERO).unwrap();
let tx = cache.conn.transaction().unwrap();
tx.query_row("SELECT count(*) FROM embeddings", [], |r| {
r.get::<_, i64>(0)
})
.unwrap();
writer
.lock()
.unwrap()
.execute(
"INSERT INTO meta(key,value) VALUES('interleaving_control','committed')",
[],
)
.unwrap();
let error = tx
.execute(
"INSERT INTO meta(key,value) VALUES('stale_snapshot','write')",
[],
)
.unwrap_err();
match error {
rusqlite::Error::SqliteFailure(code, _) => {
assert_eq!(code.extended_code, rusqlite::ffi::SQLITE_BUSY_SNAPSHOT)
}
other => panic!("expected BUSY_SNAPSHOT, got {other}"),
}
tx.rollback().unwrap();
let attempts = Arc::new(AtomicUsize::new(0));
let observed = attempts.clone();
let external = writer.clone();
cache.after_vector_read = Some(Box::new(move || {
let error = external
.lock()
.unwrap()
.execute(
"INSERT INTO meta(key,value) VALUES('external_writer','committed')",
[],
)
.unwrap_err();
assert!(sqlite_busy(&error));
observed.fetch_add(1, Ordering::SeqCst);
}));
let engine = crate::SearchEngine::builder()
.embedder(SyntheticEmbedder(crate::ModelIdentity {
model_id: "test/interleaving".into(),
revision: "1".into(),
variant: "synthetic".into(),
dimension: crate::Dimension::D128,
normalization: "l2".into(),
query_prompt: String::new(),
document_prompt: String::new(),
max_tokens: 512,
}))
.cache(cache)
.build()
.unwrap();
let result = engine.index().paths([&file]).execute().unwrap();
assert_eq!(result.sync.resources_updated, 1);
assert!(attempts.load(Ordering::SeqCst) > 0);
writer
.lock()
.unwrap()
.execute(
"INSERT INTO meta(key,value) VALUES('external_writer','committed')",
[],
)
.unwrap();
assert_eq!(
engine.with_cache(|c| c.stats()).unwrap().profiles[0].embeddings,
1
);
}
#[test]
fn scrub_scheduler_uses_quiet_and_retry_deadlines_without_busy_polling() {
let temp = tempfile::tempdir().unwrap();
let mut cache = SqliteCache::open(temp.path().join("cache.sqlite3")).unwrap();
let quiet = Duration::from_secs(60);
cache.last_change = Instant::now() - Duration::from_secs(40);
let delay = cache.scrub_delay(quiet);
assert!(delay > Duration::from_secs(19) && delay <= Duration::from_secs(20));
cache.last_attempt = Some(Instant::now() - Duration::from_secs(5));
let retry = cache.scrub_delay(quiet);
assert!(retry > Duration::from_secs(54) && retry <= Duration::from_secs(55));
cache.last_change = Instant::now() - Duration::from_secs(61);
cache.last_attempt = None;
assert_eq!(
cache.scrub_delay(quiet),
quiet,
"idle worker checks once per minute"
);
assert_eq!(MemoryCache::new().scrub_delay(quiet), quiet);
}
#[test]
fn local_writes_keep_their_deadline_but_external_writes_restart_it() {
let temp = tempfile::tempdir().unwrap();
let path = temp.path().join("cache.sqlite3");
let mut cache = SqliteCache::open(&path).unwrap();
let profile = CacheProfile {
model: ModelIdentity {
model_id: "test/deadline".into(),
revision: "1".into(),
variant: "synthetic".into(),
dimension: crate::Dimension::D128,
normalization: "l2".into(),
query_prompt: String::new(),
document_prompt: String::new(),
max_tokens: 512,
},
segmenter_version: "test".into(),
};
let pid = profile.id();
cache.register_profile(&pid, &profile).unwrap();
cache
.put_embeddings(&pid, &[("input".into(), vec![1.0; 128])])
.unwrap();
cache
.put_embeddings(&pid, &[("input".into(), vec![2.0; 128])])
.unwrap();
assert!(cache.observed_generation.is_some());
let quiet = Duration::from_secs(60);
cache.last_change = Instant::now() - Duration::from_secs(61);
assert!(
!cache.scrub_if_idle(quiet).unwrap().unwrap().pending,
"own marker must not add a second quiet period at the first wake"
);
cache
.put_embeddings(&pid, &[("input".into(), vec![3.0; 128])])
.unwrap();
cache
.put_embeddings(&pid, &[("new-input".into(), vec![4.0; 128])])
.unwrap();
let mut writer = Connection::open(&path).unwrap();
let tx = writer
.transaction_with_behavior(rusqlite::TransactionBehavior::Immediate)
.unwrap();
refresh_pending_activity(&tx).unwrap();
tx.commit().unwrap();
cache.last_change = Instant::now() - Duration::from_secs(61);
cache.last_attempt = None;
assert!(cache.scrub_if_idle(quiet).unwrap().is_none());
assert!(cache.last_change.elapsed() < Duration::from_secs(1));
let reopened = SqliteCache::open(&path).unwrap();
assert_eq!(
reopened.observed_generation, cache.observed_generation,
"open observes pending generation before starting its quiet period"
);
}
#[test]
fn idle_scrub_observes_other_writers_and_throttles_busy_retries() {
let temp = tempfile::tempdir().unwrap();
let path = temp.path().join("cache.sqlite3");
let mut cache = SqliteCache::open(&path).unwrap();
let tx = cache.conn.transaction().unwrap();
mark_cleanup(&tx, None).unwrap();
tx.commit().unwrap();
let quiet = Duration::from_secs(60);
assert!(cache.scrub_if_idle(quiet).unwrap().is_none());
cache.last_change = Instant::now() - Duration::from_secs(61);
let mut writer = Connection::open(&path).unwrap();
let tx = writer.transaction().unwrap();
refresh_pending_activity(&tx).unwrap();
tx.commit().unwrap();
assert!(cache.scrub_if_idle(quiet).unwrap().is_none());
assert!(cache.last_change.elapsed() < quiet);
let reader = Connection::open(&path).unwrap();
reader.execute_batch("BEGIN; SELECT * FROM meta;").unwrap();
cache.last_change = Instant::now() - Duration::from_secs(61);
assert!(cache.scrub_if_idle(quiet).unwrap().unwrap().pending);
assert!(
cache.scrub_if_idle(quiet).unwrap().is_none(),
"do not optimize again on every tick with a busy reader"
);
assert!(cache.cleanup_pending().unwrap());
reader.execute_batch("ROLLBACK").unwrap();
cache.last_attempt = Some(Instant::now() - Duration::from_secs(61));
assert!(!cache.scrub_if_idle(quiet).unwrap().unwrap().pending);
assert!(!cache.cleanup_pending().unwrap());
}
}