use std::sync::{Mutex, Once};
use anyhow::{Context, Result, bail};
use chrono::{DateTime, Utc};
use rusqlite::{Connection, OptionalExtension, params};
use uuid::Uuid;
use crate::entities::attachment::{AttachmentChunk, AttachmentHit};
use crate::entities::note::Note;
use crate::entities::rag::{RagDocument, RagHit, RagSourceInfo, RagStoredSource};
use crate::entities::self_model::SelfModel;
use crate::shared::storage::schema::DB_SCHEMA;
static REGISTER_VEC: Once = Once::new();
fn register_sqlite_vec() {
#[allow(clippy::missing_transmute_annotations)]
REGISTER_VEC.call_once(|| unsafe {
rusqlite::ffi::sqlite3_auto_extension(Some(std::mem::transmute(
sqlite_vec::sqlite3_vec_init as *const (),
)));
});
}
pub struct Db {
conn: Mutex<Connection>,
}
impl Db {
pub fn open(path: &std::path::Path) -> Result<Self> {
register_sqlite_vec();
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent).ok();
}
let conn = Connection::open(path).with_context(|| format!("opening {}", path.display()))?;
Self::from_conn(conn)
}
#[cfg(test)]
pub fn open_in_memory() -> Result<Self> {
register_sqlite_vec();
Self::from_conn(Connection::open_in_memory()?)
}
#[cfg(test)]
pub(crate) fn batch<T>(&self, f: impl FnOnce() -> T) -> T {
self.conn.lock().unwrap().execute_batch("BEGIN").unwrap();
let out = f();
self.conn.lock().unwrap().execute_batch("COMMIT").unwrap();
out
}
fn from_conn(conn: Connection) -> Result<Self> {
let mut conn = conn;
migrate(&mut conn)?;
Ok(Self {
conn: Mutex::new(conn),
})
}
}
fn migrate(conn: &mut Connection) -> Result<()> {
baseline_ddl(conn)?;
let from = read_user_version(conn)?;
if from > DB_SCHEMA {
bail!("data.db is from a newer app version (schema {from}, supported is {DB_SCHEMA})");
}
if from == 0 {
set_user_version(conn, DB_SCHEMA)?;
}
apply_db_steps(conn, DB_STEPS, from)?;
Ok(())
}
#[allow(dead_code)] struct DbStep {
to: u32,
summary: &'static str,
apply: fn(&Connection) -> Result<()>,
}
const DB_STEPS: &[DbStep] = &[];
fn apply_db_steps(conn: &mut Connection, steps: &[DbStep], from: u32) -> Result<()> {
for step in steps.iter().filter(|s| s.to > from.max(1)) {
let tx = conn.transaction()?;
(step.apply)(&tx).with_context(|| format!("data.db migration → v{}", step.to))?;
set_user_version(&tx, step.to)?;
tx.commit()?;
}
Ok(())
}
fn read_user_version(conn: &Connection) -> Result<u32> {
Ok(conn.pragma_query_value(None, "user_version", |r| r.get::<_, i64>(0))? as u32)
}
fn set_user_version(conn: &Connection, v: u32) -> Result<()> {
conn.execute_batch(&format!("PRAGMA user_version = {v};"))?;
Ok(())
}
pub fn peek_user_version(path: &std::path::Path) -> Result<u32> {
if !path.exists() {
return Ok(0);
}
let conn = Connection::open(path).with_context(|| format!("opening {}", path.display()))?;
read_user_version(&conn)
}
pub fn needs_step_migration(user_version: u32) -> bool {
DB_STEPS.iter().any(|s| s.to > user_version.max(1))
}
pub fn vacuum_into(src: &std::path::Path, dest: &std::path::Path) -> Result<()> {
use rusqlite::OpenFlags;
register_sqlite_vec();
if !is_sqlite_file(src) {
bail!("{} is not a SQLite database", src.display());
}
let dest_str = dest
.to_str()
.with_context(|| format!("non-UTF-8 destination path {}", dest.display()))?;
if dest.exists() {
std::fs::remove_file(dest)
.with_context(|| format!("removing a stale {}", dest.display()))?;
}
let conn = Connection::open_with_flags(
src,
OpenFlags::SQLITE_OPEN_READ_ONLY | OpenFlags::SQLITE_OPEN_NO_MUTEX,
)
.with_context(|| format!("opening {} read-only", src.display()))?;
conn.execute("VACUUM INTO ?1", [dest_str])
.with_context(|| format!("compacting {} into {}", src.display(), dest.display()))?;
Ok(())
}
pub fn vacuum(path: &std::path::Path) -> Result<()> {
register_sqlite_vec();
if !is_sqlite_file(path) {
bail!("{} is not a SQLite database", path.display());
}
let conn = Connection::open(path).with_context(|| format!("opening {}", path.display()))?;
conn.execute_batch("VACUUM")
.with_context(|| format!("compacting {}", path.display()))?;
Ok(())
}
fn is_sqlite_file(path: &std::path::Path) -> bool {
use std::io::Read;
const MAGIC: &[u8; 16] = b"SQLite format 3\0";
let Ok(mut f) = std::fs::File::open(path) else {
return false;
};
let mut head = [0u8; 16];
f.read_exact(&mut head).is_ok() && &head == MAGIC
}
fn baseline_ddl(conn: &Connection) -> Result<()> {
conn.execute_batch(
"CREATE TABLE IF NOT EXISTS meta (key TEXT PRIMARY KEY, value TEXT NOT NULL);
CREATE TABLE IF NOT EXISTS notes (
id TEXT PRIMARY KEY,
profile_id TEXT NOT NULL,
content TEXT NOT NULL,
tags TEXT NOT NULL,
created_at TEXT NOT NULL,
updated_at TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_notes_profile ON notes(profile_id);
CREATE TABLE IF NOT EXISTS note_vectors (
note_id TEXT PRIMARY KEY,
profile_id TEXT NOT NULL,
embedding TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_note_vectors_profile ON note_vectors(profile_id);
CREATE TABLE IF NOT EXISTS rag_documents (
rowid INTEGER PRIMARY KEY,
id TEXT NOT NULL UNIQUE,
profile_id TEXT NOT NULL,
source TEXT NOT NULL,
chunk_text TEXT NOT NULL,
created_at TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_rag_profile ON rag_documents(profile_id);
CREATE TABLE IF NOT EXISTS rag_sources (
profile_id TEXT NOT NULL,
source TEXT NOT NULL,
content TEXT NOT NULL,
created_at TEXT NOT NULL,
PRIMARY KEY (profile_id, source)
);
CREATE TABLE IF NOT EXISTS self_models (
profile_id TEXT PRIMARY KEY,
data TEXT NOT NULL,
version INTEGER NOT NULL,
updated_at TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS note_links (
profile_id TEXT NOT NULL,
from_id TEXT NOT NULL,
to_id TEXT NOT NULL,
relation TEXT NOT NULL,
created_at TEXT NOT NULL,
PRIMARY KEY (profile_id, from_id, to_id, relation)
);
CREATE INDEX IF NOT EXISTS idx_note_links_from ON note_links(profile_id, from_id);
CREATE INDEX IF NOT EXISTS idx_note_links_to ON note_links(profile_id, to_id);
CREATE TABLE IF NOT EXISTS note_superseded (
note_id TEXT PRIMARY KEY,
profile_id TEXT NOT NULL,
superseded_by TEXT NOT NULL,
superseded_at TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS note_rag_links (
profile_id TEXT NOT NULL,
note_id TEXT NOT NULL,
source TEXT NOT NULL,
created_at TEXT NOT NULL,
PRIMARY KEY (profile_id, note_id, source)
);
CREATE INDEX IF NOT EXISTS idx_note_rag_note ON note_rag_links(profile_id, note_id);
CREATE INDEX IF NOT EXISTS idx_note_rag_source ON note_rag_links(profile_id, source);
CREATE TABLE IF NOT EXISTS llm_history (
rowid INTEGER PRIMARY KEY,
profile_id TEXT NOT NULL,
changed_at TEXT NOT NULL,
model TEXT NOT NULL,
mode TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_llm_history_profile ON llm_history(profile_id);
CREATE TABLE IF NOT EXISTS attachment_documents (
rowid INTEGER PRIMARY KEY,
id TEXT NOT NULL UNIQUE,
chat_id TEXT NOT NULL,
attachment_id TEXT NOT NULL,
name TEXT NOT NULL,
chunk_text TEXT NOT NULL,
created_at TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_attachment_docs_chat ON attachment_documents(chat_id);
CREATE INDEX IF NOT EXISTS idx_attachment_docs_att
ON attachment_documents(chat_id, attachment_id);",
)?;
for table in ["note_vectors", "rag_documents", "attachment_documents"] {
add_column_if_missing(conn, table, "embed_gen", "INTEGER")?;
}
Ok(())
}
fn column_exists(conn: &Connection, table: &str, column: &str) -> Result<bool> {
let mut stmt = conn.prepare(&format!("PRAGMA table_info({table})"))?;
let mut rows = stmt.query([])?;
while let Some(row) = rows.next()? {
if row.get::<_, String>(1)? == column {
return Ok(true);
}
}
Ok(false)
}
fn add_column_if_missing(conn: &Connection, table: &str, column: &str, decl: &str) -> Result<()> {
if !column_exists(conn, table, column)? {
conn.execute_batch(&format!("ALTER TABLE {table} ADD COLUMN {column} {decl}"))?;
}
Ok(())
}
fn table_exists(conn: &Connection, name: &str) -> Result<bool> {
let found: Option<i64> = conn
.query_row(
"SELECT 1 FROM sqlite_master WHERE type IN ('table','view') AND name = ?1",
[name],
|r| r.get(0),
)
.optional()?;
Ok(found.is_some())
}
const KEY_RAG_DIM: &str = "rag_dim";
const KEY_EMBED_GEN: &str = "embed_gen";
const KEY_EMBED_CANARY: &str = "embed_canary";
const KEY_EMBED_MODEL_ID: &str = "embed_model_id";
const KEY_EMBED_CONVENTION: &str = "embed_convention";
const KEY_EMBED_CAL_UNRELATED: &str = "embed_cal_unrelated";
const KEY_EMBED_CAL_PARAPHRASE: &str = "embed_cal_paraphrase";
const KEY_RAG_STALE_PROFILES: &str = "rag_stale_profiles";
fn meta_get(conn: &Connection, key: &str) -> Result<Option<String>> {
let value: Option<String> = conn
.query_row("SELECT value FROM meta WHERE key = ?1", [key], |r| r.get(0))
.optional()?;
Ok(value)
}
fn meta_set(conn: &Connection, key: &str, value: &str) -> Result<()> {
conn.execute(
"INSERT INTO meta(key, value) VALUES (?1, ?2)
ON CONFLICT(key) DO UPDATE SET value = excluded.value",
params![key, value],
)?;
Ok(())
}
fn meta_del(conn: &Connection, key: &str) -> Result<()> {
conn.execute("DELETE FROM meta WHERE key = ?1", [key])?;
Ok(())
}
fn current_embed_gen(conn: &Connection) -> Result<u32> {
Ok(meta_get(conn, KEY_EMBED_GEN)?
.and_then(|v| v.parse().ok())
.unwrap_or(FIRST_EMBED_GEN))
}
const FIRST_EMBED_GEN: u32 = 1;
const NULL_EMBED_GEN: i64 = -1;
fn vec_dim(conn: &Connection) -> Result<Option<usize>> {
Ok(meta_get(conn, KEY_RAG_DIM)?.map(|d| d.parse().unwrap_or(0)))
}
fn ensure_dim(conn: &Connection, dim: usize) -> Result<()> {
if dim == 0 {
bail!("refusing to index an empty embedding");
}
match vec_dim(conn)? {
Some(existing) if existing == dim => Ok(()),
Some(existing) => bail!("embedding dim mismatch: table is {existing}, got {dim}"),
None => meta_set(conn, KEY_RAG_DIM, &dim.to_string()),
}
}
fn ensure_vec_table(conn: &Connection, dim: usize) -> Result<()> {
ensure_dim(conn, dim)?;
conn.execute(
&format!(
"CREATE VIRTUAL TABLE IF NOT EXISTS rag_vectors USING vec0(
profile_id TEXT partition key,
embedding float[{dim}]
)"
),
[],
)?;
Ok(())
}
fn ensure_attachment_vec_table(conn: &Connection, dim: usize) -> Result<()> {
ensure_dim(conn, dim)?;
conn.execute(
&format!(
"CREATE VIRTUAL TABLE IF NOT EXISTS attachment_vectors USING vec0(
chat_id TEXT partition key,
embedding float[{dim}]
)"
),
[],
)?;
Ok(())
}
fn row_to_note(r: &rusqlite::Row) -> rusqlite::Result<Note> {
Ok(Note {
id: parse_uuid(r.get::<_, String>(0)?),
profile_id: parse_uuid(r.get::<_, String>(1)?),
content: r.get(2)?,
tags: serde_json::from_str(&r.get::<_, String>(3)?).unwrap_or_default(),
created_at: parse_dt(r.get::<_, String>(4)?),
updated_at: parse_dt(r.get::<_, String>(5)?),
})
}
fn cosine(a: &[f32], b: &[f32]) -> f32 {
if a.len() != b.len() || a.is_empty() {
return 0.0;
}
let dot: f32 = a.iter().zip(b).map(|(x, y)| x * y).sum();
let na: f32 = a.iter().map(|x| x * x).sum();
let nb: f32 = b.iter().map(|x| x * x).sum();
if na == 0.0 || nb == 0.0 {
return 0.0;
}
dot / (na.sqrt() * nb.sqrt())
}
fn parse_uuid(s: String) -> Uuid {
Uuid::parse_str(&s).unwrap_or(Uuid::nil())
}
fn parse_dt(s: String) -> DateTime<Utc> {
DateTime::parse_from_rfc3339(&s)
.map(|d| d.with_timezone(&Utc))
.unwrap_or_else(|_| Utc::now())
}
mod attachments;
mod embed_gen;
mod graph;
mod llm_history;
mod notes;
mod rag;
mod self_model;
mod stats;
pub use embed_gen::{ReembedPending, ReembedRow};
pub use stats::{DbStats, KeyedRow, stats_of_file, stats_of_image};
#[cfg(test)]
mod migrate_tests {
use super::*;
fn table_exists(conn: &Connection, name: &str) -> bool {
super::table_exists(conn, name).unwrap()
}
#[test]
fn baseline_stamps_fresh_db_to_v1() {
let db = Db::open_in_memory().unwrap();
let conn = db.conn.lock().unwrap();
assert_eq!(read_user_version(&conn).unwrap(), DB_SCHEMA);
assert!(table_exists(&conn, "notes"));
}
#[test]
fn migrate_is_idempotent() {
let mut conn = Connection::open_in_memory().unwrap();
migrate(&mut conn).unwrap();
migrate(&mut conn).unwrap();
assert_eq!(read_user_version(&conn).unwrap(), DB_SCHEMA);
}
#[test]
fn migrate_refuses_downgrade() {
let mut conn = Connection::open_in_memory().unwrap();
set_user_version(&conn, DB_SCHEMA + 5).unwrap();
assert!(migrate(&mut conn).is_err());
}
#[test]
fn peek_user_version_zero_for_missing_file() {
let dir = tempfile::tempdir().unwrap();
assert_eq!(peek_user_version(&dir.path().join("nope.db")).unwrap(), 0);
}
#[test]
fn no_pending_step_migration_at_v1() {
assert!(!needs_step_migration(0));
assert!(!needs_step_migration(1));
}
#[test]
fn apply_db_steps_commits_and_rolls_back_transactionally() {
fn good(c: &Connection) -> Result<()> {
c.execute_batch("CREATE TABLE t_ok(x)")?;
Ok(())
}
fn bad(c: &Connection) -> Result<()> {
c.execute_batch("CREATE TABLE t_bad(x)")?;
bail!("deliberate step failure");
}
let mut conn = Connection::open_in_memory().unwrap();
set_user_version(&conn, 1).unwrap();
apply_db_steps(
&mut conn,
&[DbStep {
to: 2,
summary: "ok",
apply: good,
}],
1,
)
.unwrap();
assert_eq!(read_user_version(&conn).unwrap(), 2);
assert!(table_exists(&conn, "t_ok"));
let res = apply_db_steps(
&mut conn,
&[DbStep {
to: 3,
summary: "bad",
apply: bad,
}],
2,
);
assert!(res.is_err());
assert_eq!(read_user_version(&conn).unwrap(), 2);
assert!(!table_exists(&conn, "t_bad"));
}
}
#[cfg(test)]
mod compact_tests {
use super::*;
use crate::entities::rag::RagDocument;
fn fragmented_db(path: &std::path::Path) -> Uuid {
let profile = Uuid::new_v4();
let db = Db::open(path).unwrap();
db.batch(|| {
for i in 0..200 {
let text = format!("scratch {i} {}", "x".repeat(500));
db.rag_insert(&RagDocument::new(profile, "scratch", text, vec![0.0, 1.0]))
.unwrap();
}
db.rag_insert(&RagDocument::new(
profile,
"keep",
"the kept chunk",
vec![1.0, 0.0],
))
.unwrap();
db.rag_insert(&RagDocument::new(
profile,
"keep",
"another kept",
vec![0.9, 0.1],
))
.unwrap();
});
db.batch(|| db.rag_delete_by_source(profile, "scratch").unwrap());
profile
}
fn freelist(path: &std::path::Path) -> i64 {
let conn = Connection::open(path).unwrap();
conn.pragma_query_value(None, "freelist_count", |r| r.get(0))
.unwrap()
}
fn len(path: &std::path::Path) -> u64 {
std::fs::metadata(path).unwrap().len()
}
#[test]
fn vacuum_into_preserves_the_vector_index() {
let dir = tempfile::tempdir().unwrap();
let src = dir.path().join("data.db");
let profile = fragmented_db(&src);
let dest = dir.path().join("compact.db");
vacuum_into(&src, &dest).unwrap();
assert!(freelist(&src) > 0, "the source keeps its free pages");
assert_eq!(freelist(&dest), 0, "the copy has none");
assert!(len(&dest) < len(&src), "and is therefore smaller");
let copy = Db::open(&dest).unwrap();
let hits = copy.rag_search(profile, &[1.0, 0.0], 5).unwrap();
assert_eq!(hits.len(), 2, "both surviving chunks are still indexed");
assert_eq!(hits[0].chunk_text, "the kept chunk");
assert_eq!(hits[0].source, "keep");
assert_eq!(hits[1].chunk_text, "another kept");
}
#[test]
fn vacuum_into_preserves_the_schema_version() {
let dir = tempfile::tempdir().unwrap();
let src = dir.path().join("data.db");
Db::open(&src).unwrap();
let dest = dir.path().join("compact.db");
vacuum_into(&src, &dest).unwrap();
assert_eq!(peek_user_version(&dest).unwrap(), DB_SCHEMA);
}
#[test]
fn vacuum_into_overwrites_a_stale_destination() {
let dir = tempfile::tempdir().unwrap();
let src = dir.path().join("data.db");
Db::open(&src).unwrap();
let dest = dir.path().join("compact.db");
std::fs::write(&dest, b"leftovers").unwrap();
vacuum_into(&src, &dest).unwrap();
assert_eq!(peek_user_version(&dest).unwrap(), DB_SCHEMA);
}
#[test]
fn vacuum_reclaims_free_pages_in_place() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("data.db");
let profile = fragmented_db(&path);
let before = len(&path);
assert!(freelist(&path) > 0);
vacuum(&path).unwrap();
assert_eq!(freelist(&path), 0);
assert!(len(&path) < before);
let db = Db::open(&path).unwrap();
assert_eq!(db.rag_search(profile, &[1.0, 0.0], 5).unwrap().len(), 2);
}
#[test]
fn vacuum_into_does_not_touch_the_source_directory() {
let dir = tempfile::tempdir().unwrap();
let real = dir.path().join("real.db");
fragmented_db(&real);
let valid = std::fs::read(&real).unwrap();
for (label, content) in [
("placeholder", b"SQLITE".to_vec()),
("empty", Vec::new()),
("a real database", valid),
] {
let case = tempfile::tempdir().unwrap();
let src = case.path().join("data.db");
let wal = case.path().join("data.db-wal");
std::fs::write(&src, &content).unwrap();
std::fs::write(&wal, b"stale wal").unwrap();
let _ = vacuum_into(&src, &case.path().join("out.db"));
assert!(src.is_file(), "{label}: the source must survive");
assert!(wal.is_file(), "{label}: and so must its sidecar");
assert_eq!(std::fs::read(&wal).unwrap(), b"stale wal", "{label}");
assert_eq!(std::fs::read(&src).unwrap(), content, "{label}");
}
}
#[test]
fn compacting_a_non_database_fails() {
let dir = tempfile::tempdir().unwrap();
let junk = dir.path().join("data.db");
std::fs::write(&junk, b"definitely not a database").unwrap();
assert!(vacuum_into(&junk, &dir.path().join("out.db")).is_err());
assert!(vacuum(&junk).is_err());
}
}