use crate::migrations::run_migrations;
use crate::models::{
BranchRow, CodeChunkRow, MemoryRow, MessageRow, SessionChunkRow, SessionRow, SessionTextHit,
UsageTotals,
};
use rusqlite::Connection;
use whycodes_core::types::Usage;
pub struct Database {
conn: Connection,
}
pub type MessageInsert = (
String,
String,
String,
Option<String>,
Option<String>,
String,
);
pub type BranchInsert = (Option<String>, Vec<String>);
const BUSY_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5);
impl Database {
pub fn open(path: &str) -> crate::error::Result<Self> {
let conn = Connection::open(path)?;
conn.busy_timeout(BUSY_TIMEOUT)?;
let _ = conn.execute_batch("PRAGMA journal_mode=WAL;");
conn.execute_batch("PRAGMA foreign_keys=ON;")?;
run_migrations(&conn)?;
Ok(Self { conn })
}
#[cfg(any(test, feature = "test-util"))]
pub fn drop_sessions_table_for_test(&self) -> crate::error::Result<()> {
self.conn.execute_batch("DROP TABLE sessions;")?;
Ok(())
}
pub fn open_existing_readonly(path: &str) -> crate::error::Result<Option<Self>> {
if !std::path::Path::new(path).is_file() {
return Ok(None);
}
let flags =
rusqlite::OpenFlags::SQLITE_OPEN_READ_ONLY | rusqlite::OpenFlags::SQLITE_OPEN_NO_MUTEX;
let conn = Connection::open_with_flags(path, flags)?;
conn.busy_timeout(BUSY_TIMEOUT)?;
Ok(Some(Self { conn }))
}
pub fn open_in_memory() -> crate::error::Result<Self> {
let conn = Connection::open_in_memory()?;
conn.execute_batch("PRAGMA foreign_keys=ON;")?;
run_migrations(&conn)?;
Ok(Self { conn })
}
pub fn upsert_session(
&self,
id: &str,
title: &str,
project_path: &str,
created_at: &str,
updated_at: &str,
usage: &Usage,
) -> crate::error::Result<()> {
self.conn.execute(
"INSERT INTO sessions (
id, title, created_at, updated_at, project_path,
input_tokens, output_tokens,
cache_creation_input_tokens, cache_read_input_tokens
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9)
ON CONFLICT(id) DO UPDATE SET
title = excluded.title,
updated_at = excluded.updated_at,
project_path = excluded.project_path,
input_tokens = excluded.input_tokens,
output_tokens = excluded.output_tokens,
cache_creation_input_tokens = excluded.cache_creation_input_tokens,
cache_read_input_tokens = excluded.cache_read_input_tokens",
rusqlite::params![
id,
title,
created_at,
updated_at,
project_path,
usage.input_tokens as i64,
usage.output_tokens as i64,
usage.cache_creation_input_tokens.map(|v| v as i64),
usage.cache_read_input_tokens.map(|v| v as i64),
],
)?;
Ok(())
}
pub fn create_session(
&self,
id: &str,
title: &str,
project_path: &str,
) -> crate::error::Result<()> {
let now = chrono::Utc::now().to_rfc3339();
self.upsert_session(id, title, project_path, &now, &now, &Usage::default())
}
pub fn get_session(&self, id: &str) -> crate::error::Result<Option<SessionRow>> {
let mut stmt = self.conn.prepare(
"SELECT id, title, created_at, updated_at, project_path,
input_tokens, output_tokens,
cache_creation_input_tokens, cache_read_input_tokens
FROM sessions WHERE id = ?1",
)?;
let mut rows = stmt.query_map(rusqlite::params![id], map_session_row)?;
Ok(rows.next().transpose()?)
}
pub fn list_sessions(&self) -> crate::error::Result<Vec<SessionRow>> {
let mut stmt = self.conn.prepare(
"SELECT id, title, created_at, updated_at, project_path,
input_tokens, output_tokens,
cache_creation_input_tokens, cache_read_input_tokens
FROM sessions ORDER BY updated_at DESC",
)?;
let rows = stmt.query_map([], map_session_row)?;
let mut result = Vec::new();
for row in rows {
result.push(row?);
}
Ok(result)
}
pub fn search_sessions_text(
&self,
terms: &[String],
limit: usize,
) -> crate::error::Result<Vec<SessionTextHit>> {
if terms.is_empty() || limit == 0 {
return Ok(Vec::new());
}
let mut sql = String::from(
"SELECT s.id, s.title, s.created_at, s.updated_at, s.project_path,
s.input_tokens, s.output_tokens,
s.cache_creation_input_tokens, s.cache_read_input_tokens,
COUNT(*), MIN(m.content)
FROM messages m JOIN sessions s ON s.id = m.session_id
WHERE m.content LIKE ?1 ESCAPE '\\'",
);
for i in 2..=terms.len() {
sql.push_str(&format!(" AND m.content LIKE ?{i} ESCAPE '\\'"));
}
sql.push_str(&format!(
" GROUP BY s.id ORDER BY s.updated_at DESC LIMIT {limit}"
));
let patterns: Vec<String> = terms.iter().map(|t| like_pattern(t)).collect();
let mut stmt = self.conn.prepare(&sql)?;
let rows = stmt.query_map(rusqlite::params_from_iter(&patterns), map_session_text_hit)?;
let mut result = Vec::new();
for row in rows {
result.push(row?);
}
Ok(result)
}
pub fn usage_totals(&self) -> crate::error::Result<UsageTotals> {
let (session_count, input, output, cache_create, cache_read): (
i64,
i64,
i64,
Option<i64>,
Option<i64>,
) = self.conn.query_row(
"SELECT
COUNT(*),
COALESCE(SUM(input_tokens), 0),
COALESCE(SUM(output_tokens), 0),
SUM(cache_creation_input_tokens),
SUM(cache_read_input_tokens)
FROM sessions",
[],
|row| {
Ok((
row.get(0)?,
row.get(1)?,
row.get(2)?,
row.get(3)?,
row.get(4)?,
))
},
)?;
let message_count: i64 =
self.conn
.query_row("SELECT COUNT(*) FROM messages", [], |row| row.get(0))?;
let usage = Usage {
input_tokens: input as u64,
output_tokens: output as u64,
cache_creation_input_tokens: cache_create.map(|v| v as u64),
cache_read_input_tokens: cache_read.map(|v| v as u64),
};
Ok(UsageTotals {
session_count: session_count as usize,
message_count: message_count as usize,
usage,
})
}
pub fn update_title(&self, id: &str, title: &str) -> crate::error::Result<()> {
let now = chrono::Utc::now().to_rfc3339();
self.conn.execute(
"UPDATE sessions SET title = ?1, updated_at = ?2 WHERE id = ?3",
rusqlite::params![title, now, id],
)?;
Ok(())
}
pub fn delete_session(&self, id: &str) -> crate::error::Result<()> {
self.conn.execute(
"DELETE FROM messages WHERE session_id = ?1",
rusqlite::params![id],
)?;
delete_branches(&self.conn, id)?;
self.conn
.execute("DELETE FROM sessions WHERE id = ?1", rusqlite::params![id])?;
Ok(())
}
pub fn insert_message(
&self,
msg_id: &str,
session_id: &str,
role: &str,
content: &str,
tool_call_id: Option<&str>,
name: Option<&str>,
) -> crate::error::Result<()> {
let now = chrono::Utc::now().to_rfc3339();
self.conn.execute(
"INSERT INTO messages (id, session_id, role, content, tool_call_id, name, created_at) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
rusqlite::params![msg_id, session_id, role, content, tool_call_id, name, now],
)?;
Ok(())
}
pub fn delete_messages(&self, session_id: &str) -> crate::error::Result<()> {
self.conn.execute(
"DELETE FROM messages WHERE session_id = ?1",
rusqlite::params![session_id],
)?;
Ok(())
}
pub fn replace_messages(
&self,
session_id: &str,
messages: &[MessageInsert],
) -> crate::error::Result<()> {
let tx = self.conn.unchecked_transaction()?;
write_messages(&tx, session_id, messages)?;
tx.commit()?;
Ok(())
}
pub fn replace_transcript(
&self,
session_id: &str,
messages: &[MessageInsert],
branches: &[BranchInsert],
) -> crate::error::Result<()> {
let tx = self.conn.unchecked_transaction()?;
write_messages(&tx, session_id, messages)?;
delete_branches(&tx, session_id)?;
for (ord, (name, rows)) in branches.iter().enumerate() {
let ord = ord as i64;
tx.execute(
"INSERT INTO session_branches (session_id, ord, name) VALUES (?1, ?2, ?3)",
rusqlite::params![session_id, ord, name],
)?;
for (seq, content) in rows.iter().enumerate() {
let seq = seq as i64;
tx.execute(
"INSERT INTO branch_messages (session_id, ord, seq, content) \
VALUES (?1, ?2, ?3, ?4)",
rusqlite::params![session_id, ord, seq, content],
)?;
}
}
tx.commit()?;
Ok(())
}
pub fn get_branches(&self, session_id: &str) -> crate::error::Result<Vec<BranchRow>> {
let mut branches = std::collections::BTreeMap::new();
let mut heads = self
.conn
.prepare("SELECT ord, name FROM session_branches WHERE session_id = ?1")?;
for head in heads.query_map(rusqlite::params![session_id], map_branch_head)? {
let (ord, name) = head?;
let row = BranchRow {
name,
messages: Vec::new(),
};
branches.insert(ord, row);
}
let mut rows = self.conn.prepare(
"SELECT ord, content FROM branch_messages WHERE session_id = ?1 ORDER BY ord, seq",
)?;
for row in rows.query_map(rusqlite::params![session_id], map_branch_message)? {
let (ord, content) = row?;
let branch: &mut BranchRow = branches.entry(ord).or_default();
branch.messages.push(content);
}
Ok(branches.into_values().collect())
}
pub fn message_counts_by_session(
&self,
) -> crate::error::Result<std::collections::HashMap<String, usize>> {
let mut stmt = self
.conn
.prepare("SELECT session_id, COUNT(*) FROM messages GROUP BY session_id")?;
let rows = stmt.query_map([], map_session_message_count)?;
let mut map = std::collections::HashMap::new();
for row in rows {
let (id, n) = row?;
map.insert(id, n);
}
Ok(map)
}
pub fn get_messages(&self, session_id: &str) -> crate::error::Result<Vec<MessageRow>> {
let mut stmt = self.conn.prepare(
"SELECT id, session_id, role, content, tool_call_id, name, created_at FROM messages WHERE session_id = ?1 ORDER BY created_at ASC"
)?;
let rows = stmt.query_map(rusqlite::params![session_id], map_message_row)?;
let mut result = Vec::new();
for row in rows {
result.push(row?);
}
Ok(result)
}
pub fn message_count(&self, session_id: &str) -> crate::error::Result<usize> {
let count: i64 = self.conn.query_row(
"SELECT COUNT(*) FROM messages WHERE session_id = ?1",
rusqlite::params![session_id],
|row| row.get(0),
)?;
Ok(count as usize)
}
pub fn get_state(&self, key: &str) -> crate::error::Result<Option<String>> {
let mut stmt = self
.conn
.prepare("SELECT value FROM state WHERE key = ?1")?;
let mut rows = stmt.query_map(rusqlite::params![key], |row| row.get::<_, String>(0))?;
Ok(rows.next().transpose()?)
}
pub fn set_state(&self, key: &str, value: &str) -> crate::error::Result<()> {
self.conn.execute(
"INSERT OR REPLACE INTO state (key, value) VALUES (?1, ?2)",
rusqlite::params![key, value],
)?;
Ok(())
}
pub fn delete_state(&self, key: &str) -> crate::error::Result<()> {
self.conn
.execute("DELETE FROM state WHERE key = ?1", rusqlite::params![key])?;
Ok(())
}
pub fn insert_memory(
&self,
id: &str,
project_key: &str,
text: &str,
embedding: &[u8],
source_session: Option<&str>,
) -> crate::error::Result<()> {
let now = chrono::Utc::now().to_rfc3339();
self.conn.execute(
"INSERT INTO memories (
id, project_key, text, embedding, source_session,
created_at, last_recalled_at, recall_count
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, NULL, 0)",
rusqlite::params![id, project_key, text, embedding, source_session, now],
)?;
Ok(())
}
pub fn get_memory(&self, id: &str) -> crate::error::Result<Option<MemoryRow>> {
let mut stmt = self.conn.prepare(
"SELECT id, project_key, text, embedding, source_session,
created_at, last_recalled_at, recall_count
FROM memories WHERE id = ?1",
)?;
let mut rows = stmt.query_map(rusqlite::params![id], map_memory_row)?;
Ok(rows.next().transpose()?)
}
pub fn list_memories(
&self,
project_key: &str,
limit: usize,
) -> crate::error::Result<Vec<MemoryRow>> {
let mut stmt = self.conn.prepare(
"SELECT id, project_key, text, embedding, source_session,
created_at, last_recalled_at, recall_count
FROM memories WHERE project_key = ?1
ORDER BY created_at DESC
LIMIT ?2",
)?;
let rows = stmt.query_map(rusqlite::params![project_key, limit as i64], map_memory_row)?;
let mut result = Vec::new();
for row in rows {
result.push(row?);
}
Ok(result)
}
pub fn delete_memory(&self, id: &str) -> crate::error::Result<bool> {
let n = self
.conn
.execute("DELETE FROM memories WHERE id = ?1", rusqlite::params![id])?;
Ok(n > 0)
}
pub fn clear_memories(&self, project_key: &str) -> crate::error::Result<usize> {
let n = self.conn.execute(
"DELETE FROM memories WHERE project_key = ?1",
rusqlite::params![project_key],
)?;
Ok(n)
}
pub fn touch_memory_recall(&self, id: &str) -> crate::error::Result<()> {
let now = chrono::Utc::now().to_rfc3339();
self.conn.execute(
"UPDATE memories SET last_recalled_at = ?1, recall_count = recall_count + 1
WHERE id = ?2",
rusqlite::params![now, id],
)?;
Ok(())
}
pub fn count_memories(&self, project_key: &str) -> crate::error::Result<usize> {
let n: i64 = self.conn.query_row(
"SELECT COUNT(*) FROM memories WHERE project_key = ?1",
rusqlite::params![project_key],
|row| row.get(0),
)?;
Ok(n as usize)
}
#[allow(clippy::too_many_arguments)]
pub fn insert_code_chunk(
&self,
id: &str,
project_key: &str,
path: &str,
start_line: i64,
end_line: i64,
text: &str,
embedding: &[u8],
) -> crate::error::Result<()> {
let now = chrono::Utc::now().to_rfc3339();
self.conn.execute(
"INSERT INTO code_chunks (
id, project_key, path, start_line, end_line, text, embedding, updated_at
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)",
rusqlite::params![
id,
project_key,
path,
start_line,
end_line,
text,
embedding,
now
],
)?;
Ok(())
}
pub fn list_code_chunks(
&self,
project_key: &str,
limit: usize,
) -> crate::error::Result<Vec<CodeChunkRow>> {
let mut stmt = self.conn.prepare(
"SELECT id, project_key, path, start_line, end_line, text, embedding, updated_at
FROM code_chunks WHERE project_key = ?1
ORDER BY path ASC, start_line ASC
LIMIT ?2",
)?;
let rows = stmt.query_map((project_key, limit as i64), map_code_chunk_row)?;
let mut result = Vec::new();
for row in rows {
result.push(row?);
}
Ok(result)
}
pub fn clear_code_chunks(&self, project_key: &str) -> crate::error::Result<usize> {
let n = self.conn.execute(
"DELETE FROM code_chunks WHERE project_key = ?1",
rusqlite::params![project_key],
)?;
Ok(n)
}
pub fn insert_session_chunk(
&self,
id: &str,
project_key: &str,
session_id: &str,
turn_index: i64,
text: &str,
embedding: &[u8],
) -> crate::error::Result<()> {
let now = chrono::Utc::now().to_rfc3339();
self.conn.execute(
"INSERT INTO session_chunks (
id, project_key, session_id, turn_index, text, embedding, created_at
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
rusqlite::params![
id,
project_key,
session_id,
turn_index,
text,
embedding,
now
],
)?;
Ok(())
}
pub fn list_session_chunks(
&self,
project_key: &str,
limit: usize,
) -> crate::error::Result<Vec<SessionChunkRow>> {
let mut stmt = self.conn.prepare(
"SELECT id, project_key, session_id, turn_index, text, embedding, created_at
FROM session_chunks WHERE project_key = ?1
ORDER BY created_at DESC
LIMIT ?2",
)?;
let rows = stmt.query_map((project_key, limit as i64), map_session_chunk_row)?;
let mut result = Vec::new();
for row in rows {
result.push(row?);
}
Ok(result)
}
pub fn search_memories(
&self,
project_key: &str,
query: &str,
limit: usize,
) -> crate::error::Result<FtsPage<MemoryRow>> {
self.fts_page(
"memories",
"t.id, t.project_key, t.text, t.embedding, t.source_session,
t.created_at, t.last_recalled_at, t.recall_count",
"t.created_at DESC",
(project_key, query, limit),
map_memory_row,
)
}
pub fn search_code_chunks(
&self,
project_key: &str,
query: &str,
limit: usize,
) -> crate::error::Result<FtsPage<CodeChunkRow>> {
self.fts_page(
"code_chunks",
"t.id, t.project_key, t.path, t.start_line, t.end_line, t.text,
t.embedding, t.updated_at",
"t.path ASC, t.start_line ASC",
(project_key, query, limit),
map_code_chunk_row,
)
}
pub fn search_session_chunks(
&self,
project_key: &str,
query: &str,
limit: usize,
) -> crate::error::Result<FtsPage<SessionChunkRow>> {
self.fts_page(
"session_chunks",
"t.id, t.project_key, t.session_id, t.turn_index, t.text,
t.embedding, t.created_at",
"t.created_at DESC",
(project_key, query, limit),
map_session_chunk_row,
)
}
fn fts_page<T>(
&self,
table: &str,
cols: &str,
browse_order: &str,
(project_key, query, limit): (&str, &str, usize),
map: fn(&rusqlite::Row<'_>) -> rusqlite::Result<T>,
) -> crate::error::Result<FtsPage<T>> {
let count = format!("SELECT COUNT(*) FROM {table} WHERE project_key = ?1");
let corpus: i64 = self
.conn
.query_row(&count, [project_key], |row| row.get(0))?;
let corpus = corpus as usize;
let browse = |n: usize| self.browse(table, cols, browse_order, project_key, n, map);
let Some(fts_query) = fts_match_query(query) else {
let rows = browse(limit)?;
return Ok(FtsPage {
rows,
matched: corpus,
corpus,
});
};
match self.fts_rows(table, cols, (project_key, &fts_query, limit), map) {
Ok((rows, matched)) => Ok(FtsPage {
rows,
matched,
corpus,
}),
Err(e) => {
tracing::warn!(table, error = %e, "fts index unavailable; scanning the bank");
let rows = browse(FTS_SCAN_LIMIT)?;
Ok(FtsPage {
matched: rows.len(),
rows,
corpus,
})
}
}
}
fn browse<T>(
&self,
table: &str,
cols: &str,
order: &str,
project_key: &str,
limit: usize,
map: fn(&rusqlite::Row<'_>) -> rusqlite::Result<T>,
) -> crate::error::Result<Vec<T>> {
let mut stmt = self.conn.prepare(&format!(
"SELECT {cols} FROM {table} t WHERE t.project_key = ?1 ORDER BY {order} LIMIT ?2"
))?;
let rows = stmt.query_map((project_key, limit as i64), map)?;
Ok(rows.collect::<rusqlite::Result<Vec<T>>>()?)
}
fn fts_rows<T>(
&self,
table: &str,
cols: &str,
(project_key, fts_query, limit): (&str, &str, usize),
map: fn(&rusqlite::Row<'_>) -> rusqlite::Result<T>,
) -> rusqlite::Result<(Vec<T>, usize)> {
let from = format!(
"FROM {table} t JOIN {table}_fts ON {table}_fts.rowid = t.rowid
WHERE {table}_fts MATCH ?1 AND t.project_key = ?2"
);
let matched: i64 = self.conn.query_row(
&format!("SELECT COUNT(*) {from}"),
(fts_query, project_key),
|row| row.get(0),
)?;
let mut stmt = self.conn.prepare(&format!(
"SELECT {cols} {from} ORDER BY bm25({table}_fts) LIMIT ?3"
))?;
let rows = stmt.query_map((fts_query, project_key, limit as i64), map)?;
Ok((
rows.collect::<rusqlite::Result<Vec<T>>>()?,
matched as usize,
))
}
#[cfg(any(test, feature = "test-util"))]
pub fn drop_fts_for_test(&self) -> crate::error::Result<()> {
for table in crate::migrations::FTS_TABLES {
self.conn.execute_batch(&format!(
"DROP TRIGGER {table}_fts_ai; DROP TRIGGER {table}_fts_ad;
DROP TRIGGER {table}_fts_au; DROP TABLE {table}_fts;"
))?;
}
Ok(())
}
}
pub const FTS_SCAN_LIMIT: usize = 50_000;
#[derive(Debug, Clone)]
pub struct FtsPage<T> {
pub rows: Vec<T>,
pub matched: usize,
pub corpus: usize,
}
pub fn fts_match_query(text: &str) -> Option<String> {
let mut words: Vec<String> = Vec::new();
for word in text.split(|c: char| !c.is_alphanumeric()) {
let word = word.to_lowercase();
if word.chars().count() >= 2 && !words.contains(&word) && words.len() < 32 {
words.push(word);
}
}
if words.is_empty() {
return None;
}
let quoted: Vec<String> = words.iter().map(|w| format!("\"{w}\"")).collect();
Some(quoted.join(" OR "))
}
fn map_message_row(row: &rusqlite::Row<'_>) -> rusqlite::Result<MessageRow> {
Ok(MessageRow {
id: row.get(0)?,
session_id: row.get(1)?,
role: row.get(2)?,
content: row.get(3)?,
tool_call_id: row.get(4)?,
name: row.get(5)?,
created_at: row.get(6)?,
})
}
fn delete_branches(conn: &Connection, session_id: &str) -> crate::error::Result<()> {
conn.execute(
"DELETE FROM branch_messages WHERE session_id = ?1",
rusqlite::params![session_id],
)?;
conn.execute(
"DELETE FROM session_branches WHERE session_id = ?1",
rusqlite::params![session_id],
)?;
Ok(())
}
fn write_messages(
tx: &rusqlite::Transaction<'_>,
session_id: &str,
messages: &[MessageInsert],
) -> crate::error::Result<()> {
tx.execute(
"DELETE FROM messages WHERE session_id = ?1",
rusqlite::params![session_id],
)?;
for (msg_id, role, content, tool_call_id, name, created_at) in messages {
tx.execute(
"INSERT INTO messages (id, session_id, role, content, tool_call_id, name, created_at) \
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
rusqlite::params![
msg_id,
session_id,
role,
content,
tool_call_id,
name,
created_at
],
)?;
}
Ok(())
}
fn map_branch_head(row: &rusqlite::Row<'_>) -> rusqlite::Result<(i64, Option<String>)> {
Ok((row.get(0)?, row.get(1)?))
}
fn map_branch_message(row: &rusqlite::Row<'_>) -> rusqlite::Result<(i64, String)> {
Ok((row.get(0)?, row.get(1)?))
}
fn map_session_message_count(row: &rusqlite::Row<'_>) -> rusqlite::Result<(String, usize)> {
let id: String = row.get(0)?;
let n: i64 = row.get(1)?;
Ok((id, n as usize))
}
fn map_memory_row(row: &rusqlite::Row<'_>) -> rusqlite::Result<MemoryRow> {
Ok(MemoryRow {
id: row.get(0)?,
project_key: row.get(1)?,
text: row.get(2)?,
embedding: row.get(3)?,
source_session: row.get(4)?,
created_at: row.get(5)?,
last_recalled_at: row.get(6)?,
recall_count: row.get(7)?,
})
}
fn map_session_chunk_row(row: &rusqlite::Row<'_>) -> rusqlite::Result<SessionChunkRow> {
Ok(SessionChunkRow {
id: row.get(0)?,
project_key: row.get(1)?,
session_id: row.get(2)?,
turn_index: row.get(3)?,
text: row.get(4)?,
embedding: row.get(5)?,
created_at: row.get(6)?,
})
}
fn map_code_chunk_row(row: &rusqlite::Row<'_>) -> rusqlite::Result<CodeChunkRow> {
Ok(CodeChunkRow {
id: row.get(0)?,
project_key: row.get(1)?,
path: row.get(2)?,
start_line: row.get(3)?,
end_line: row.get(4)?,
text: row.get(5)?,
embedding: row.get(6)?,
updated_at: row.get(7)?,
})
}
fn like_pattern(term: &str) -> String {
let mut out = String::with_capacity(term.len() + 2);
out.push('%');
for c in term.chars() {
if matches!(c, '%' | '_' | '\\') {
out.push('\\');
}
out.push(c);
}
out.push('%');
out
}
fn map_session_text_hit(row: &rusqlite::Row<'_>) -> rusqlite::Result<SessionTextHit> {
let hits: i64 = row.get(9)?;
Ok(SessionTextHit {
session: map_session_row(row)?,
hits: hits as usize,
sample: row.get(10)?,
})
}
fn map_session_row(row: &rusqlite::Row<'_>) -> rusqlite::Result<SessionRow> {
let input: i64 = row.get(5)?;
let output: i64 = row.get(6)?;
let cache_create: Option<i64> = row.get(7)?;
let cache_read: Option<i64> = row.get(8)?;
Ok(SessionRow {
id: row.get(0)?,
title: row.get(1)?,
created_at: row.get(2)?,
updated_at: row.get(3)?,
project_path: row.get(4)?,
usage: Usage {
input_tokens: input as u64,
output_tokens: output as u64,
cache_creation_input_tokens: cache_create.map(|v| v as u64),
cache_read_input_tokens: cache_read.map(|v| v as u64),
},
})
}
#[cfg(test)]
mod tests {
use super::*;
fn test_db() -> Database {
Database::open_in_memory().unwrap()
}
#[test]
fn test_open_while_another_connection_holds_a_write_transaction() {
let path = std::env::temp_dir().join(format!("whycodes-test-{}.db", uuid::Uuid::new_v4()));
let path_str = path.to_string_lossy().to_string();
let seed = Connection::open(&path_str).unwrap();
run_migrations(&seed).unwrap();
seed.execute_batch("BEGIN IMMEDIATE;").unwrap();
let opened = Database::open(&path_str);
let _ = seed.execute_batch("ROLLBACK;");
drop(seed);
for suffix in ["db", "db-wal", "db-shm"] {
let _ = std::fs::remove_file(path.with_extension(suffix));
}
opened.expect("open failed while another connection held a write transaction");
}
#[test]
fn test_create_and_get_session() {
let db = test_db();
db.create_session("s1", "Test", "/tmp").unwrap();
let s = db.get_session("s1").unwrap().unwrap();
assert_eq!(s.title, "Test");
assert!(s.usage.is_empty());
}
#[test]
fn test_upsert_preserves_created_at_and_stores_usage() {
let db = test_db();
let created = "2020-01-01T00:00:00+00:00";
let usage = Usage {
input_tokens: 10,
output_tokens: 20,
cache_creation_input_tokens: Some(1),
cache_read_input_tokens: Some(2),
};
db.upsert_session("s1", "Old", "/tmp", created, created, &usage)
.unwrap();
let updated = "2021-06-15T12:00:00+00:00";
let usage2 = Usage {
input_tokens: 100,
output_tokens: 200,
cache_creation_input_tokens: Some(3),
cache_read_input_tokens: Some(4),
};
db.upsert_session("s1", "New", "/proj", created, updated, &usage2)
.unwrap();
let s = db.get_session("s1").unwrap().unwrap();
assert_eq!(s.title, "New");
assert_eq!(s.created_at, created);
assert_eq!(s.updated_at, updated);
assert_eq!(s.project_path, "/proj");
assert_eq!(s.usage.input_tokens, 100);
assert_eq!(s.usage.output_tokens, 200);
assert_eq!(s.usage.cache_creation_input_tokens, Some(3));
assert_eq!(s.usage.cache_read_input_tokens, Some(4));
}
#[test]
fn test_usage_totals() {
let db = test_db();
let now = chrono::Utc::now().to_rfc3339();
db.upsert_session(
"a",
"A",
"/a",
&now,
&now,
&Usage {
input_tokens: 10,
output_tokens: 5,
cache_creation_input_tokens: Some(1),
cache_read_input_tokens: None,
},
)
.unwrap();
db.upsert_session(
"b",
"B",
"/b",
&now,
&now,
&Usage {
input_tokens: 20,
output_tokens: 7,
cache_creation_input_tokens: None,
cache_read_input_tokens: Some(4),
},
)
.unwrap();
db.insert_message("m1", "a", "user", "hi", None, None)
.unwrap();
db.insert_message("m2", "a", "assistant", "yo", None, None)
.unwrap();
let t = db.usage_totals().unwrap();
assert_eq!(t.session_count, 2);
assert_eq!(t.message_count, 2);
assert_eq!(t.usage.input_tokens, 30);
assert_eq!(t.usage.output_tokens, 12);
assert_eq!(t.usage.cache_creation_input_tokens, Some(1));
assert_eq!(t.usage.cache_read_input_tokens, Some(4));
assert_eq!(t.usage.total(), 30 + 12 + 1 + 4);
}
#[test]
fn test_list_sessions() {
let db = test_db();
db.create_session("a", "A", "/a").unwrap();
db.create_session("b", "B", "/b").unwrap();
assert_eq!(db.list_sessions().unwrap().len(), 2);
}
#[test]
fn test_update_title() {
let db = test_db();
db.create_session("s1", "Old", "/tmp").unwrap();
db.update_title("s1", "New").unwrap();
assert_eq!(db.get_session("s1").unwrap().unwrap().title, "New");
}
#[test]
fn test_delete_session() {
let db = test_db();
db.create_session("s1", "Test", "/tmp").unwrap();
db.insert_message("m1", "s1", "user", "hi", None, None)
.unwrap();
db.delete_session("s1").unwrap();
assert!(db.get_session("s1").unwrap().is_none());
assert!(db.get_messages("s1").unwrap().is_empty());
}
#[test]
fn test_insert_and_get_messages() {
let db = test_db();
db.create_session("s1", "Test", "/tmp").unwrap();
db.insert_message("m1", "s1", "user", "hello", None, None)
.unwrap();
db.insert_message("m2", "s1", "assistant", "hi", None, None)
.unwrap();
let msgs = db.get_messages("s1").unwrap();
assert_eq!(msgs.len(), 2);
assert_eq!(msgs[0].role, "user");
}
#[test]
fn test_delete_messages() {
let db = test_db();
db.create_session("s1", "Test", "/tmp").unwrap();
db.insert_message("m1", "s1", "user", "hi", None, None)
.unwrap();
db.delete_messages("s1").unwrap();
assert_eq!(db.message_count("s1").unwrap(), 0);
}
#[test]
fn test_message_count() {
let db = test_db();
db.create_session("s1", "Test", "/tmp").unwrap();
assert_eq!(db.message_count("s1").unwrap(), 0);
db.insert_message("m1", "s1", "user", "hi", None, None)
.unwrap();
assert_eq!(db.message_count("s1").unwrap(), 1);
}
#[test]
fn test_state_get_set_delete() {
let db = test_db();
assert!(db.get_state("k1").unwrap().is_none());
db.set_state("k1", "v1").unwrap();
assert_eq!(db.get_state("k1").unwrap().unwrap(), "v1");
db.delete_state("k1").unwrap();
assert!(db.get_state("k1").unwrap().is_none());
}
#[test]
fn test_nonexistent_session() {
let db = test_db();
assert!(db.get_session("no").unwrap().is_none());
}
#[test]
fn test_open_file_roundtrip_and_bad_paths() {
let tmp = tempfile::tempdir().unwrap();
let path = tmp.path().join("whycodes.db");
let path_str = path.to_str().unwrap();
let db = Database::open(path_str).unwrap();
db.create_session("s1", "File", "/proj").unwrap();
drop(db);
let db = Database::open(path_str).unwrap();
assert_eq!(db.get_session("s1").unwrap().unwrap().title, "File");
assert!(Database::open(tmp.path().to_str().unwrap()).is_err());
assert!(Database::open("/no/such/whycodes-storage-dir/db.sqlite").is_err());
}
#[test]
fn search_sessions_text_ands_terms_and_escapes_wildcards() {
let db = test_db();
db.create_session("s1", "Retry", "/tmp").unwrap();
db.create_session("s2", "Other", "/tmp").unwrap();
db.insert_message("a", "s1", "user", "fix the RETRY loop", None, None)
.unwrap();
db.insert_message("b", "s1", "assistant", "retry loop done", None, None)
.unwrap();
db.insert_message("c", "s2", "user", "retry only", None, None)
.unwrap();
db.insert_message("d", "s2", "user", r"50%_off \path", None, None)
.unwrap();
let terms = |t: &[&str]| t.iter().map(|s| s.to_string()).collect::<Vec<_>>();
let hits = db
.search_sessions_text(&terms(&["retry", "loop"]), 10)
.unwrap();
assert_eq!(hits.len(), 1);
assert_eq!(hits[0].session.title, "Retry");
assert_eq!(hits[0].hits, 2);
assert_eq!(hits[0].sample, "fix the RETRY loop");
assert_eq!(
db.search_sessions_text(&terms(&["retry"]), 1)
.unwrap()
.len(),
1
);
assert_eq!(
db.search_sessions_text(&terms(&["retry"]), 10)
.unwrap()
.len(),
2
);
assert!(db.search_sessions_text(&[], 10).unwrap().is_empty());
assert!(
db.search_sessions_text(&terms(&["retry"]), 0)
.unwrap()
.is_empty()
);
for literal in ["50%_off", r"\path"] {
assert_eq!(
db.search_sessions_text(&terms(&[literal]), 10)
.unwrap()
.len(),
1
);
}
assert!(
db.search_sessions_text(&terms(&["5%f"]), 10)
.unwrap()
.is_empty()
);
assert!(
db.search_sessions_text(&terms(&["50_"]), 10)
.unwrap()
.is_empty()
);
}
#[test]
fn test_replace_messages_and_counts() {
let db = test_db();
db.create_session("s1", "Test", "/tmp").unwrap();
db.insert_message("old", "s1", "user", "stale", None, None)
.unwrap();
db.insert_message("tool1", "s1", "tool", "out", Some("call-1"), Some("bash"))
.unwrap();
let created = "2020-01-01T00:00:00+00:00".to_string();
db.replace_messages(
"s1",
&[
(
"m1".into(),
"user".into(),
"hello".into(),
None,
None,
created.clone(),
),
(
"m2".into(),
"tool".into(),
"done".into(),
Some("tc1".into()),
Some("grep".into()),
created,
),
],
)
.unwrap();
let msgs = db.get_messages("s1").unwrap();
assert_eq!(msgs.len(), 2);
assert_eq!(msgs[0].id, "m1");
assert_eq!(msgs[1].tool_call_id.as_deref(), Some("tc1"));
assert_eq!(msgs[1].name.as_deref(), Some("grep"));
let counts = db.message_counts_by_session().unwrap();
assert_eq!(counts.get("s1"), Some(&2));
db.replace_messages("s1", &[]).unwrap();
assert!(db.get_messages("s1").unwrap().is_empty());
assert!(db.message_counts_by_session().unwrap().is_empty());
}
#[test]
fn test_replace_messages_duplicate_id_fails() {
let db = test_db();
db.create_session("s1", "Test", "/tmp").unwrap();
let created = "2020-01-01T00:00:00+00:00".to_string();
let err = db.replace_messages(
"s1",
&[
(
"m1".into(),
"user".into(),
"a".into(),
None,
None,
created.clone(),
),
("m1".into(), "user".into(), "b".into(), None, None, created),
],
);
assert!(err.is_err());
}
fn row(id: &str, content: &str) -> MessageInsert {
let created = "2020-01-01T00:00:00+00:00".to_string();
(
id.into(),
"user".into(),
content.into(),
None,
None,
created,
)
}
fn branch(name: Option<&str>, msgs: &[&str]) -> BranchInsert {
(
name.map(str::to_string),
msgs.iter().map(|m| m.to_string()).collect(),
)
}
#[test]
fn replace_transcript_round_trips_branches_in_order() {
let db = test_db();
db.create_session("s1", "Test", "/tmp").unwrap();
db.create_session("s2", "Other", "/tmp").unwrap();
db.replace_transcript(
"s1",
&[row("m1", "a"), row("m2", "b")],
&[
branch(Some("try"), &["a", "x"]),
branch(None, &["a", "y", "z"]),
],
)
.unwrap();
db.replace_transcript("s2", &[row("o1", "q")], &[branch(None, &["q", "r"])])
.unwrap();
assert_eq!(db.get_messages("s1").unwrap().len(), 2);
let branches = db.get_branches("s1").unwrap();
assert_eq!(
branches,
vec![
BranchRow {
name: Some("try".into()),
messages: vec!["a".into(), "x".into()],
},
BranchRow {
name: None,
messages: vec!["a".into(), "y".into(), "z".into()],
},
]
);
db.replace_transcript("s1", &[row("m1", "a")], &[branch(None, &["b"])])
.unwrap();
assert_eq!(db.get_branches("s1").unwrap().len(), 1);
assert_eq!(db.get_branches("s2").unwrap().len(), 1);
db.replace_messages("s1", &[row("m9", "c")]).unwrap();
assert_eq!(db.get_branches("s1").unwrap().len(), 1);
db.delete_session("s1").unwrap();
assert!(db.get_branches("s1").unwrap().is_empty());
assert_eq!(db.get_branches("s2").unwrap().len(), 1);
}
#[test]
fn get_branches_keeps_message_rows_without_a_header() {
let db = test_db();
db.conn
.execute_batch(
"INSERT INTO branch_messages (session_id, ord, seq, content)
VALUES ('s', 3, 1, 'second'), ('s', 3, 0, 'first');",
)
.unwrap();
assert_eq!(
db.get_branches("s").unwrap(),
vec![BranchRow {
name: None,
messages: vec!["first".into(), "second".into()],
}]
);
}
#[test]
fn database_from_the_previous_schema_migrates_forward() {
let tmp = tempfile::tempdir().unwrap();
let path = tmp.path().join("old.db");
{
let conn = Connection::open(&path).unwrap();
conn.execute_batch(
"
CREATE TABLE sessions (
id TEXT PRIMARY KEY, title TEXT, created_at TEXT,
updated_at TEXT, project_path TEXT,
input_tokens INTEGER NOT NULL DEFAULT 0,
output_tokens INTEGER NOT NULL DEFAULT 0,
cache_creation_input_tokens INTEGER,
cache_read_input_tokens INTEGER
);
CREATE TABLE messages (
id TEXT PRIMARY KEY, session_id TEXT, role TEXT, content TEXT,
tool_call_id TEXT, name TEXT, created_at TEXT,
FOREIGN KEY (session_id) REFERENCES sessions(id)
);
CREATE INDEX idx_messages_session ON messages(session_id);
CREATE TABLE state (key TEXT PRIMARY KEY, value TEXT);
INSERT INTO sessions (id, title, created_at, updated_at, project_path)
VALUES ('s', 'old', '2020-01-01T00:00:00+00:00',
'2020-01-01T00:00:00+00:00', '/p');
INSERT INTO messages (id, session_id, role, content, created_at) VALUES
('m1', 's', 'user', 'one', '2020-01-01T00:00:01+00:00'),
('m2', 's', 'assistant', 'two', '2020-01-01T00:00:02+00:00');
",
)
.unwrap();
}
let db = Database::open(&path.to_string_lossy()).unwrap();
let msgs = db.get_messages("s").unwrap();
assert_eq!(
msgs.iter().map(|m| m.content.as_str()).collect::<Vec<_>>(),
["one", "two"]
);
assert!(db.get_branches("s").unwrap().is_empty());
db.replace_transcript("s", &[row("m1", "one")], &[branch(None, &["one", "x"])])
.unwrap();
drop(db);
let db = Database::open(&path.to_string_lossy()).unwrap();
assert_eq!(db.get_messages("s").unwrap().len(), 1);
assert_eq!(db.get_branches("s").unwrap()[0].messages, ["one", "x"]);
}
#[test]
fn branch_writes_and_reads_report_sql_errors() {
let broken = |sql: &str| {
let db = test_db();
db.create_session("s", "t", "/p").unwrap();
db.conn.execute_batch(sql).unwrap();
db
};
let msgs = [row("m", "a")];
let one = [branch(None, &["a"])];
let db = broken(
"DROP TABLE session_branches;
CREATE TABLE session_branches (session_id TEXT);",
);
db.replace_messages("s", &[row("old", "kept")]).unwrap();
assert!(db.replace_transcript("s", &msgs, &one).is_err());
assert_eq!(db.get_messages("s").unwrap()[0].content, "kept");
assert!(db.get_branches("s").is_err());
let db = broken(
"DROP TABLE branch_messages;
CREATE TABLE branch_messages (session_id TEXT);",
);
assert!(db.replace_transcript("s", &msgs, &one).is_err());
assert!(db.get_branches("s").is_err());
let db = broken("DROP TABLE branch_messages;");
assert!(db.replace_transcript("s", &msgs, &[]).is_err());
assert!(db.delete_session("s").is_err());
let db = broken("DROP TABLE session_branches;");
assert!(db.replace_transcript("s", &msgs, &[]).is_err());
assert!(db.get_branches("s").is_err());
let db = broken("DROP TABLE messages;");
assert!(db.replace_transcript("s", &msgs, &[]).is_err());
}
#[test]
fn test_insert_message_foreign_key() {
let db = test_db();
assert!(
db.insert_message("m1", "missing", "user", "hi", None, None)
.is_err()
);
}
#[test]
fn test_usage_totals_empty_database() {
let db = test_db();
let t = db.usage_totals().unwrap();
assert_eq!(t.session_count, 0);
assert_eq!(t.message_count, 0);
assert!(t.usage.is_empty());
}
#[test]
fn test_memories_crud() {
let db = test_db();
let emb = [1_u8, 2, 3, 4];
db.insert_memory("mem1", "proj", "note", &emb, Some("s1"))
.unwrap();
db.insert_memory("mem2", "proj", "other", &emb, None)
.unwrap();
db.insert_memory("mem3", "elsewhere", "x", &emb, None)
.unwrap();
let got = db.get_memory("mem1").unwrap().unwrap();
assert_eq!(got.project_key, "proj");
assert_eq!(got.text, "note");
assert_eq!(got.embedding, emb);
assert_eq!(got.source_session.as_deref(), Some("s1"));
assert!(got.last_recalled_at.is_none());
assert_eq!(got.recall_count, 0);
assert!(db.get_memory("missing").unwrap().is_none());
let listed = db.list_memories("proj", 1).unwrap();
assert_eq!(listed.len(), 1);
assert_eq!(db.count_memories("proj").unwrap(), 2);
db.touch_memory_recall("mem1").unwrap();
let got = db.get_memory("mem1").unwrap().unwrap();
assert_eq!(got.recall_count, 1);
assert!(got.last_recalled_at.is_some());
assert!(db.delete_memory("mem2").unwrap());
assert!(!db.delete_memory("mem2").unwrap());
assert_eq!(db.count_memories("proj").unwrap(), 1);
assert_eq!(db.clear_memories("proj").unwrap(), 1);
assert_eq!(db.count_memories("proj").unwrap(), 0);
assert!(db.list_memories("proj", 10).unwrap().is_empty());
}
#[test]
fn test_code_chunks() {
let db = test_db();
let emb = [9_u8, 8, 7];
db.insert_code_chunk("c1", "proj", "src/a.rs", 1, 10, "fn a() {}", &emb)
.unwrap();
db.insert_code_chunk("c2", "proj", "src/b.rs", 2, 4, "fn b() {}", &emb)
.unwrap();
db.insert_code_chunk("c3", "other", "x.rs", 1, 1, "x", &emb)
.unwrap();
let rows = db.list_code_chunks("proj", 10).unwrap();
assert_eq!(rows.len(), 2);
assert_eq!(rows[0].path, "src/a.rs");
assert_eq!(rows[0].start_line, 1);
assert_eq!(rows[0].end_line, 10);
assert_eq!(rows[0].text, "fn a() {}");
assert_eq!(rows[0].embedding, emb);
assert_eq!(db.list_code_chunks("proj", 1).unwrap().len(), 1);
assert_eq!(db.clear_code_chunks("proj").unwrap(), 2);
assert!(db.list_code_chunks("proj", 10).unwrap().is_empty());
}
#[test]
fn test_session_chunks() {
let db = test_db();
let emb = [4_u8, 5];
db.insert_session_chunk("sc1", "proj", "s1", 0, "turn0", &emb)
.unwrap();
db.insert_session_chunk("sc2", "proj", "s1", 1, "turn1", &emb)
.unwrap();
db.insert_session_chunk("sc3", "other", "s2", 0, "x", &emb)
.unwrap();
let rows = db.list_session_chunks("proj", 10).unwrap();
assert_eq!(rows.len(), 2);
assert!(rows.iter().all(|r| r.project_key == "proj"));
assert!(rows.iter().any(|r| r.turn_index == 0 && r.text == "turn0"));
assert_eq!(rows.iter().find(|r| r.id == "sc1").unwrap().embedding, emb);
assert_eq!(db.list_session_chunks("proj", 1).unwrap().len(), 1);
}
#[test]
fn test_sql_error_paths_after_dropping_schema() {
let db = test_db();
db.drop_sessions_table_for_test().unwrap();
assert!(db.list_sessions().is_err());
db.conn
.execute_batch(
"
DROP TABLE IF EXISTS session_chunks;
DROP TABLE IF EXISTS code_chunks;
DROP TABLE IF EXISTS memories;
DROP TABLE IF EXISTS messages;
DROP TABLE IF EXISTS state;
",
)
.unwrap();
let usage = Usage::default();
let now = "2020-01-01T00:00:00+00:00";
assert!(db.upsert_session("s", "t", "/p", now, now, &usage).is_err());
assert!(db.create_session("s", "t", "/p").is_err());
assert!(db.get_session("s").is_err());
assert!(db.list_sessions().is_err());
assert!(db.usage_totals().is_err());
assert!(db.update_title("s", "t").is_err());
assert!(db.delete_session("s").is_err());
assert!(
db.insert_message("m", "s", "user", "hi", None, None)
.is_err()
);
assert!(db.delete_messages("s").is_err());
assert!(db.replace_messages("s", &[]).is_err());
assert!(db.message_counts_by_session().is_err());
assert!(db.get_messages("s").is_err());
assert!(db.message_count("s").is_err());
assert!(db.get_state("k").is_err());
assert!(db.set_state("k", "v").is_err());
assert!(db.delete_state("k").is_err());
assert!(db.insert_memory("id", "pk", "t", &[], None).is_err());
assert!(db.get_memory("id").is_err());
assert!(db.list_memories("pk", 10).is_err());
assert!(db.delete_memory("id").is_err());
assert!(db.clear_memories("pk").is_err());
assert!(db.touch_memory_recall("id").is_err());
assert!(db.count_memories("pk").is_err());
assert!(
db.insert_code_chunk("id", "pk", "p.rs", 1, 2, "t", &[])
.is_err()
);
assert!(db.list_code_chunks("pk", 10).is_err());
assert!(db.clear_code_chunks("pk").is_err());
assert!(
db.insert_session_chunk("id", "pk", "s", 0, "t", &[])
.is_err()
);
assert!(db.list_session_chunks("pk", 10).is_err());
}
#[test]
fn open_existing_readonly_none_when_missing() {
let tmp = tempfile::tempdir().unwrap();
let path = tmp.path().join("no-such.db");
assert!(
Database::open_existing_readonly(&path.to_string_lossy())
.unwrap()
.is_none()
);
}
#[test]
fn open_existing_readonly_lists_sessions() {
let tmp = tempfile::tempdir().unwrap();
let path = tmp.path().join("whycodes.db");
{
let db = Database::open(path.to_str().unwrap()).unwrap();
db.create_session("abc123456789", "t", "/p").unwrap();
}
let ro = Database::open_existing_readonly(path.to_str().unwrap())
.unwrap()
.expect("existing db");
let sessions = ro.list_sessions().unwrap();
assert_eq!(sessions.len(), 1);
assert_eq!(sessions[0].id, "abc123456789");
}
#[test]
fn fts_match_query_quotes_words_and_drops_noise() {
assert_eq!(
fts_match_query("Bundled SQLite, bundled! a").as_deref(),
Some("\"bundled\" OR \"sqlite\"")
);
assert_eq!(
fts_match_query("NOT x* OR \"y").as_deref(),
Some("\"not\" OR \"or\"")
);
assert_eq!(fts_match_query(" ? ! a "), None);
let many = (0..40)
.map(|i| format!("w{i}"))
.collect::<Vec<_>>()
.join(" ");
assert_eq!(fts_match_query(&many).unwrap().matches(" OR ").count(), 31);
}
#[test]
fn search_memories_ranks_matches_and_counts_without_loading_the_bank() {
let db = test_db();
db.insert_memory("m1", "p", "windows ci links bundled sqlite", b"", None)
.unwrap();
db.insert_memory("m2", "p", "sqlite sqlite sqlite wal mode", b"", None)
.unwrap();
db.insert_memory("m3", "p", "release notes live in docs", b"", None)
.unwrap();
db.insert_memory("m4", "other", "sqlite in another bank", b"", None)
.unwrap();
let page = db.search_memories("p", "sqlite", 1).unwrap();
assert_eq!((page.matched, page.corpus, page.rows.len()), (2, 3, 1));
assert_eq!(page.rows[0].id, "m2", "bm25 ranks the denser match first");
let none = db.search_memories("p", "kubernetes", 10).unwrap();
assert_eq!((none.matched, none.corpus), (0, 3));
assert!(none.rows.is_empty());
let browse = db.search_memories("p", "?", 2).unwrap();
assert_eq!(
(browse.matched, browse.corpus, browse.rows.len()),
(3, 3, 2)
);
db.delete_memory("m2").unwrap();
assert_eq!(db.search_memories("p", "sqlite", 5).unwrap().matched, 1);
db.clear_memories("p").unwrap();
assert_eq!(db.search_memories("p", "sqlite", 5).unwrap().corpus, 0);
}
#[test]
fn search_code_and_session_chunks_use_their_fts_index() {
let db = test_db();
db.insert_code_chunk("c1", "p", "b.rs", 1, 9, "fn list_memories()", b"")
.unwrap();
db.insert_code_chunk("c2", "p", "a.rs", 3, 7, "fn other()", b"")
.unwrap();
let code = db.search_code_chunks("p", "memories", 5).unwrap();
assert_eq!((code.matched, code.corpus), (1, 2));
assert_eq!(code.rows[0].path, "b.rs");
let browse = db.search_code_chunks("p", "", 5).unwrap();
assert_eq!(browse.rows[0].path, "a.rs", "browse is path order");
db.clear_code_chunks("p").unwrap();
assert_eq!(
db.search_code_chunks("p", "memories", 5).unwrap().matched,
0
);
db.insert_session_chunk("s1", "p", "ses", 0, "User: retry loop", b"")
.unwrap();
db.insert_session_chunk("s2", "p", "ses", 1, "User: unrelated", b"")
.unwrap();
let sess = db.search_session_chunks("p", "retry", 5).unwrap();
assert_eq!((sess.matched, sess.corpus), (1, 2));
assert_eq!(sess.rows[0].id, "s1");
assert_eq!(db.search_session_chunks("p", "", 5).unwrap().rows.len(), 2);
}
#[test]
fn search_without_fts_index_scans_the_bank() {
let db = test_db();
db.insert_memory("m1", "p", "sqlite", b"", None).unwrap();
db.drop_fts_for_test().unwrap();
db.insert_memory("m2", "p", "unrelated", b"", None).unwrap();
let page = db.search_memories("p", "sqlite", 1).unwrap();
assert_eq!((page.rows.len(), page.matched, page.corpus), (2, 2, 2));
}
}