crates/knut-graphdb/src/cozo.rs
//! The **embedded CozoDB** backend — an in-process Datalog store (the `mem`
//! engine: no container, no redis, no socket). It answers the same [`GraphDb`] ops
//! as FalkorDB, compiled to **CozoScript/Datalog** instead of Cypher.
//!
//! ## Schema — genuinely multi-graph
//! Two stored relations mirror the property-graph shape a viewer renders, keyed by
//! the graph NAME so one embedded instance serves the same
//! `graph_names`/`open`/`read` surface a FalkorDB endpoint does:
//! ```text
//! node {g: String, id: Int => label: String, nk: String?} -- nk = raw natural key
//! edge {g: String, src: Int, dst: Int, rel: String} -- all-key ⇒ idempotent :put
//! ```
//!
//! ## Cypher: none, and that is a complete backend
//! Cozo speaks Datalog. It does not implement [`crate::Cypher`] AT ALL — there is no
//! stub here to panic from, and [`GraphDb::cypher_support`] answers
//! [`CypherSupport::None`] by taking the trait's default. A caller offering a query
//! bar gets [`crate::no_cypher_here`], a typed value.
//!
//! ## Scoped
//! Implemented: `graph_names`, `open`, `read`, `neighbours`, `upsert`. **Shortest
//! path** is deferred: it returns [`GraphError::Unsupported`] (a clear, non-panicking
//! signal), not a `todo!()`.
use std::collections::BTreeSet;
use std::sync::OnceLock;
use async_trait::async_trait;
use cozo::{DbInstance, ScriptMutability};
use crate::{
GraphDb, GraphEdge, GraphError, GraphNode, GraphSchema, SubGraph, UpsertEdge, UpsertNode, cap,
};
/// The graph name a caller gets when it does not name one — the embedded store's
/// default graph, the Cozo analogue of FalkorDB's first `GRAPH.LIST` entry.
pub const DEFAULT_GRAPH: &str = "korp";
/// An embedded CozoDB graph store (mem engine). `DbInstance` is `Arc`-backed, so
/// this is cheap to clone and shares one in-memory store across clones.
#[derive(Clone)]
pub struct CozoBackend {
db: DbInstance,
}
impl CozoBackend {
/// Open a fresh in-memory Cozo store and create the `node`/`edge` relations.
/// No file, no container — it lives entirely in this process.
pub fn new() -> Result<Self, GraphError> {
let db = DbInstance::new("mem", "", "")
.map_err(|e| GraphError::Connect(format!("cozo mem: {e:?}")))?;
let this = Self { db };
// Each `:create` is its own Cozo transaction — run them separately.
this.run_mut(":create node {g: String, id: Int => label: String, nk: String?}")?;
this.run_mut(":create edge {g: String, src: Int, dst: Int, rel: String}")?;
Ok(this)
}
/// The **process-global** shared store, created + demo-seeded once. A consumer's
/// Cozo backend must persist across ops (a per-op mem DB would always be empty),
/// and starts with a small built-in demo scene so switching to it renders a real
/// graph with no container and no populate step.
pub fn shared() -> Result<Self, GraphError> {
static SHARED: OnceLock<Result<CozoBackend, String>> = OnceLock::new();
SHARED
.get_or_init(|| {
let be = CozoBackend::new().map_err(|e| e.to_string())?;
be.seed_demo().map_err(|e| e.to_string())?;
Ok(be)
})
.clone()
.map_err(GraphError::Connect)
}
/// Run a mutating CozoScript, discarding the result rows.
fn run_mut(&self, script: &str) -> Result<(), GraphError> {
self.db
.run_script(script, Default::default(), ScriptMutability::Mutable)
.map(|_| ())
.map_err(|e| GraphError::Backend(format!("{e:?}")))
}
/// Run a read-only CozoScript and hand back its rows.
fn run_rows(&self, script: &str) -> Result<Vec<Vec<cozo::DataValue>>, GraphError> {
Ok(self
.db
.run_script(script, Default::default(), ScriptMutability::Immutable)
.map_err(|e| GraphError::Backend(format!("{e:?}")))?
.rows)
}
/// Run a read-only CozoScript and shape the result into a [`SubGraph`]. The
/// query MUST project `src, dst, rel, sl, dl, sk, dk` (in that order).
fn run_scene(&self, script: &str) -> Result<SubGraph, GraphError> {
let rows = self.run_rows(script)?;
let mut nodes: Vec<GraphNode> = Vec::new();
let mut edges: Vec<GraphEdge> = Vec::new();
let mut seen = std::collections::HashSet::new();
for row in rows {
if row.len() < 7 {
return Err(GraphError::Malformed(format!(
"scene row has {} columns, expected 7 (src,dst,rel,sl,dl,sk,dk)",
row.len()
)));
}
let src = row[0]
.get_int()
.ok_or_else(|| GraphError::Malformed("src not int".into()))?;
let dst = row[1]
.get_int()
.ok_or_else(|| GraphError::Malformed("dst not int".into()))?;
let rel = row[2].get_str().unwrap_or("").to_string();
let sl = row[3].get_str().unwrap_or("").to_string();
let dl = row[4].get_str().unwrap_or("").to_string();
let sk = row[5].get_str().map(str::to_string);
let dk = row[6].get_str().map(str::to_string);
if seen.insert(src) {
nodes.push(GraphNode::new(src, sl).with_key(sk));
}
if seen.insert(dst) {
nodes.push(GraphNode::new(dst, dl).with_key(dk));
}
edges.push(GraphEdge {
src_id: src,
dst_id: dst,
rel,
});
}
Ok(SubGraph { nodes, edges })
}
/// Seed a small demo scene (Persons + a Company + a Vessel and their links) into
/// [`DEFAULT_GRAPH`] so a container-free Cozo backend is non-empty.
pub fn seed_demo(&self) -> Result<u64, GraphError> {
let nodes = [
UpsertNode::new(1, "Person"),
UpsertNode::new(2, "Person"),
UpsertNode::new(3, "Company").with_key("5560000001"),
UpsertNode::new(4, "Vessel").with_key("265547890"),
UpsertNode::new(5, "Address"),
];
let edges = [
UpsertEdge::new(1, 3, "WORKS_AT"),
UpsertEdge::new(2, 3, "WORKS_AT"),
UpsertEdge::new(1, 4, "OWNS"),
UpsertEdge::new(1, 5, "LIVES_AT"),
UpsertEdge::new(2, 5, "LIVES_AT"),
];
self.upsert_sync(DEFAULT_GRAPH, &nodes, &edges)
}
/// The synchronous upsert body (the async trait method just awaits this).
fn upsert_sync(
&self,
graph: &str,
nodes: &[UpsertNode],
edges: &[UpsertEdge],
) -> Result<u64, GraphError> {
let g = escape(graph);
let mut written = 0u64;
if !nodes.is_empty() {
let rows: Vec<String> = nodes
.iter()
.map(|n| {
let nk = match &n.key {
Some(k) => format!("'{}'", escape(k)),
None => "null".to_string(),
};
format!("['{g}', {}, '{}', {}]", n.id, escape(&n.label), nk)
})
.collect();
let script = format!("?[g, id, label, nk] <- [{}]\n:put node", rows.join(", "));
self.run_mut(&script)?;
written += nodes.len() as u64;
}
if !edges.is_empty() {
let rows: Vec<String> = edges
.iter()
.map(|e| format!("['{g}', {}, {}, '{}']", e.src_id, e.dst_id, escape(&e.rel)))
.collect();
let script = format!("?[g, src, dst, rel] <- [{}]\n:put edge", rows.join(", "));
self.run_mut(&script)?;
written += edges.len() as u64;
}
Ok(written)
}
/// Every graph name that carries at least one node OR one edge.
fn names(&self) -> Result<Vec<String>, GraphError> {
let mut out: BTreeSet<String> = BTreeSet::new();
for script in [
"?[g] := *node{g}",
"?[g] := *edge{g}",
] {
for row in self.run_rows(script)? {
if let Some(s) = row.first().and_then(|v| v.get_str()) {
out.insert(s.to_string());
}
}
}
Ok(out.into_iter().collect())
}
}
#[async_trait]
impl GraphDb for CozoBackend {
fn kind(&self) -> &'static str {
"cozo"
}
/// An embedded store has no socket. Its "endpoint" says exactly that, so a
/// diagnostic never prints a host that does not exist.
fn endpoint(&self) -> &str {
"cozo:mem"
}
async fn graph_names(&self) -> Result<Vec<String>, GraphError> {
self.names()
}
/// Cozo has no catalogue procedures, so the schema is derived from the rows the
/// graph actually holds. Property keys stay empty — honest, not a failure.
///
/// **Lists before it probes**, exactly like the FalkorDB backend: a name that is
/// not in [`graph_names`](GraphDb::graph_names) is [`GraphError::NotFound`], never
/// an empty schema, and no relation row is written by the check.
async fn open(&self, graph: &str) -> Result<GraphSchema, GraphError> {
if !self.names()?.iter().any(|n| n == graph) {
return Err(GraphError::NotFound {
graph: graph.to_string(),
endpoint: self.endpoint().to_string(),
});
}
// Node/edge counts are the WHOLE graph, not a capped read — a schema that
// reported the cap instead of the size would silently lie about big graphs.
let g = escape(graph);
let labels: BTreeSet<String> = self
.run_rows(&format!("?[label] := *node{{g: '{g}', label}}"))?
.into_iter()
.filter_map(|r| r.first().and_then(|v| v.get_str()).map(str::to_string))
.filter(|s| !s.is_empty())
.collect();
let rels: BTreeSet<String> = self
.run_rows(&format!("?[rel] := *edge{{g: '{g}', rel}}"))?
.into_iter()
.filter_map(|r| r.first().and_then(|v| v.get_str()).map(str::to_string))
.filter(|s| !s.is_empty())
.collect();
let node_count = self
.run_rows(&format!("?[id] := *node{{g: '{g}', id}}"))?
.len() as u64;
let edge_count = self
.run_rows(&format!("?[src, dst, rel] := *edge{{g: '{g}', src, dst, rel}}"))?
.len() as u64;
Ok(GraphSchema {
graph: graph.to_string(),
labels: labels.into_iter().collect(),
relationship_types: rels.into_iter().collect(),
property_keys: Vec::new(),
node_count,
edge_count,
}
.normalised())
}
async fn read(&self, graph: &str, limit: usize) -> Result<SubGraph, GraphError> {
let g = escape(graph);
let limit = cap(limit);
let script = format!(
"?[src, dst, rel, sl, dl, sk, dk] := \
*edge{{g: '{g}', src, dst, rel}}, \
*node{{g: '{g}', id: src, label: sl, nk: sk}}, \
*node{{g: '{g}', id: dst, label: dl, nk: dk}}\n:limit {limit}"
);
self.run_scene(&script)
}
async fn neighbours(
&self,
graph: &str,
node_id: i64,
limit: usize,
) -> Result<SubGraph, GraphError> {
let g = escape(graph);
let limit = cap(limit);
let script = format!(
"nb[src, dst, rel] := *edge{{g: '{g}', src, dst, rel}}, src == {node_id}\n\
nb[src, dst, rel] := *edge{{g: '{g}', src, dst, rel}}, dst == {node_id}\n\
?[src, dst, rel, sl, dl, sk, dk] := \
nb[src, dst, rel], \
*node{{g: '{g}', id: src, label: sl, nk: sk}}, \
*node{{g: '{g}', id: dst, label: dl, nk: dk}}\n:limit {limit}"
);
self.run_scene(&script)
}
async fn shortest_path(
&self,
_graph: &str,
_from_id: i64,
_to_id: i64,
_max_hops: usize,
) -> Result<SubGraph, GraphError> {
// Deferred op — a clear, non-panicking signal (not todo!()).
Err(GraphError::Unsupported {
backend: "cozo",
op: "shortest_path",
})
}
async fn upsert(
&self,
graph: &str,
nodes: &[UpsertNode],
edges: &[UpsertEdge],
) -> Result<u64, GraphError> {
self.upsert_sync(graph, nodes, edges)
}
}
/// Escape a single-quoted CozoScript string literal.
fn escape(s: &str) -> String {
s.replace('\\', "\\\\").replace('\'', "\\'")
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{CypherSupport, no_cypher_here};
/// The container-free round-trip: create nodes + edges, read them back, expand
/// neighbours — proving the embedded Cozo backend works with no container.
#[tokio::test]
async fn cozo_roundtrip_upsert_read_neighbours() {
let be = CozoBackend::new().unwrap();
let nodes = [
UpsertNode::new(1, "Person"),
UpsertNode::new(2, "Company").with_key("5561234567"),
UpsertNode::new(3, "Vessel"),
];
let edges = [
UpsertEdge::new(1, 2, "WORKS_AT"),
UpsertEdge::new(1, 3, "OWNS"),
];
let n = be.upsert("g1", &nodes, &edges).await.unwrap();
assert_eq!(n, 5, "3 nodes + 2 edges written");
let scene = be.read("g1", 100).await.unwrap();
assert_eq!(scene.edges.len(), 2, "two edges in the scene");
assert_eq!(scene.nodes.len(), 3, "three distinct nodes");
let company = scene.nodes.iter().find(|n| n.label == "Company").unwrap();
assert_eq!(company.natural_key.as_deref(), Some("5561234567"));
let nb = be.neighbours("g1", 1, 100).await.unwrap();
assert_eq!(nb.edges.len(), 2);
let nb2 = be.neighbours("g1", 2, 100).await.unwrap();
assert_eq!(nb2.edges.len(), 1);
}
/// **The fold's Cozo half: `graph_names` + `open` really work on an embedded
/// store, and they are SCOPED per graph.**
///
/// RED when broken: drop the `g` key from the relations and the two graphs bleed
/// into each other — `read("a")` would return b's edge and `open("a").node_count`
/// would be 4, not 2.
#[tokio::test]
async fn cozo_lists_opens_and_scopes_two_graphs() {
let be = CozoBackend::new().unwrap();
be.upsert(
"alpha",
&[UpsertNode::new(1, "Person"), UpsertNode::new(2, "Company")],
&[UpsertEdge::new(1, 2, "WORKS_AT")],
)
.await
.unwrap();
be.upsert(
"beta",
&[UpsertNode::new(9, "Vessel"), UpsertNode::new(8, "Person")],
&[UpsertEdge::new(8, 9, "OWNS")],
)
.await
.unwrap();
assert_eq!(be.graph_names().await.unwrap(), vec!["alpha", "beta"]);
let a = be.open("alpha").await.unwrap();
assert_eq!(a.labels, vec!["Company".to_string(), "Person".to_string()]);
assert_eq!(a.relationship_types, vec!["WORKS_AT".to_string()]);
assert_eq!((a.node_count, a.edge_count), (2, 1));
let b = be.open("beta").await.unwrap();
assert_eq!(b.relationship_types, vec!["OWNS".to_string()]);
assert_eq!((b.node_count, b.edge_count), (2, 1));
// The scenes do not bleed.
let sa = be.read("alpha", 0).await.unwrap();
assert_eq!(sa.edges.len(), 1);
assert_eq!(sa.edges[0].rel, "WORKS_AT");
}
/// A graph that is not there is `NotFound`, never an empty schema — and the
/// check does not CREATE it (the FalkorDB bug's Cozo twin).
#[tokio::test]
async fn an_absent_graph_is_not_found_and_is_not_created() {
let be = CozoBackend::new().unwrap();
be.upsert("real", &[UpsertNode::new(1, "Person")], &[])
.await
.unwrap();
let e = be.open("ghost").await.unwrap_err();
assert!(matches!(e, GraphError::NotFound { .. }), "{e}");
assert_eq!(
be.graph_names().await.unwrap(),
vec!["real"],
"the failed open left no stray graph behind"
);
}
/// Cozo is a COMPLETE backend with no query language: it does not implement
/// `Cypher` at all, so there is no stub to panic from.
#[tokio::test]
async fn cozo_offers_no_cypher_and_says_so() {
let be = CozoBackend::new().unwrap();
assert_eq!(be.cypher_support(), CypherSupport::None);
let refusal = no_cypher_here(&be);
assert!(matches!(
refusal,
GraphError::NoQueryLanguage { backend: "cozo", .. }
));
assert!(refusal.to_string().contains("cozo:mem"));
}
#[tokio::test]
async fn cozo_shortest_path_is_unsupported_not_panic() {
let be = CozoBackend::new().unwrap();
let err = be.shortest_path("g", 1, 2, 3).await.unwrap_err();
assert!(matches!(
err,
GraphError::Unsupported {
backend: "cozo",
..
}
));
}
#[tokio::test]
async fn cozo_seed_demo_renders_a_scene() {
let be = CozoBackend::new().unwrap();
be.seed_demo().unwrap();
assert_eq!(be.graph_names().await.unwrap(), vec![DEFAULT_GRAPH]);
let scene = be.read(DEFAULT_GRAPH, 1000).await.unwrap();
assert!(!scene.nodes.is_empty(), "seeded scene is non-empty");
assert!(!scene.edges.is_empty());
}
}