use std::collections::{HashMap, HashSet, VecDeque};
use std::sync::Arc;
use async_trait::async_trait;
use chrono::{DateTime, TimeZone, Utc};
use rusqlite::OptionalExtension;
use uuid::Uuid;
use khive_storage::error::StorageError;
use khive_storage::types::{
BatchWriteSummary, DeleteMode, DirectedNeighborHit, Direction, Edge, EdgeFilter, EdgeSeekPage,
EdgeSortField, GraphPath, GuardedBatchOutcome, GuardedBatchRefusal, GuardedWriteOutcome,
MissingEndpoints, NeighborHit, NeighborQuery, Page, PageRequest, PathNode, SeekCursor,
SeekPage, SortDirection, SortOrder, SqlStatement, SqlValue, TraversalExecutionBudget,
TraversalOptions, TraversalRequest,
};
use khive_storage::GraphStore;
use khive_storage::LinkId;
use khive_storage::StorageCapability;
use khive_types::EdgeRelation;
use crate::error::SqliteError;
use crate::pool::ConnectionPool;
use crate::sql_bridge::bind_params;
use crate::writer_task::WriterTaskHandle;
fn map_err(e: rusqlite::Error, op: &'static str) -> StorageError {
StorageError::driver(StorageCapability::Graph, op, e)
}
fn map_sqlite_err(e: SqliteError, op: &'static str) -> StorageError {
StorageError::driver(StorageCapability::Graph, op, e)
}
const NAMESPACE_COUNT_CHUNK_SIZE: usize = 500;
const EDGE_NATURAL_KEY_CONFLICT_SET: &str = "weight = excluded.weight, \
updated_at = excluded.updated_at, \
deleted_at = NULL, \
metadata = excluded.metadata, \
target_backend = excluded.target_backend";
fn endpoint_exists_clause(id_param: &str) -> String {
format!(
"EXISTS (SELECT 1 FROM entities WHERE id = {id_param} AND deleted_at IS NULL) \
OR EXISTS (SELECT 1 FROM notes WHERE id = {id_param} AND deleted_at IS NULL) \
OR EXISTS (SELECT 1 FROM events WHERE id = {id_param}) \
OR EXISTS (SELECT 1 FROM graph_edges WHERE id = {id_param} AND deleted_at IS NULL)"
)
}
pub fn edge_upsert_statement(edge: &Edge) -> SqlStatement {
let (source_id, target_id) =
canonical_edge_endpoints(edge.relation, edge.source_id, edge.target_id);
let metadata_str = edge
.metadata
.as_ref()
.map(|v| serde_json::to_string(v).unwrap_or_default());
SqlStatement {
sql: format!(
"INSERT INTO graph_edges \
(namespace, id, source_id, target_id, relation, weight, \
created_at, updated_at, deleted_at, metadata, target_backend) \
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11) \
ON CONFLICT(namespace, id) DO UPDATE SET \
source_id = excluded.source_id, \
target_id = excluded.target_id, \
relation = excluded.relation, \
{EDGE_NATURAL_KEY_CONFLICT_SET} \
ON CONFLICT(namespace, source_id, target_id, relation) DO UPDATE SET \
{EDGE_NATURAL_KEY_CONFLICT_SET}"
),
params: vec![
SqlValue::Text(edge.namespace.clone()),
SqlValue::Text(Uuid::from(edge.id).to_string()),
SqlValue::Text(source_id.to_string()),
SqlValue::Text(target_id.to_string()),
SqlValue::Text(edge.relation.to_string()),
SqlValue::Float(edge.weight),
SqlValue::Integer(edge.created_at.timestamp_micros()),
SqlValue::Integer(edge.updated_at.timestamp_micros()),
match edge.deleted_at {
Some(t) => SqlValue::Integer(t.timestamp_micros()),
None => SqlValue::Null,
},
match metadata_str {
Some(m) => SqlValue::Text(m),
None => SqlValue::Null,
},
match &edge.target_backend {
Some(b) => SqlValue::Text(b.clone()),
None => SqlValue::Null,
},
],
label: Some("edge-upsert".to_string()),
}
}
pub fn edge_replace_if_unchanged_statement(
edge: &Edge,
expected_updated_at: DateTime<Utc>,
expected_deleted_at: Option<DateTime<Utc>>,
) -> SqlStatement {
let (source_id, target_id) =
canonical_edge_endpoints(edge.relation, edge.source_id, edge.target_id);
let metadata_str = edge
.metadata
.as_ref()
.map(|v| serde_json::to_string(v).unwrap_or_default());
SqlStatement {
sql: "UPDATE graph_edges SET \
namespace = ?1, source_id = ?2, target_id = ?3, relation = ?4, weight = ?5, \
updated_at = ?6, deleted_at = ?7, metadata = ?8, target_backend = ?9 \
WHERE id = ?10 AND updated_at = ?11 AND deleted_at IS ?12 \
AND ?6 > updated_at"
.to_string(),
params: vec![
SqlValue::Text(edge.namespace.clone()),
SqlValue::Text(source_id.to_string()),
SqlValue::Text(target_id.to_string()),
SqlValue::Text(edge.relation.to_string()),
SqlValue::Float(edge.weight),
SqlValue::Integer(edge.updated_at.timestamp_micros()),
match edge.deleted_at {
Some(t) => SqlValue::Integer(t.timestamp_micros()),
None => SqlValue::Null,
},
match metadata_str {
Some(m) => SqlValue::Text(m),
None => SqlValue::Null,
},
match &edge.target_backend {
Some(b) => SqlValue::Text(b.clone()),
None => SqlValue::Null,
},
SqlValue::Text(Uuid::from(edge.id).to_string()),
SqlValue::Integer(expected_updated_at.timestamp_micros()),
match expected_deleted_at {
Some(value) => SqlValue::Integer(value.timestamp_micros()),
None => SqlValue::Null,
},
],
label: Some("edge-replace-if-unchanged".to_string()),
}
}
#[allow(clippy::too_many_arguments)]
pub fn edge_insert_guarded_by_endpoints_statement(
namespace: &str,
edge_id: Uuid,
source_id: Uuid,
target_id: Uuid,
relation: EdgeRelation,
weight: f64,
now: i64,
metadata: Option<&str>,
) -> SqlStatement {
let src_exists = endpoint_exists_clause("?3");
let tgt_exists = endpoint_exists_clause("?4");
SqlStatement {
sql: format!(
"INSERT INTO graph_edges \
(namespace, id, source_id, target_id, relation, weight, \
created_at, updated_at, metadata) \
SELECT ?1, ?2, ?3, ?4, ?5, ?6, ?7, ?7, ?8 \
WHERE ({src_exists}) AND ({tgt_exists}) \
ON CONFLICT(namespace, source_id, target_id, relation) DO UPDATE SET \
{EDGE_NATURAL_KEY_CONFLICT_SET}"
),
params: vec![
SqlValue::Text(namespace.to_string()),
SqlValue::Text(edge_id.to_string()),
SqlValue::Text(source_id.to_string()),
SqlValue::Text(target_id.to_string()),
SqlValue::Text(relation.as_str().to_string()),
SqlValue::Float(weight),
SqlValue::Integer(now),
match metadata {
Some(m) => SqlValue::Text(m.to_string()),
None => SqlValue::Null,
},
],
label: Some("atomic-link-insert-edge-where-exists".to_string()),
}
}
pub fn edge_soft_delete_statement(id: Uuid, now: i64) -> SqlStatement {
SqlStatement {
sql: "UPDATE graph_edges SET deleted_at = ?2, updated_at = ?2 \
WHERE id = ?1 AND deleted_at IS NULL"
.to_string(),
params: vec![SqlValue::Text(id.to_string()), SqlValue::Integer(now)],
label: Some("edge-delete-soft".to_string()),
}
}
pub fn edge_hard_delete_statement(id: Uuid) -> SqlStatement {
SqlStatement {
sql: "DELETE FROM graph_edges WHERE id = ?1".to_string(),
params: vec![SqlValue::Text(id.to_string())],
label: Some("edge-delete-hard".to_string()),
}
}
pub fn purge_incident_edges_statement(node_id: Uuid) -> SqlStatement {
SqlStatement {
sql: "DELETE FROM graph_edges WHERE source_id = ?1 OR target_id = ?1".to_string(),
params: vec![SqlValue::Text(node_id.to_string())],
label: Some("edge-purge-incident".to_string()),
}
}
pub const EDGE_SYMMETRIC_CONFLICT_PROBE_SQL: &str = "SELECT id FROM graph_edges \
WHERE namespace = ?1 AND source_id = ?2 AND target_id = ?3 \
AND relation = ?4 AND id != ?5";
pub const EDGE_SYMMETRIC_DELETE_NONCANONICAL_SQL: &str =
"DELETE FROM graph_edges WHERE namespace = ?1 AND id = ?2";
pub const EDGE_SYMMETRIC_DELETE_NONCANONICAL_GUARDED_SQL: &str =
"DELETE FROM graph_edges WHERE namespace = ?1 AND id = ?2 \
AND updated_at = ?3 AND deleted_at IS ?4";
pub const EDGE_SYMMETRIC_UPDATE_INPLACE_SQL: &str = "UPDATE graph_edges SET \
source_id = ?1, target_id = ?2, relation = ?3, \
weight = ?4, updated_at = ?5, metadata = ?6 \
WHERE namespace = ?7 AND id = ?8 \
AND updated_at = ?9 AND deleted_at IS ?10 \
AND ?5 > updated_at";
pub fn edge_symmetric_conflict_probe_statement(
namespace: &str,
canon_src: Uuid,
canon_tgt: Uuid,
relation: EdgeRelation,
exclude_id: Uuid,
) -> SqlStatement {
SqlStatement {
sql: EDGE_SYMMETRIC_CONFLICT_PROBE_SQL.to_string(),
params: vec![
SqlValue::Text(namespace.to_string()),
SqlValue::Text(canon_src.to_string()),
SqlValue::Text(canon_tgt.to_string()),
SqlValue::Text(relation.to_string()),
SqlValue::Text(exclude_id.to_string()),
],
label: Some("edge-symmetric-conflict-probe".to_string()),
}
}
pub fn edge_symmetric_delete_noncanonical_statement(namespace: &str, id: Uuid) -> SqlStatement {
SqlStatement {
sql: EDGE_SYMMETRIC_DELETE_NONCANONICAL_SQL.to_string(),
params: vec![
SqlValue::Text(namespace.to_string()),
SqlValue::Text(id.to_string()),
],
label: Some("edge-symmetric-delete-noncanonical".to_string()),
}
}
#[allow(clippy::too_many_arguments)]
pub fn edge_symmetric_update_inplace_statement(
namespace: &str,
id: Uuid,
canon_src: Uuid,
canon_tgt: Uuid,
relation: EdgeRelation,
weight: f64,
updated_at_micros: i64,
metadata: Option<&str>,
expected_updated_at_micros: i64,
expected_deleted_at_micros: Option<i64>,
) -> SqlStatement {
SqlStatement {
sql: EDGE_SYMMETRIC_UPDATE_INPLACE_SQL.to_string(),
params: vec![
SqlValue::Text(canon_src.to_string()),
SqlValue::Text(canon_tgt.to_string()),
SqlValue::Text(relation.to_string()),
SqlValue::Float(weight),
SqlValue::Integer(updated_at_micros),
match metadata {
Some(m) => SqlValue::Text(m.to_string()),
None => SqlValue::Null,
},
SqlValue::Text(namespace.to_string()),
SqlValue::Text(id.to_string()),
SqlValue::Integer(expected_updated_at_micros),
match expected_deleted_at_micros {
Some(value) => SqlValue::Integer(value),
None => SqlValue::Null,
},
],
label: Some("edge-symmetric-update-inplace".to_string()),
}
}
#[allow(clippy::too_many_arguments)]
pub fn edge_symmetric_delete_if_conflict_statement(
namespace: &str,
id: Uuid,
canon_src: Uuid,
canon_tgt: Uuid,
relation: EdgeRelation,
expected_updated_at_micros: i64,
expected_deleted_at_micros: Option<i64>,
) -> SqlStatement {
SqlStatement {
sql: "DELETE FROM graph_edges \
WHERE namespace = ?1 AND id = ?2 \
AND updated_at = ?6 AND deleted_at IS ?7 \
AND EXISTS ( \
SELECT 1 FROM graph_edges \
WHERE namespace = ?1 AND source_id = ?3 AND target_id = ?4 \
AND relation = ?5 AND id != ?2 \
)"
.to_string(),
params: vec![
SqlValue::Text(namespace.to_string()),
SqlValue::Text(id.to_string()),
SqlValue::Text(canon_src.to_string()),
SqlValue::Text(canon_tgt.to_string()),
SqlValue::Text(relation.to_string()),
SqlValue::Integer(expected_updated_at_micros),
match expected_deleted_at_micros {
Some(value) => SqlValue::Integer(value),
None => SqlValue::Null,
},
],
label: Some("edge-symmetric-delete-if-conflict".to_string()),
}
}
#[allow(clippy::too_many_arguments)]
pub fn edge_symmetric_absorb_or_update_inplace_statement(
namespace: &str,
id: Uuid,
canon_src: Uuid,
canon_tgt: Uuid,
relation: EdgeRelation,
weight: f64,
updated_at_micros: i64,
metadata: Option<&str>,
target_backend: Option<&str>,
expected_updated_at_micros: i64,
expected_deleted_at_micros: Option<i64>,
) -> SqlStatement {
SqlStatement {
sql: "UPDATE graph_edges SET \
source_id = CASE WHEN id = ?2 THEN ?3 ELSE source_id END, \
target_id = CASE WHEN id = ?2 THEN ?4 ELSE target_id END, \
relation = CASE WHEN id = ?2 THEN ?5 ELSE relation END, \
weight = CASE WHEN id = ?2 THEN ?6 ELSE weight END, \
updated_at = CASE WHEN id = ?2 THEN ?7 ELSE updated_at END, \
deleted_at = CASE WHEN id = ?2 THEN NULL ELSE deleted_at END, \
metadata = CASE WHEN id = ?2 THEN ?8 ELSE metadata END, \
target_backend = CASE WHEN id = ?2 THEN ?9 ELSE target_backend END \
WHERE namespace = ?1 \
AND ( \
(id = ?2 AND changes() = 0 AND updated_at = ?10 AND deleted_at IS ?11 \
AND ?7 > updated_at) \
OR (source_id = ?3 AND target_id = ?4 AND relation = ?5 \
AND id != ?2 AND changes() = 1) \
)"
.to_string(),
params: vec![
SqlValue::Text(namespace.to_string()),
SqlValue::Text(id.to_string()),
SqlValue::Text(canon_src.to_string()),
SqlValue::Text(canon_tgt.to_string()),
SqlValue::Text(relation.to_string()),
SqlValue::Float(weight),
SqlValue::Integer(updated_at_micros),
match metadata {
Some(m) => SqlValue::Text(m.to_string()),
None => SqlValue::Null,
},
match target_backend {
Some(b) => SqlValue::Text(b.to_string()),
None => SqlValue::Null,
},
SqlValue::Integer(expected_updated_at_micros),
match expected_deleted_at_micros {
Some(value) => SqlValue::Integer(value),
None => SqlValue::Null,
},
],
label: Some("edge-symmetric-absorb-or-update-inplace".to_string()),
}
}
pub struct SqlGraphStore {
pool: Arc<ConnectionPool>,
is_file_backed: bool,
namespace: String,
writer_task: Option<WriterTaskHandle>,
}
impl SqlGraphStore {
pub fn new_scoped(
pool: Arc<ConnectionPool>,
is_file_backed: bool,
namespace: impl Into<String>,
) -> Self {
let writer_task = pool.writer_task_handle().ok().flatten();
Self {
pool,
is_file_backed,
namespace: namespace.into(),
writer_task,
}
}
fn open_standalone_writer(&self) -> Result<rusqlite::Connection, StorageError> {
self.pool
.open_standalone_writer()
.map_err(|e| map_sqlite_err(e, "open_graph_writer"))
}
fn current_writer_task(
&self,
operation: &'static str,
) -> Result<Option<WriterTaskHandle>, StorageError> {
self.pool
.writer_task_for_write(self.writer_task.as_ref(), operation)
}
async fn with_writer<F, R>(&self, op: &'static str, f: F) -> Result<R, StorageError>
where
F: FnOnce(&rusqlite::Connection) -> Result<R, rusqlite::Error> + Send + 'static,
R: Send + 'static,
{
if let Some(writer_task) = self.current_writer_task(op)? {
return writer_task
.send_bounded(move |conn| f(conn).map_err(|e| map_err(e, op)))
.await;
}
self.pool
.record_direct_route(crate::timeout_sink::Site::DirectRouteGraphGeneralWrite);
if self.is_file_backed {
let conn = self.open_standalone_writer()?;
let db = crate::timeout_sink::db_label(&self.pool);
tokio::task::spawn_blocking(move || {
f(&conn).map_err(|e| {
crate::timeout_sink::maybe_emit_busy(
&db,
crate::timeout_sink::Site::StandaloneGraph,
&e,
);
map_err(e, op)
})
})
.await
.map_err(|e| StorageError::driver(StorageCapability::Graph, op, e))?
} else {
let pool = Arc::clone(&self.pool);
tokio::task::spawn_blocking(move || {
let guard = pool.try_writer().map_err(|e| map_sqlite_err(e, op))?;
f(guard.conn()).map_err(|e| map_err(e, op))
})
.await
.map_err(|e| StorageError::driver(StorageCapability::Graph, op, e))?
}
}
async fn with_reader<F, R>(&self, op: &'static str, f: F) -> Result<R, StorageError>
where
F: FnOnce(&rusqlite::Connection) -> Result<R, rusqlite::Error> + Send + 'static,
R: Send + 'static,
{
if self.is_file_backed {
let pool = Arc::clone(&self.pool);
crate::read_cancellation::run_declared_interruptible_read(
StorageCapability::Graph,
op,
move |scope| {
scope.ensure_active()?;
let conn = pool
.open_standalone_reader()
.map_err(|error| map_sqlite_err(error, op))?;
scope.run(&conn, || f(&conn).map_err(|e| map_err(e, op)))
},
)
.await
} else {
let pool = Arc::clone(&self.pool);
crate::read_cancellation::run_declared_interruptible_read(
StorageCapability::Graph,
op,
move |scope| {
let mut guard = pool.resolve_reader_checkout(
StorageCapability::Graph,
op,
pool.reader_until(|| scope.should_stop()),
)?;
scope.run_pooled_reader(&mut guard, |conn| f(conn).map_err(|e| map_err(e, op)))
},
)
.await
}
}
}
fn report_graph_usage(queries: &std::sync::atomic::AtomicU64, rows: &std::sync::atomic::AtomicU64) {
khive_storage::usage::count(
khive_storage::usage::UsageUnit::DbRoundTrips,
queries.load(std::sync::atomic::Ordering::Relaxed),
);
khive_storage::usage::count(
khive_storage::usage::UsageUnit::GraphHops,
rows.load(std::sync::atomic::Ordering::Relaxed),
);
}
fn read_edge(row: &rusqlite::Row<'_>) -> Result<Edge, rusqlite::Error> {
let namespace: String = row.get(0)?;
let id_str: String = row.get(1)?;
let source_str: String = row.get(2)?;
let target_str: String = row.get(3)?;
let relation_str: String = row.get(4)?;
let weight: f64 = row.get(5)?;
let created_micros: i64 = row.get(6)?;
let updated_micros: i64 = row.get(7)?;
let deleted_micros: Option<i64> = row.get(8)?;
let metadata_str: Option<String> = row.get(9)?;
let target_backend: Option<String> = row.get(10)?;
let id = parse_uuid(&id_str)?;
let source_id = parse_uuid(&source_str)?;
let target_id = parse_uuid(&target_str)?;
let created_at = micros_to_datetime(created_micros);
let relation = relation_str.parse::<EdgeRelation>().map_err(|e| {
rusqlite::Error::FromSqlConversionFailure(4, rusqlite::types::Type::Text, Box::new(e))
})?;
let metadata = match metadata_str {
Some(s) => {
let v = serde_json::from_str(&s).map_err(|e| {
rusqlite::Error::FromSqlConversionFailure(
9,
rusqlite::types::Type::Text,
Box::new(e),
)
})?;
Some(v)
}
None => None,
};
Ok(Edge {
id: id.into(),
namespace,
source_id,
target_id,
relation,
weight,
created_at,
updated_at: micros_to_datetime(updated_micros),
deleted_at: deleted_micros.map(micros_to_datetime),
metadata,
target_backend,
})
}
fn parse_uuid(s: &str) -> Result<Uuid, rusqlite::Error> {
Uuid::parse_str(s).map_err(|e| {
rusqlite::Error::FromSqlConversionFailure(0, rusqlite::types::Type::Text, Box::new(e))
})
}
fn neighbor_extra_clause(
query: &NeighborQuery,
start_param_idx: usize,
) -> (String, String, Vec<Box<dyn rusqlite::types::ToSql>>) {
let mut conditions: Vec<String> = Vec::new();
let mut extra_params: Vec<Box<dyn rusqlite::types::ToSql>> = Vec::new();
let mut param_idx = start_param_idx;
if let Some(ref rels) = query.relations {
if !rels.is_empty() {
let placeholders: Vec<String> = rels
.iter()
.map(|r| {
extra_params.push(Box::new(r.to_string()));
let p = format!("?{}", param_idx);
param_idx += 1;
p
})
.collect();
conditions.push(format!("relation IN ({})", placeholders.join(",")));
}
}
if let Some(min_w) = query.min_weight {
extra_params.push(Box::new(min_w));
conditions.push(format!("weight >= ?{}", param_idx));
param_idx += 1;
}
let where_extra = if conditions.is_empty() {
String::new()
} else {
format!(" WHERE {}", conditions.join(" AND "))
};
let limit_clause = if let Some(lim) = query.limit {
extra_params.push(Box::new(lim as i64));
format!(" LIMIT ?{}", param_idx)
} else {
String::new()
};
(where_extra, limit_clause, extra_params)
}
#[cfg(test)]
static NEIGHBOR_SELECT_COUNT: std::sync::atomic::AtomicUsize =
std::sync::atomic::AtomicUsize::new(0);
#[cfg(test)]
fn count_neighbor_select() {
NEIGHBOR_SELECT_COUNT.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
}
#[cfg(not(test))]
fn count_neighbor_select() {}
#[cfg(test)]
pub(crate) fn reset_neighbor_select_count() {
NEIGHBOR_SELECT_COUNT.store(0, std::sync::atomic::Ordering::Relaxed);
}
#[cfg(test)]
pub(crate) fn neighbor_select_count() -> usize {
NEIGHBOR_SELECT_COUNT.load(std::sync::atomic::Ordering::Relaxed)
}
fn micros_to_datetime(micros: i64) -> DateTime<Utc> {
Utc.timestamp_micros(micros)
.single()
.unwrap_or_else(Utc::now)
}
fn edge_order_clause(sort: &[SortOrder<EdgeSortField>]) -> String {
if sort.is_empty() {
return " ORDER BY created_at DESC, id DESC".to_string();
}
let mut parts: Vec<String> = sort
.iter()
.map(|s| {
let dir = match s.direction {
SortDirection::Asc => "ASC",
SortDirection::Desc => "DESC",
};
format!("{} {}", edge_sort_col(&s.field), dir)
})
.collect();
let dir = match sort.last().map(|s| &s.direction) {
Some(SortDirection::Asc) => "ASC",
_ => "DESC",
};
parts.push(format!("id {dir}"));
format!(" ORDER BY {}", parts.join(", "))
}
fn build_edge_filter_sql(
namespace: &str,
filter: &EdgeFilter,
) -> (String, Vec<Box<dyn rusqlite::types::ToSql>>) {
build_edge_filter_sql_for_namespaces(&[namespace.to_string()], filter)
}
fn build_edge_filter_sql_for_namespaces(
namespaces: &[String],
filter: &EdgeFilter,
) -> (String, Vec<Box<dyn rusqlite::types::ToSql>>) {
let params: Vec<Box<dyn rusqlite::types::ToSql>> = namespaces
.iter()
.map(|namespace| -> Box<dyn rusqlite::types::ToSql> { Box::new(namespace.clone()) })
.collect();
let namespace_condition = match namespaces.len() {
0 => "0".to_string(),
1 => "namespace = ?1".to_string(),
_ => {
let placeholders: Vec<String> =
(1..=namespaces.len()).map(|i| format!("?{i}")).collect();
format!("namespace IN ({})", placeholders.join(", "))
}
};
build_edge_filter_conditions(namespace_condition, params, filter)
}
fn build_edge_filter_sql_for_namespaces_json(
namespaces_json: &str,
filter: &EdgeFilter,
) -> (String, Vec<Box<dyn rusqlite::types::ToSql>>) {
let params: Vec<Box<dyn rusqlite::types::ToSql>> = vec![Box::new(namespaces_json.to_string())];
let namespace_condition = "namespace IN (SELECT value FROM json_each(?1))".to_string();
build_edge_filter_conditions(namespace_condition, params, filter)
}
fn build_edge_filter_conditions(
namespace_condition: String,
mut params: Vec<Box<dyn rusqlite::types::ToSql>>,
filter: &EdgeFilter,
) -> (String, Vec<Box<dyn rusqlite::types::ToSql>>) {
let mut conditions = vec![namespace_condition, "deleted_at IS NULL".to_string()];
if !filter.ids.is_empty() {
let placeholders: Vec<String> = filter
.ids
.iter()
.map(|id| {
params.push(Box::new(id.to_string()));
format!("?{}", params.len())
})
.collect();
conditions.push(format!("id IN ({})", placeholders.join(",")));
}
if !filter.source_ids.is_empty() {
let placeholders: Vec<String> = filter
.source_ids
.iter()
.map(|id| {
params.push(Box::new(id.to_string()));
format!("?{}", params.len())
})
.collect();
conditions.push(format!("source_id IN ({})", placeholders.join(",")));
}
if !filter.target_ids.is_empty() {
let placeholders: Vec<String> = filter
.target_ids
.iter()
.map(|id| {
params.push(Box::new(id.to_string()));
format!("?{}", params.len())
})
.collect();
conditions.push(format!("target_id IN ({})", placeholders.join(",")));
}
if !filter.relations.is_empty() {
let placeholders: Vec<String> = filter
.relations
.iter()
.map(|r| {
params.push(Box::new(r.to_string()));
format!("?{}", params.len())
})
.collect();
conditions.push(format!("relation IN ({})", placeholders.join(",")));
}
if let Some(min_w) = filter.min_weight {
params.push(Box::new(min_w));
conditions.push(format!("weight >= ?{}", params.len()));
}
if let Some(max_w) = filter.max_weight {
params.push(Box::new(max_w));
conditions.push(format!("weight <= ?{}", params.len()));
}
if let Some(ref time_range) = filter.created_at {
if let Some(start) = time_range.start {
params.push(Box::new(start.timestamp_micros()));
conditions.push(format!("created_at >= ?{}", params.len()));
}
if let Some(end) = time_range.end {
params.push(Box::new(end.timestamp_micros()));
conditions.push(format!("created_at < ?{}", params.len()));
}
}
let clause = format!(" WHERE {}", conditions.join(" AND "));
(clause, params)
}
fn edge_sort_col(field: &EdgeSortField) -> &'static str {
match field {
EdgeSortField::CreatedAt => "created_at",
EdgeSortField::Weight => "weight",
EdgeSortField::Relation => "relation",
}
}
fn canonical_edge_endpoints(
relation: EdgeRelation,
source_id: Uuid,
target_id: Uuid,
) -> (Uuid, Uuid) {
if relation.is_symmetric() && target_id < source_id {
(target_id, source_id)
} else {
(source_id, target_id)
}
}
fn batch_upsert_edges(
conn: &rusqlite::Connection,
edges: &[Edge],
attempted: u64,
) -> Result<BatchWriteSummary, rusqlite::Error> {
let mut affected = 0u64;
for edge in edges {
let statement = edge_upsert_statement(edge);
let mut stmt = conn.prepare(&statement.sql)?;
bind_params(&mut stmt, &statement.params)?;
stmt.raw_execute()?;
affected += 1;
}
Ok(BatchWriteSummary {
attempted,
affected,
failed: 0,
first_error: String::new(),
})
}
fn edge_endpoints_exist(
conn: &rusqlite::Connection,
source_id: Uuid,
target_id: Uuid,
) -> Result<MissingEndpoints, rusqlite::Error> {
let src_exists = endpoint_exists_clause("?1");
let tgt_exists = endpoint_exists_clause("?2");
let sql = format!("SELECT ({src_exists}), ({tgt_exists})");
conn.query_row(
&sql,
rusqlite::params![source_id.to_string(), target_id.to_string()],
|row| {
let src_exists: bool = row.get(0)?;
let tgt_exists: bool = row.get(1)?;
Ok(MissingEndpoints {
source: !src_exists,
target: !tgt_exists,
})
},
)
}
fn edge_insert_guarded(
conn: &rusqlite::Connection,
statement: &SqlStatement,
source_id: Uuid,
target_id: Uuid,
) -> Result<GuardedWriteOutcome, rusqlite::Error> {
let mut stmt = conn.prepare(&statement.sql)?;
bind_params(&mut stmt, &statement.params)?;
if stmt.raw_execute()? > 0 {
return Ok(GuardedWriteOutcome::Written);
}
#[cfg(test)]
tests::insert_probe_seam::hook((source_id, target_id));
let missing = edge_endpoints_exist(conn, source_id, target_id)?;
Ok(GuardedWriteOutcome::Refused(missing))
}
fn batch_upsert_edges_guarded(
conn: &rusqlite::Connection,
edges: &[Edge],
attempted: u64,
) -> Result<GuardedBatchOutcome, rusqlite::Error> {
for (index, edge) in edges.iter().enumerate() {
let (source_id, target_id) =
canonical_edge_endpoints(edge.relation, edge.source_id, edge.target_id);
let missing = edge_endpoints_exist(conn, source_id, target_id)?;
if missing.any() {
return Ok(GuardedBatchOutcome {
summary: BatchWriteSummary {
attempted,
affected: 0,
failed: attempted,
first_error: format!(
"batch entry {index}: edge endpoint no longer exists at write time: source {source_id} or target {target_id}"
),
},
refused: Some(GuardedBatchRefusal {
entry_index: index,
missing,
}),
});
}
}
let mut affected = 0u64;
for edge in edges {
let statement = edge_upsert_statement(edge);
let mut stmt = conn.prepare(&statement.sql)?;
bind_params(&mut stmt, &statement.params)?;
stmt.raw_execute()?;
affected += 1;
}
Ok(GuardedBatchOutcome {
summary: BatchWriteSummary {
attempted,
affected,
failed: 0,
first_error: String::new(),
},
refused: None,
})
}
fn traversal_neighbor_sql(
direction: Direction,
relation_count: usize,
has_min_weight: bool,
) -> String {
let (node_column, endpoint_column, index) = match direction {
Direction::Out => ("target_id", "source_id", "idx_graph_edges_ns_src_rel"),
Direction::In => ("source_id", "target_id", "idx_graph_edges_ns_tgt_rel"),
Direction::Both => unreachable!("Direction::Both is split into indexed Out/In seeks"),
};
let mut sql = format!(
"SELECT {node_column}, id, weight \
FROM graph_edges INDEXED BY {index} \
WHERE namespace = ?1 AND {endpoint_column} = ?2 AND deleted_at IS NULL"
);
if relation_count > 0 {
let placeholders = (0..relation_count)
.map(|offset| format!("?{}", 4 + offset))
.collect::<Vec<_>>()
.join(",");
sql.push_str(&format!(" AND relation IN ({placeholders})"));
}
if has_min_weight {
sql.push_str(&format!(" AND weight >= ?{}", 4 + relation_count));
}
sql.push_str(" LIMIT ?3");
sql
}
fn traversal_timeout_error(budget: &TraversalExecutionBudget) -> StorageError {
StorageError::Timeout {
operation: format!(
"traverse ({}ms execution budget)",
budget.max_duration().as_millis()
)
.into(),
}
}
fn traversal_work_error(budget: &TraversalExecutionBudget) -> StorageError {
StorageError::InvalidInput {
capability: StorageCapability::Graph,
operation: "traverse".into(),
message: format!(
"traversal work budget exceeded after {} adjacency rows; \
narrow roots, depth, relations, or result limit",
budget.work_limit()
),
}
}
#[derive(Clone, Copy)]
struct TraversalFrontierNode {
node_id: Uuid,
depth: usize,
total_weight: f64,
}
#[allow(clippy::too_many_arguments)]
fn run_bounded_traversal(
conn: &rusqlite::Connection,
roots: Vec<Uuid>,
opts: TraversalOptions,
include_roots: bool,
namespace: String,
origin: khive_storage::tx_registry::TxOrigin,
budget: TraversalExecutionBudget,
counted_rows: &std::sync::atomic::AtomicU64,
counted_queries: &std::sync::atomic::AtomicU64,
) -> Result<Vec<GraphPath>, StorageError> {
let progress_timed_out = Arc::new(std::sync::atomic::AtomicBool::new(false));
let callback_timed_out = Arc::clone(&progress_timed_out);
let callback_budget = budget.clone();
#[cfg(test)]
let progress_seam_root = roots.first().copied();
conn.progress_handler(
1_000,
Some(move || {
if crate::read_cancellation::current_read_should_interrupt() {
return true;
}
#[cfg(test)]
if tests::traverse_progress_seam::hook(progress_seam_root) {
callback_timed_out.store(true, std::sync::atomic::Ordering::Relaxed);
return true;
}
let expired = callback_budget.is_expired();
if expired {
callback_timed_out.store(true, std::sync::atomic::Ordering::Relaxed);
}
expired
}),
)
.map_err(|e| map_err(e, "traverse_progress_handler"))?;
let result = (|| {
let result_limit = opts.effective_limit() as usize;
let relation_count = opts.relations.as_ref().map_or(0, Vec::len);
let directions = match opts.direction {
Direction::Out => vec![Direction::Out],
Direction::In => vec![Direction::In],
Direction::Both => vec![Direction::Out, Direction::In],
};
let statements = directions
.into_iter()
.map(|direction| {
traversal_neighbor_sql(direction, relation_count, opts.min_weight.is_some())
})
.collect::<Vec<_>>();
let map_sql_error = |error| {
if progress_timed_out.load(std::sync::atomic::Ordering::Relaxed) {
traversal_timeout_error(&budget)
} else {
map_err(error, "traverse")
}
};
let mut all_paths = Vec::with_capacity(roots.len());
for root_id in roots {
let mut seen = HashSet::new();
seen.insert(root_id);
let mut frontier = VecDeque::from([TraversalFrontierNode {
node_id: root_id,
depth: 0,
total_weight: 0.0,
}]);
let mut nodes = Vec::with_capacity(result_limit + usize::from(include_roots));
if include_roots {
nodes.push(PathNode {
node_id: root_id,
via_edge: None,
depth: 0,
name: None,
kind: None,
properties: None,
weight: 0.0,
});
}
let mut non_root_count = 0usize;
'root_walk: while non_root_count < result_limit {
let Some(current) = frontier.pop_front() else {
break;
};
if current.depth >= opts.max_depth {
continue;
}
if budget.is_expired() {
return Err(traversal_timeout_error(&budget));
}
for sql in &statements {
let row_cap = budget.remaining_work().saturating_add(1);
let mut params: Vec<Box<dyn rusqlite::types::ToSql>> = vec![
Box::new(namespace.clone()),
Box::new(current.node_id.to_string()),
Box::new(row_cap as i64),
];
if let Some(relations) = &opts.relations {
params.extend(relations.iter().map(|relation| {
Box::new(relation.to_string()) as Box<dyn rusqlite::types::ToSql>
}));
}
if let Some(min_weight) = opts.min_weight {
params.push(Box::new(min_weight));
}
let param_refs = params
.iter()
.map(|param| param.as_ref())
.collect::<Vec<&dyn rusqlite::types::ToSql>>();
let _snapshot = khive_storage::tx_registry::register_scoped(
Some("graph_traverse_read".to_string()),
origin.clone(),
);
let mut stmt = conn.prepare(sql).map_err(&map_sql_error)?;
counted_queries.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
let mut rows = stmt.query(param_refs.as_slice()).map_err(&map_sql_error)?;
while let Some(row) = rows.next().map_err(&map_sql_error)? {
#[cfg(test)]
tests::traverse_snapshot_seam::hook(current.node_id);
counted_rows.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
if budget.is_expired() {
return Err(traversal_timeout_error(&budget));
}
if !budget.try_consume_row() {
return Err(traversal_work_error(&budget));
}
let node_str: String = row.get(0).map_err(&map_sql_error)?;
let edge_str: String = row.get(1).map_err(&map_sql_error)?;
let edge_weight: f64 = row.get(2).map_err(&map_sql_error)?;
let node_id = parse_uuid(&node_str).map_err(&map_sql_error)?;
if !seen.insert(node_id) {
continue;
}
let via_edge = parse_uuid(&edge_str).map_err(&map_sql_error)?;
let depth = current.depth + 1;
let total_weight = current.total_weight + edge_weight;
nodes.push(PathNode {
node_id,
via_edge: Some(via_edge),
depth,
name: None,
kind: None,
properties: None,
weight: total_weight,
});
non_root_count += 1;
if depth < opts.max_depth {
frontier.push_back(TraversalFrontierNode {
node_id,
depth,
total_weight,
});
}
if non_root_count == result_limit {
break 'root_walk;
}
}
}
}
if !nodes.is_empty() {
let total_weight = nodes.iter().map(|node| node.weight).fold(0.0_f64, f64::max);
all_paths.push(GraphPath {
root_id,
nodes,
total_weight,
});
}
}
Ok(all_paths)
})();
conn.progress_handler(0, None::<fn() -> bool>)
.map_err(|e| map_err(e, "traverse_progress_handler_clear"))?;
result
}
#[async_trait]
impl GraphStore for SqlGraphStore {
async fn upsert_edge(&self, edge: Edge) -> Result<(), StorageError> {
let statement = edge_upsert_statement(&edge);
self.with_writer("upsert_edge", move |conn| {
let mut stmt = conn.prepare(&statement.sql)?;
bind_params(&mut stmt, &statement.params)?;
stmt.raw_execute()?;
Ok(())
})
.await
}
async fn replace_edge_if_unchanged(
&self,
edge: Edge,
expected_updated_at: DateTime<Utc>,
expected_deleted_at: Option<DateTime<Utc>>,
) -> Result<bool, StorageError> {
let statement =
edge_replace_if_unchanged_statement(&edge, expected_updated_at, expected_deleted_at);
self.with_writer("replace_edge_if_unchanged", move |conn| {
let mut stmt = conn.prepare(&statement.sql)?;
bind_params(&mut stmt, &statement.params)?;
Ok(stmt.raw_execute()? > 0)
})
.await
}
async fn upsert_edges(&self, edges: Vec<Edge>) -> Result<BatchWriteSummary, StorageError> {
let attempted = edges.len() as u64;
if let Some(writer_task) = self.current_writer_task("upsert_edges")? {
return writer_task
.send_bounded(move |conn| {
batch_upsert_edges(conn, &edges, attempted)
.map_err(|e| map_err(e, "upsert_edges"))
})
.await;
}
let origin = self.pool.origin();
self.with_writer("upsert_edges", move |conn| {
conn.execute_batch("BEGIN IMMEDIATE")?;
let _tx_handle = khive_storage::tx_registry::register_scoped(
Some("graph_upsert_edges".to_string()),
origin,
);
let summary = match batch_upsert_edges(conn, &edges, attempted) {
Ok(summary) => summary,
Err(e) => {
let _ = conn.execute_batch("ROLLBACK");
return Err(e);
}
};
if let Err(e) = conn.execute_batch("COMMIT") {
let _ = conn.execute_batch("ROLLBACK");
return Err(e);
}
Ok(summary)
})
.await
}
async fn upsert_edge_guarded(&self, edge: Edge) -> Result<GuardedWriteOutcome, StorageError> {
let (source_id, target_id) =
canonical_edge_endpoints(edge.relation, edge.source_id, edge.target_id);
let metadata_str = edge
.metadata
.as_ref()
.map(|v| serde_json::to_string(v).unwrap_or_default());
let statement = edge_insert_guarded_by_endpoints_statement(
&edge.namespace,
Uuid::from(edge.id),
source_id,
target_id,
edge.relation,
edge.weight,
edge.created_at.timestamp_micros(),
metadata_str.as_deref(),
);
if let Some(writer_task) = self.current_writer_task("upsert_edge_guarded")? {
return writer_task
.send_bounded(move |conn| {
edge_insert_guarded(conn, &statement, source_id, target_id)
.map_err(|e| map_err(e, "upsert_edge_guarded"))
})
.await;
}
let origin = self.pool.origin();
self.with_writer("upsert_edge_guarded", move |conn| {
conn.execute_batch("BEGIN IMMEDIATE")?;
let _tx_handle = khive_storage::tx_registry::register_scoped(
Some("graph_upsert_edge_guarded".to_string()),
origin,
);
let outcome = match edge_insert_guarded(conn, &statement, source_id, target_id) {
Ok(outcome) => outcome,
Err(e) => {
let _ = conn.execute_batch("ROLLBACK");
return Err(e);
}
};
if let Err(e) = conn.execute_batch("COMMIT") {
let _ = conn.execute_batch("ROLLBACK");
return Err(e);
}
Ok(outcome)
})
.await
}
async fn upsert_edges_guarded(
&self,
edges: Vec<Edge>,
) -> Result<GuardedBatchOutcome, StorageError> {
let attempted = edges.len() as u64;
if let Some(writer_task) = self.current_writer_task("upsert_edges_guarded")? {
return writer_task
.send_bounded(move |conn| {
batch_upsert_edges_guarded(conn, &edges, attempted)
.map_err(|e| map_err(e, "upsert_edges_guarded"))
})
.await;
}
let origin = self.pool.origin();
self.with_writer("upsert_edges_guarded", move |conn| {
conn.execute_batch("BEGIN IMMEDIATE")?;
let _tx_handle = khive_storage::tx_registry::register_scoped(
Some("graph_upsert_edges_guarded".to_string()),
origin,
);
let summary = match batch_upsert_edges_guarded(conn, &edges, attempted) {
Ok(summary) => summary,
Err(e) => {
let _ = conn.execute_batch("ROLLBACK");
return Err(e);
}
};
if let Err(e) = conn.execute_batch("COMMIT") {
let _ = conn.execute_batch("ROLLBACK");
return Err(e);
}
Ok(summary)
})
.await
}
async fn get_edge(&self, id: LinkId) -> Result<Option<Edge>, StorageError> {
let id_str = Uuid::from(id).to_string();
self.with_reader("get_edge", move |conn| {
let mut stmt = conn.prepare(
"SELECT namespace, id, source_id, target_id, relation, weight, \
created_at, updated_at, deleted_at, metadata, target_backend \
FROM graph_edges WHERE id = ?1 AND deleted_at IS NULL",
)?;
let mut rows = stmt.query(rusqlite::params![id_str])?;
match rows.next()? {
Some(row) => Ok(Some(read_edge(row)?)),
None => Ok(None),
}
})
.await
}
async fn get_edge_including_deleted(&self, id: LinkId) -> Result<Option<Edge>, StorageError> {
let id_str = Uuid::from(id).to_string();
self.with_reader("get_edge_including_deleted", move |conn| {
let mut stmt = conn.prepare(
"SELECT namespace, id, source_id, target_id, relation, weight, \
created_at, updated_at, deleted_at, metadata, target_backend \
FROM graph_edges WHERE id = ?1",
)?;
let mut rows = stmt.query(rusqlite::params![id_str])?;
match rows.next()? {
Some(row) => Ok(Some(read_edge(row)?)),
None => Ok(None),
}
})
.await
}
async fn edge_sequence(&self, id: Uuid) -> Result<Option<i64>, StorageError> {
let id = id.to_string();
self.with_reader("edge_sequence", move |conn| {
conn.query_row(
"SELECT seq FROM graph_edges_seq WHERE edge_id = ?1",
rusqlite::params![id],
|row| row.get(0),
)
.optional()
})
.await
}
async fn edge_sequences(&self, ids: &[Uuid]) -> Result<Vec<(Uuid, i64)>, StorageError> {
if ids.is_empty() {
return Ok(Vec::new());
}
let ids = ids.to_vec();
self.with_reader("edge_sequences", move |conn| {
const CHUNK: usize = 900;
let mut resolved = Vec::with_capacity(ids.len());
for chunk in ids.chunks(CHUNK) {
let placeholders = (1..=chunk.len())
.map(|index| format!("?{index}"))
.collect::<Vec<_>>()
.join(", ");
let sql = format!(
"SELECT edge_id, seq FROM graph_edges_seq WHERE edge_id IN ({placeholders})"
);
let strings = chunk.iter().map(Uuid::to_string).collect::<Vec<_>>();
let params = strings
.iter()
.map(|id| id as &dyn rusqlite::types::ToSql)
.collect::<Vec<_>>();
let mut stmt = conn.prepare(&sql)?;
let rows = stmt.query_map(params.as_slice(), |row| {
let id: String = row.get(0)?;
Ok((parse_uuid(&id)?, row.get::<_, i64>(1)?))
})?;
resolved.extend(rows.collect::<Result<Vec<_>, _>>()?);
}
Ok(resolved)
})
.await
}
async fn get_edge_by_natural_key_including_deleted(
&self,
namespace: &str,
source_id: Uuid,
target_id: Uuid,
relation: EdgeRelation,
) -> Result<Option<Edge>, StorageError> {
let namespace = namespace.to_string();
let source_str = source_id.to_string();
let target_str = target_id.to_string();
let relation_str = relation.to_string();
self.with_reader("get_edge_by_natural_key_including_deleted", move |conn| {
let mut stmt = conn.prepare(
"SELECT namespace, id, source_id, target_id, relation, weight, \
created_at, updated_at, deleted_at, metadata, target_backend \
FROM graph_edges \
WHERE namespace = ?1 AND source_id = ?2 AND target_id = ?3 AND relation = ?4",
)?;
let mut rows = stmt.query(rusqlite::params![
namespace,
source_str,
target_str,
relation_str
])?;
match rows.next()? {
Some(row) => Ok(Some(read_edge(row)?)),
None => Ok(None),
}
})
.await
}
async fn get_edges(&self, ids: &[LinkId]) -> Result<Vec<Edge>, StorageError> {
if ids.is_empty() {
return Ok(Vec::new());
}
const CHUNK: usize = 900;
let id_strs: Vec<String> = ids.iter().map(|id| Uuid::from(*id).to_string()).collect();
let mut result: Vec<Edge> = Vec::with_capacity(ids.len());
for chunk in id_strs.chunks(CHUNK) {
let chunk_owned: Vec<String> = chunk.to_vec();
let edges = self
.with_reader("get_edges", move |conn| {
let placeholders: Vec<String> =
(1..=chunk_owned.len()).map(|i| format!("?{}", i)).collect();
let sql = format!(
"SELECT namespace, id, source_id, target_id, relation, weight, \
created_at, updated_at, deleted_at, metadata, target_backend \
FROM graph_edges WHERE id IN ({}) AND deleted_at IS NULL",
placeholders.join(",")
);
let mut stmt = conn.prepare(&sql)?;
let params: Vec<&dyn rusqlite::types::ToSql> = chunk_owned
.iter()
.map(|s| s as &dyn rusqlite::types::ToSql)
.collect();
let rows = stmt.query_map(params.as_slice(), read_edge)?;
let mut edges = Vec::new();
for row in rows {
edges.push(row?);
}
Ok(edges)
})
.await?;
result.extend(edges);
}
Ok(result)
}
async fn batch_neighbors(
&self,
sources: &[Uuid],
query: NeighborQuery,
) -> Result<Vec<(Uuid, NeighborHit)>, StorageError> {
use khive_storage::types::Direction;
if sources.is_empty() {
return Ok(Vec::new());
}
let mut seen_sources = HashSet::with_capacity(sources.len());
let unique_sources: Vec<Uuid> = sources
.iter()
.copied()
.filter(|source| seen_sources.insert(*source))
.collect();
const CHUNK_SIZE: usize = 880;
let namespace = self.namespace.clone();
let mut result: Vec<(Uuid, NeighborHit)> = Vec::new();
let counted_queries = Arc::new(std::sync::atomic::AtomicU64::new(0));
let counted_rows = Arc::new(std::sync::atomic::AtomicU64::new(0));
for chunk in unique_sources.chunks(CHUNK_SIZE) {
let chunk_owned: Vec<Uuid> = chunk.to_vec();
let query_clone = query.clone();
let ns = namespace.clone();
let closure_queries = Arc::clone(&counted_queries);
let closure_rows = Arc::clone(&counted_rows);
let chunk_result = self
.with_reader("batch_neighbors", move |conn| {
let src_strs: Vec<String> = chunk_owned.iter().map(|u| u.to_string()).collect();
let sources_json = serde_json::to_string(&src_strs).map_err(|error| {
rusqlite::Error::ToSqlConversionFailure(Box::new(error))
})?;
let build_inner_sql =
|direction_out: bool,
q: &NeighborQuery|
-> (String, Vec<String>, Option<f64>) {
let (filter_col, node_col) = if direction_out {
("source_id", "target_id")
} else {
("target_id", "source_id")
};
let mut rel_params: Vec<String> = Vec::new();
let mut conditions: Vec<String> = Vec::new();
let mut param_idx = 3;
if let Some(ref rels) = q.relations {
if !rels.is_empty() {
let ps: Vec<String> = rels
.iter()
.map(|r| {
rel_params.push(r.to_string());
let p = format!("?{param_idx}");
param_idx += 1;
p
})
.collect();
conditions
.push(format!("edges.relation IN ({})", ps.join(",")));
}
}
let min_weight_val = if let Some(min_w) = q.min_weight {
conditions.push(format!("edges.weight >= ?{param_idx}"));
Some(min_w)
} else {
None
};
let where_extra = if conditions.is_empty() {
String::new()
} else {
format!(" AND {}", conditions.join(" AND "))
};
let sql = format!(
"SELECT requested.origin_id, edges.{node_col} AS node_id, \
edges.id AS edge_id, edges.relation, edges.weight \
FROM requested CROSS JOIN graph_edges AS edges \
ON edges.{filter_col} = requested.origin_id \
WHERE edges.namespace = ?1 \
AND edges.deleted_at IS NULL{where_extra}",
);
(sql, rel_params, min_weight_val)
};
let mut all_params: Vec<Box<dyn rusqlite::types::ToSql>> = Vec::new();
all_params.push(Box::new(ns.to_string()));
all_params.push(Box::new(sources_json));
let (combined_inner, rel_params, min_weight_val) = match query_clone.direction {
Direction::Out => build_inner_sql(true, &query_clone),
Direction::In => build_inner_sql(false, &query_clone),
Direction::Both => {
let (out_sql, rel_params, min_weight_val) =
build_inner_sql(true, &query_clone);
let (in_sql, _, _) = build_inner_sql(false, &query_clone);
(
format!("{out_sql} UNION ALL {in_sql}"),
rel_params,
min_weight_val,
)
}
};
for relation in rel_params {
all_params.push(Box::new(relation));
}
if let Some(min_weight) = min_weight_val {
all_params.push(Box::new(min_weight));
}
let limit_param_idx = all_params.len() + 1;
let full_sql = if let Some(lim) = query_clone.limit {
all_params.push(Box::new(lim as i64));
format!(
"WITH requested(origin_id) AS (\
SELECT value FROM json_each(?2)\
) SELECT origin_id, node_id, edge_id, relation, weight \
FROM (SELECT *, ROW_NUMBER() OVER (PARTITION BY origin_id \
ORDER BY weight DESC, node_id ASC) AS rn \
FROM ({combined_inner})) WHERE rn <= ?{limit_param_idx}",
)
} else {
format!(
"WITH requested(origin_id) AS (\
SELECT value FROM json_each(?2)\
) SELECT origin_id, node_id, edge_id, relation, weight \
FROM ({combined_inner})",
)
};
let param_refs: Vec<&dyn rusqlite::types::ToSql> =
all_params.iter().map(|p| p.as_ref()).collect();
let mut stmt = conn.prepare(&full_sql)?;
closure_queries.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
let rows = stmt.query_map(param_refs.as_slice(), |row| {
let origin_str: String = row.get(0)?;
let nid_str: String = row.get(1)?;
let eid_str: String = row.get(2)?;
let relation_str: String = row.get(3)?;
let weight: f64 = row.get(4)?;
Ok((origin_str, nid_str, eid_str, relation_str, weight))
})?;
let mut pairs = Vec::new();
for row in rows {
let (origin_str, nid_str, eid_str, relation_str, weight) = row?;
closure_rows.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
let origin = parse_uuid(&origin_str)?;
let node_id = parse_uuid(&nid_str)?;
let edge_id = parse_uuid(&eid_str)?;
let relation = relation_str.parse::<EdgeRelation>().map_err(|e| {
rusqlite::Error::FromSqlConversionFailure(
3,
rusqlite::types::Type::Text,
Box::new(e),
)
})?;
pairs.push((
origin,
NeighborHit {
node_id,
edge_id,
relation,
weight,
name: None,
kind: None,
entity_type: None,
},
));
}
Ok(pairs)
})
.await;
let pairs = match chunk_result {
Ok(pairs) => pairs,
Err(e) => {
report_graph_usage(&counted_queries, &counted_rows);
return Err(e);
}
};
result.extend(pairs);
}
report_graph_usage(&counted_queries, &counted_rows);
let requested: HashSet<Uuid> = unique_sources.iter().copied().collect();
let mut grouped: HashMap<Uuid, Vec<NeighborHit>> =
HashMap::with_capacity(unique_sources.len());
for (origin, hit) in result {
if !requested.contains(&origin) {
return Err(StorageError::Internal(format!(
"batch_neighbors returned unrequested origin {origin}"
)));
}
grouped.entry(origin).or_default().push(hit);
}
for hits in grouped.values_mut() {
hits.sort_by(|a, b| {
b.weight
.partial_cmp(&a.weight)
.unwrap_or(std::cmp::Ordering::Equal)
.then(a.node_id.cmp(&b.node_id))
.then(a.edge_id.cmp(&b.edge_id))
});
}
let mut ordered = Vec::new();
for &source in sources {
if let Some(hits) = grouped.get(&source) {
ordered.extend(hits.iter().cloned().map(|hit| (source, hit)));
}
}
Ok(ordered)
}
async fn delete_edge(&self, id: LinkId, mode: DeleteMode) -> Result<bool, StorageError> {
let id = Uuid::from(id);
let statement = match mode {
DeleteMode::Soft => {
edge_soft_delete_statement(id, chrono::Utc::now().timestamp_micros())
}
DeleteMode::Hard => edge_hard_delete_statement(id),
};
self.with_writer("delete_edge", move |conn| {
let mut stmt = conn.prepare(&statement.sql)?;
bind_params(&mut stmt, &statement.params)?;
Ok(stmt.raw_execute()? > 0)
})
.await
}
async fn query_edges(
&self,
filter: EdgeFilter,
sort: Vec<SortOrder<EdgeSortField>>,
page: PageRequest,
) -> Result<Page<Edge>, StorageError> {
let namespace = self.namespace.clone();
let limit_i64 = i64::from(page.limit);
let offset_i64 = i64::try_from(page.offset).map_err(|_| StorageError::InvalidInput {
capability: StorageCapability::Graph,
operation: "query_edges".into(),
message: format!(
"PageRequest: offset must be <= i64::MAX, got {}",
page.offset
),
})?;
self.with_reader("query_edges", move |conn| {
let (where_clause, filter_params) = build_edge_filter_sql(&namespace, &filter);
let count_sql = format!("SELECT COUNT(*) FROM graph_edges{}", where_clause);
let total: i64 = {
let mut stmt = conn.prepare(&count_sql)?;
let param_refs: Vec<&dyn rusqlite::types::ToSql> =
filter_params.iter().map(|p| p.as_ref()).collect();
stmt.query_row(param_refs.as_slice(), |row| row.get(0))?
};
let order_clause = edge_order_clause(&sort);
let (_, data_filter_params) = build_edge_filter_sql(&namespace, &filter);
let mut all_params: Vec<Box<dyn rusqlite::types::ToSql>> = data_filter_params;
all_params.push(Box::new(limit_i64));
all_params.push(Box::new(offset_i64));
let limit_idx = all_params.len() - 1;
let offset_idx = all_params.len();
let data_sql = format!(
"SELECT namespace, id, source_id, target_id, relation, weight, \
created_at, updated_at, deleted_at, metadata, target_backend \
FROM graph_edges{}{} LIMIT ?{} OFFSET ?{}",
where_clause, order_clause, limit_idx, offset_idx,
);
let mut stmt = conn.prepare(&data_sql)?;
let param_refs: Vec<&dyn rusqlite::types::ToSql> =
all_params.iter().map(|p| p.as_ref()).collect();
let rows = stmt.query_map(param_refs.as_slice(), read_edge)?;
let mut items = Vec::new();
for row in rows {
items.push(row?);
}
Ok(Page {
items,
total: Some(total as u64),
})
})
.await
}
async fn count_edges(&self, filter: EdgeFilter) -> Result<u64, StorageError> {
let namespace = self.namespace.clone();
self.with_reader("count_edges", move |conn| {
let (where_clause, params) = build_edge_filter_sql(&namespace, &filter);
let sql = format!("SELECT COUNT(*) FROM graph_edges{}", where_clause);
let mut stmt = conn.prepare(&sql)?;
let param_refs: Vec<&dyn rusqlite::types::ToSql> =
params.iter().map(|p| p.as_ref()).collect();
let count: i64 = stmt.query_row(param_refs.as_slice(), |row| row.get(0))?;
Ok(count as u64)
})
.await
}
async fn count_edges_in_namespaces(
&self,
namespaces: &[String],
filter: EdgeFilter,
) -> Result<u64, StorageError> {
let namespaces: Vec<String> = namespaces
.iter()
.cloned()
.collect::<HashSet<_>>()
.into_iter()
.collect();
self.with_reader("count_edges_in_namespaces", move |conn| {
let mut total = 0;
for chunk in namespaces.chunks(NAMESPACE_COUNT_CHUNK_SIZE) {
let (where_clause, params) = build_edge_filter_sql_for_namespaces(chunk, &filter);
let sql = format!("SELECT COUNT(*) FROM graph_edges{where_clause}");
let mut stmt = conn.prepare(&sql)?;
let param_refs: Vec<&dyn rusqlite::types::ToSql> =
params.iter().map(|p| p.as_ref()).collect();
let count: i64 = stmt.query_row(param_refs.as_slice(), |row| row.get(0))?;
total += count as u64;
}
Ok(total)
})
.await
}
async fn query_edges_in_namespaces(
&self,
namespaces: &[String],
filter: EdgeFilter,
sort: Vec<SortOrder<EdgeSortField>>,
page: PageRequest,
) -> Result<Page<Edge>, StorageError> {
let namespaces: Vec<String> = {
let mut seen = HashSet::new();
namespaces
.iter()
.filter(|ns| seen.insert((*ns).clone()))
.cloned()
.collect()
};
let limit_i64 = i64::from(page.limit);
let offset_i64 = i64::try_from(page.offset).map_err(|_| StorageError::InvalidInput {
capability: StorageCapability::Graph,
operation: "query_edges_in_namespaces".into(),
message: format!(
"PageRequest: offset must be <= i64::MAX, got {}",
page.offset
),
})?;
self.with_reader("query_edges_in_namespaces", move |conn| {
let namespaces_json = serde_json::to_string(&namespaces)
.map_err(|error| rusqlite::Error::ToSqlConversionFailure(Box::new(error)))?;
let (where_clause, filter_params) =
build_edge_filter_sql_for_namespaces_json(&namespaces_json, &filter);
let count_sql = format!("SELECT COUNT(*) FROM graph_edges{}", where_clause);
let total: i64 = {
let mut stmt = conn.prepare(&count_sql)?;
let param_refs: Vec<&dyn rusqlite::types::ToSql> =
filter_params.iter().map(|p| p.as_ref()).collect();
stmt.query_row(param_refs.as_slice(), |row| row.get(0))?
};
let order_clause = edge_order_clause(&sort);
let (_, data_filter_params) =
build_edge_filter_sql_for_namespaces_json(&namespaces_json, &filter);
let mut all_params: Vec<Box<dyn rusqlite::types::ToSql>> = data_filter_params;
all_params.push(Box::new(limit_i64));
all_params.push(Box::new(offset_i64));
let limit_idx = all_params.len() - 1;
let offset_idx = all_params.len();
let data_sql = format!(
"SELECT namespace, id, source_id, target_id, relation, weight, \
created_at, updated_at, deleted_at, metadata, target_backend \
FROM graph_edges{}{} LIMIT ?{} OFFSET ?{}",
where_clause, order_clause, limit_idx, offset_idx,
);
let mut stmt = conn.prepare(&data_sql)?;
let param_refs: Vec<&dyn rusqlite::types::ToSql> =
all_params.iter().map(|p| p.as_ref()).collect();
let rows = stmt.query_map(param_refs.as_slice(), read_edge)?;
let mut items = Vec::new();
for row in rows {
items.push(row?);
}
Ok(Page {
items,
total: Some(total as u64),
})
})
.await
}
async fn count_edges_by_relation(&self) -> Result<Vec<(EdgeRelation, u64)>, StorageError> {
let namespace = self.namespace.clone();
self.with_reader("count_edges_by_relation", move |conn| {
let sql = "SELECT relation, COUNT(*) FROM graph_edges \
WHERE namespace = ?1 AND deleted_at IS NULL \
GROUP BY relation";
let mut stmt = conn.prepare(sql)?;
let rows = stmt.query_map([&namespace], |row| {
let relation_str: String = row.get(0)?;
let count: i64 = row.get(1)?;
Ok((relation_str, count))
})?;
let mut out = Vec::new();
for row in rows {
let (relation_str, count) = row?;
let relation = relation_str.parse::<EdgeRelation>().map_err(|e| {
rusqlite::Error::FromSqlConversionFailure(
0,
rusqlite::types::Type::Text,
Box::new(e),
)
})?;
out.push((relation, count as u64));
}
Ok(out)
})
.await
}
async fn count_edges_by_relation_in_namespaces(
&self,
namespaces: &[String],
) -> Result<Vec<(EdgeRelation, u64)>, StorageError> {
let namespaces: Vec<String> = namespaces
.iter()
.cloned()
.collect::<HashSet<_>>()
.into_iter()
.collect();
self.with_reader("count_edges_by_relation_in_namespaces", move |conn| {
let mut totals = HashMap::new();
for chunk in namespaces.chunks(NAMESPACE_COUNT_CHUNK_SIZE) {
let (where_clause, params) =
build_edge_filter_sql_for_namespaces(chunk, &EdgeFilter::default());
let sql = format!(
"SELECT relation, COUNT(*) FROM graph_edges{where_clause} GROUP BY relation"
);
let mut stmt = conn.prepare(&sql)?;
let param_refs: Vec<&dyn rusqlite::types::ToSql> =
params.iter().map(|p| p.as_ref()).collect();
let rows = stmt.query_map(param_refs.as_slice(), |row| {
let relation_str: String = row.get(0)?;
let count: i64 = row.get(1)?;
Ok((relation_str, count))
})?;
for row in rows {
let (relation_str, count) = row?;
let relation = relation_str.parse::<EdgeRelation>().map_err(|e| {
rusqlite::Error::FromSqlConversionFailure(
0,
rusqlite::types::Type::Text,
Box::new(e),
)
})?;
*totals.entry(relation).or_insert(0) += count as u64;
}
}
Ok(totals.into_iter().collect())
})
.await
}
async fn query_edges_after(
&self,
filter: EdgeFilter,
after: Option<Uuid>,
limit: u32,
) -> Result<EdgeSeekPage, StorageError> {
let namespace = self.namespace.clone();
let limit_usize = limit as usize;
let probe_limit_i64 = i64::from(limit) + 1;
self.with_reader("query_edges_after", move |conn| {
let (mut where_clause, mut params) = build_edge_filter_sql(&namespace, &filter);
if let Some(cursor) = after {
params.push(Box::new(cursor.to_string()));
where_clause.push_str(&format!(" AND id > ?{}", params.len()));
}
params.push(Box::new(probe_limit_i64));
let limit_idx = params.len();
let data_sql = format!(
"SELECT namespace, id, source_id, target_id, relation, weight, \
created_at, updated_at, deleted_at, metadata, target_backend \
FROM graph_edges{} ORDER BY id ASC LIMIT ?{}",
where_clause, limit_idx,
);
let mut stmt = conn.prepare(&data_sql)?;
let param_refs: Vec<&dyn rusqlite::types::ToSql> =
params.iter().map(|p| p.as_ref()).collect();
let rows = stmt.query_map(param_refs.as_slice(), read_edge)?;
let mut items = Vec::new();
for row in rows {
items.push(row?);
}
let has_more = items.len() > limit_usize;
if has_more {
items.truncate(limit_usize);
}
let next_after = if has_more {
items.last().map(|e| Uuid::from(e.id))
} else {
None
};
Ok(EdgeSeekPage { items, next_after })
})
.await
}
async fn query_edges_sequence_after(
&self,
filter: EdgeFilter,
after: Option<SeekCursor>,
limit: u32,
) -> Result<SeekPage<Edge>, StorageError> {
if limit == 0 {
return Ok(SeekPage::default());
}
let namespace = self.namespace.clone();
let limit_usize = limit as usize;
let probe_limit_i64 = i64::from(limit) + 1;
self.with_reader("query_edges_sequence_after", move |conn| {
let (mut where_clause, mut params) = build_edge_filter_sql(&namespace, &filter);
if let Some(cursor) = after {
params.push(Box::new(cursor.sequence));
where_clause.push_str(&format!(" AND graph_edges_seq.seq > ?{}", params.len()));
}
params.push(Box::new(probe_limit_i64));
let limit_idx = params.len();
let sql = format!(
"SELECT namespace, id, source_id, target_id, relation, weight, \
created_at, updated_at, deleted_at, metadata, target_backend, \
graph_edges_seq.seq \
FROM graph_edges_seq CROSS JOIN graph_edges \
ON graph_edges.id = graph_edges_seq.edge_id{where_clause} \
ORDER BY graph_edges_seq.seq ASC LIMIT ?{limit_idx}"
);
let mut stmt = conn.prepare(&sql)?;
let param_refs: Vec<&dyn rusqlite::types::ToSql> =
params.iter().map(|param| param.as_ref()).collect();
let rows = stmt.query_map(param_refs.as_slice(), |row| {
Ok((read_edge(row)?, row.get::<_, i64>(11)?))
})?;
let mut entries = rows.collect::<Result<Vec<_>, _>>()?;
let has_more = entries.len() > limit_usize;
if has_more {
entries.truncate(limit_usize);
}
let next_after = if has_more {
entries.last().map(|(edge, sequence)| SeekCursor {
sequence: *sequence,
id: Uuid::from(edge.id),
})
} else {
None
};
let items = entries.into_iter().map(|(edge, _)| edge).collect();
Ok(SeekPage { items, next_after })
})
.await
}
async fn neighbors(
&self,
node_id: Uuid,
query: NeighborQuery,
) -> Result<Vec<NeighborHit>, StorageError> {
count_neighbor_select();
let namespace = self.namespace.clone();
let node_str = node_id.to_string();
let counted_queries = Arc::new(std::sync::atomic::AtomicU64::new(0));
let counted_rows = Arc::new(std::sync::atomic::AtomicU64::new(0));
let closure_queries = Arc::clone(&counted_queries);
let closure_rows = Arc::clone(&counted_rows);
let result = self
.with_reader("neighbors", move |conn| {
let base_out = "SELECT target_id AS node_id, id AS edge_id, relation, weight \
FROM graph_edges \
WHERE namespace = ?1 AND source_id = ?2 AND deleted_at IS NULL";
let base_in = "SELECT source_id AS node_id, id AS edge_id, relation, weight \
FROM graph_edges \
WHERE namespace = ?1 AND target_id = ?2 AND deleted_at IS NULL";
let sql = match query.direction {
Direction::Out => base_out.to_string(),
Direction::In => base_in.to_string(),
Direction::Both => format!("{} UNION ALL {}", base_out, base_in),
};
let (where_extra, limit_clause, extra_params) = neighbor_extra_clause(&query, 3);
let full_sql = format!(
"SELECT node_id, edge_id, relation, weight FROM ({}){} \
ORDER BY weight DESC, node_id ASC{}",
sql, where_extra, limit_clause
);
let mut stmt = conn.prepare(&full_sql)?;
let mut all_params: Vec<Box<dyn rusqlite::types::ToSql>> = Vec::new();
all_params.push(Box::new(namespace.clone()));
all_params.push(Box::new(node_str.clone()));
all_params.extend(extra_params);
let param_refs: Vec<&dyn rusqlite::types::ToSql> =
all_params.iter().map(|p| p.as_ref()).collect();
closure_queries.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
let rows = stmt.query_map(param_refs.as_slice(), |row| {
let nid_str: String = row.get(0)?;
let eid_str: String = row.get(1)?;
let relation_str: String = row.get(2)?;
let weight: f64 = row.get(3)?;
Ok((nid_str, eid_str, relation_str, weight))
})?;
let mut hits = Vec::new();
for row in rows {
let (nid_str, eid_str, relation_str, weight) = row?;
closure_rows.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
let relation = relation_str.parse::<EdgeRelation>().map_err(|e| {
rusqlite::Error::FromSqlConversionFailure(
2,
rusqlite::types::Type::Text,
Box::new(e),
)
})?;
hits.push(NeighborHit {
node_id: parse_uuid(&nid_str)?,
edge_id: parse_uuid(&eid_str)?,
relation,
weight,
name: None,
kind: None,
entity_type: None,
});
}
Ok(hits)
})
.await;
report_graph_usage(&counted_queries, &counted_rows);
result
}
async fn neighbors_both_directions(
&self,
node_id: Uuid,
query: NeighborQuery,
) -> Result<Vec<DirectedNeighborHit>, StorageError> {
count_neighbor_select();
let namespace = self.namespace.clone();
let node_str = node_id.to_string();
let counted_queries = Arc::new(std::sync::atomic::AtomicU64::new(0));
let counted_rows = Arc::new(std::sync::atomic::AtomicU64::new(0));
let closure_queries = Arc::clone(&counted_queries);
let closure_rows = Arc::clone(&counted_rows);
let result = self
.with_reader("neighbors_both_directions", move |conn| {
let base_out = "SELECT target_id AS node_id, id AS edge_id, relation, weight, \
'out' AS dir \
FROM graph_edges \
WHERE namespace = ?1 AND source_id = ?2 AND deleted_at IS NULL";
let base_in = "SELECT source_id AS node_id, id AS edge_id, relation, weight, \
'in' AS dir \
FROM graph_edges \
WHERE namespace = ?1 AND target_id = ?2 AND deleted_at IS NULL";
let sql = format!("{} UNION ALL {}", base_out, base_in);
let (where_extra, limit_clause, extra_params) = neighbor_extra_clause(&query, 3);
let full_sql = format!(
"SELECT node_id, edge_id, relation, weight, dir FROM ({}){} \
ORDER BY weight DESC, node_id ASC, \
CASE dir WHEN 'out' THEN 0 ELSE 1 END ASC, edge_id ASC{}",
sql, where_extra, limit_clause
);
let mut stmt = conn.prepare(&full_sql)?;
let mut all_params: Vec<Box<dyn rusqlite::types::ToSql>> = Vec::new();
all_params.push(Box::new(namespace.clone()));
all_params.push(Box::new(node_str.clone()));
all_params.extend(extra_params);
let param_refs: Vec<&dyn rusqlite::types::ToSql> =
all_params.iter().map(|p| p.as_ref()).collect();
closure_queries.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
let rows = stmt.query_map(param_refs.as_slice(), |row| {
let nid_str: String = row.get(0)?;
let eid_str: String = row.get(1)?;
let relation_str: String = row.get(2)?;
let weight: f64 = row.get(3)?;
let dir_str: String = row.get(4)?;
Ok((nid_str, eid_str, relation_str, weight, dir_str))
})?;
let mut hits = Vec::new();
for row in rows {
let (nid_str, eid_str, relation_str, weight, dir_str) = row?;
closure_rows.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
let relation = relation_str.parse::<EdgeRelation>().map_err(|e| {
rusqlite::Error::FromSqlConversionFailure(
2,
rusqlite::types::Type::Text,
Box::new(e),
)
})?;
let direction = if dir_str == "out" {
Direction::Out
} else {
Direction::In
};
hits.push(DirectedNeighborHit {
hit: NeighborHit {
node_id: parse_uuid(&nid_str)?,
edge_id: parse_uuid(&eid_str)?,
relation,
weight,
name: None,
kind: None,
entity_type: None,
},
direction,
});
}
Ok(hits)
})
.await;
report_graph_usage(&counted_queries, &counted_rows);
result
}
async fn traverse(&self, request: TraversalRequest) -> Result<Vec<GraphPath>, StorageError> {
request
.validate()
.map_err(|message| StorageError::InvalidInput {
capability: StorageCapability::Graph,
operation: "traverse".into(),
message,
})?;
if request.roots.is_empty() {
return Ok(Vec::new());
}
let mut distinct_roots = HashSet::with_capacity(request.roots.len());
let roots = request
.roots
.iter()
.copied()
.filter(|root| distinct_roots.insert(*root))
.collect::<Vec<_>>();
let opts = request.options;
let include_roots = request.include_roots;
let namespace = self.namespace.clone();
let origin = self.pool.origin();
let budget = request.execution_budget;
let counted_rows = Arc::new(std::sync::atomic::AtomicU64::new(0));
let counted_queries = Arc::new(std::sync::atomic::AtomicU64::new(0));
let closure_rows = Arc::clone(&counted_rows);
let closure_queries = Arc::clone(&counted_queries);
let result = self
.with_reader("traverse", move |conn| {
Ok(run_bounded_traversal(
conn,
roots,
opts,
include_roots,
namespace,
origin,
budget,
closure_rows.as_ref(),
closure_queries.as_ref(),
))
})
.await
.and_then(|inner| inner);
khive_storage::usage::count(
khive_storage::usage::UsageUnit::DbRoundTrips,
counted_queries.load(std::sync::atomic::Ordering::Relaxed),
);
khive_storage::usage::count(
khive_storage::usage::UsageUnit::GraphHops,
counted_rows.load(std::sync::atomic::Ordering::Relaxed),
);
result
}
async fn purge_incident_edges(&self, node_id: Uuid) -> Result<u64, StorageError> {
let statement = purge_incident_edges_statement(node_id);
self.with_writer("purge_incident_edges", move |conn| {
let mut stmt = conn.prepare(&statement.sql)?;
bind_params(&mut stmt, &statement.params)?;
Ok(stmt.raw_execute()? as u64)
})
.await
}
}
const GRAPH_DDL: &str = include_str!("../../sql/graph-ddl.sql");
pub(crate) fn ensure_graph_schema(conn: &rusqlite::Connection) -> Result<(), rusqlite::Error> {
conn.execute_batch(GRAPH_DDL)
}
#[cfg(test)]
#[path = "graph_tests.rs"]
mod tests;