#![deny(unsafe_code)]
#![warn(missing_docs)]
mod clock;
mod error;
mod export;
mod hlc;
mod migrate;
mod model;
mod ops;
mod query;
mod search;
mod vecext;
mod writer;
use std::path::Path;
use std::sync::mpsc::{Receiver, Sender, channel};
use std::sync::{Arc, Mutex};
use std::thread::JoinHandle;
use rusqlite::{Connection, OpenFlags};
use uuid::Uuid;
pub use crate::clock::{Clock, ManualClock, SystemClock, TimestampMs};
pub use crate::error::{Error, Result};
pub use crate::export::{ImportStats, import_jsonl};
pub use crate::hlc::{Hlc, HlcGenerator};
pub use crate::migrate::SCHEMA_VERSION;
pub use crate::model::{Edge, Episode, Node, Subgraph};
pub use crate::ops::{GraphOp, OplogEntry};
#[doc(hidden)]
pub use crate::query::{EXPAND_SQL_AS_AT, EXPAND_SQL_CURRENT};
pub use crate::query::{MergeCandidate, NodeHistory, Stats, Traversal};
pub use crate::search::{Budget, Query, SearchHit, SearchKind, SearchTarget};
use crate::writer::{VecTable, WriteMsg};
pub struct Options {
pub device_id: String,
pub clock: Arc<dyn Clock>,
}
impl Options {
pub fn new(device_id: impl Into<String>) -> Self {
Self {
device_id: device_id.into(),
clock: Arc::new(SystemClock),
}
}
pub fn clock(mut self, clock: Arc<dyn Clock>) -> Self {
self.clock = clock;
self
}
}
pub struct Grit {
write_tx: Sender<WriteMsg>,
writer: Option<JoinHandle<()>>,
pub(crate) read_conn: Mutex<Connection>,
clock: Arc<dyn Clock>,
hlc: HlcGenerator,
device_id: String,
}
impl Grit {
pub fn open(path: impl AsRef<Path>, opts: Options) -> Result<Self> {
let path = path.as_ref();
if let Some(s) = path.to_str()
&& (s.is_empty()
|| s == ":memory:"
|| (s.starts_with("file:")
&& (s.contains(":memory:") || s.contains("mode=memory"))))
{
return Err(Error::InvalidOp(
"in-memory or empty database paths are not supported".into(),
));
}
vecext::register_sqlite_vec();
let mut write_conn = Connection::open(path)?;
configure(&write_conn)?;
migrate::migrate(&mut write_conn)?;
let read_conn = Connection::open_with_flags(
path,
OpenFlags::SQLITE_OPEN_READ_ONLY | OpenFlags::SQLITE_OPEN_NO_MUTEX,
)?;
read_conn.busy_timeout(std::time::Duration::from_secs(5))?;
configure_cache(&read_conn)?;
let (write_tx, write_rx): (Sender<WriteMsg>, Receiver<WriteMsg>) = channel();
let writer = writer::spawn(write_conn, Arc::clone(&opts.clock), write_rx);
Ok(Self {
write_tx,
writer: Some(writer),
read_conn: Mutex::new(read_conn),
hlc: HlcGenerator::new(opts.device_id.clone()),
clock: opts.clock,
device_id: opts.device_id,
})
}
pub fn new_id(&self) -> Uuid {
let ms = self.clock.now_ms().max(0) as u64;
let ts = uuid::Timestamp::from_unix(
uuid::NoContext,
ms / 1000,
((ms % 1000) * 1_000_000) as u32,
);
Uuid::new_v7(ts)
}
pub fn apply(&self, op: GraphOp) -> Result<OplogEntry> {
validate(&op)?;
let entry = OplogEntry {
id: self.new_id(),
hlc: self.hlc.next(self.clock.as_ref()),
device_id: self.device_id.clone(),
op,
};
writer::call(&self.write_tx, |reply| WriteMsg::Apply {
entry: entry.clone(),
reply,
})?;
Ok(entry)
}
pub fn apply_remote(&self, entry: OplogEntry) -> Result<bool> {
if entry.hlc.wall_ms < 0 {
return Err(Error::InvalidHlc(entry.hlc.encode()));
}
self.hlc.observe(&entry.hlc);
writer::call(&self.write_tx, |reply| WriteMsg::Apply { entry, reply })
}
pub fn oplog_since(&self, after_seq: i64) -> Result<Vec<(i64, OplogEntry)>> {
let conn = self.read();
let mut stmt = conn.prepare_cached(
"SELECT seq, id, hlc, device_id, op FROM oplog WHERE seq > ?1 ORDER BY seq",
)?;
let rows = stmt.query_map([after_seq], |row| {
Ok((
row.get::<_, i64>(0)?,
row.get::<_, String>(1)?,
row.get::<_, String>(2)?,
row.get::<_, String>(3)?,
row.get::<_, String>(4)?,
))
})?;
let mut out = Vec::new();
for row in rows {
let (seq, id, hlc, device_id, op) = row?;
out.push((
seq,
OplogEntry {
id: Uuid::parse_str(&id)
.map_err(|_| Error::Corrupt(format!("bad uuid in oplog: {id}")))?,
hlc: hlc.parse()?,
device_id,
op: serde_json::from_str(&op)?,
},
));
}
Ok(out)
}
pub fn register_embedding_model(
&self,
model_id: impl Into<String>,
dim: usize,
model_version: impl Into<String>,
) -> Result<()> {
writer::call(&self.write_tx, |reply| WriteMsg::RegisterModel {
model_id: model_id.into(),
dim,
version: model_version.into(),
reply,
})
}
pub fn set_node_embedding(&self, node_id: Uuid, vector: Vec<f32>) -> Result<()> {
writer::call(&self.write_tx, |reply| WriteMsg::SetEmbedding {
table: VecTable::Nodes,
id: node_id,
vector,
reply,
})
}
pub fn set_edge_embedding(&self, edge_id: Uuid, vector: Vec<f32>) -> Result<()> {
writer::call(&self.write_tx, |reply| WriteMsg::SetEmbedding {
table: VecTable::Edges,
id: edge_id,
vector,
reply,
})
}
pub(crate) fn now_ms(&self) -> TimestampMs {
self.clock.now_ms()
}
pub(crate) fn read(&self) -> std::sync::MutexGuard<'_, Connection> {
self.read_conn
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
}
}
const _: () = {
const fn assert_send_sync<T: Send + Sync>() {}
assert_send_sync::<Grit>();
};
impl Drop for Grit {
fn drop(&mut self) {
let _ = self.write_tx.send(WriteMsg::Shutdown);
if let Some(handle) = self.writer.take() {
let _ = handle.join();
}
}
}
fn validate(op: &GraphOp) -> Result<()> {
match op {
GraphOp::MergeNodes { from, into } if from == into => {
Err(Error::InvalidOp("cannot merge a node into itself".into()))
}
GraphOp::Purge { ids } if ids.is_empty() => {
Err(Error::InvalidOp("purge with no ids".into()))
}
GraphOp::UpdateNode {
name: None,
summary: None,
kind: None,
attrs: None,
..
} => Err(Error::InvalidOp("update with no fields".into())),
_ => Ok(()),
}
}
fn configure(conn: &Connection) -> Result<()> {
conn.busy_timeout(std::time::Duration::from_secs(5))?;
conn.pragma_update(None, "journal_mode", "WAL")?;
conn.pragma_update(None, "synchronous", "NORMAL")?;
configure_cache(conn)?;
Ok(())
}
fn configure_cache(conn: &Connection) -> Result<()> {
conn.pragma_update(None, "cache_size", -32_768)?; conn.pragma_update(None, "mmap_size", 268_435_456)?; Ok(())
}