1use std::collections::{HashMap, HashSet, VecDeque};
4use std::sync::Arc;
5
6use async_trait::async_trait;
7use chrono::{DateTime, TimeZone, Utc};
8use rusqlite::OptionalExtension;
9use uuid::Uuid;
10
11use khive_storage::error::StorageError;
12use khive_storage::types::{
13 BatchWriteSummary, DeleteMode, DirectedNeighborHit, Direction, Edge, EdgeFilter, EdgeSeekPage,
14 EdgeSortField, GraphPath, GuardedBatchOutcome, GuardedBatchRefusal, GuardedWriteOutcome,
15 MissingEndpoints, NeighborHit, NeighborQuery, Page, PageRequest, PathNode, SeekCursor,
16 SeekPage, SortDirection, SortOrder, SqlStatement, SqlValue, TraversalExecutionBudget,
17 TraversalOptions, TraversalRequest,
18};
19use khive_storage::GraphStore;
20use khive_storage::LinkId;
21use khive_storage::StorageCapability;
22use khive_types::EdgeRelation;
23
24use crate::error::SqliteError;
25use crate::pool::ConnectionPool;
26use crate::sql_bridge::bind_params;
27use crate::writer_task::WriterTaskHandle;
28
29fn map_err(e: rusqlite::Error, op: &'static str) -> StorageError {
31 StorageError::driver(StorageCapability::Graph, op, e)
32}
33
34fn map_sqlite_err(e: SqliteError, op: &'static str) -> StorageError {
35 StorageError::driver(StorageCapability::Graph, op, e)
36}
37
38const NAMESPACE_COUNT_CHUNK_SIZE: usize = 500;
39
40const EDGE_NATURAL_KEY_CONFLICT_SET: &str = "weight = excluded.weight, \
56 updated_at = excluded.updated_at, \
57 deleted_at = NULL, \
58 metadata = excluded.metadata, \
59 target_backend = excluded.target_backend";
60
61fn endpoint_exists_clause(id_param: &str) -> String {
72 format!(
73 "EXISTS (SELECT 1 FROM entities WHERE id = {id_param} AND deleted_at IS NULL) \
74 OR EXISTS (SELECT 1 FROM notes WHERE id = {id_param} AND deleted_at IS NULL) \
75 OR EXISTS (SELECT 1 FROM events WHERE id = {id_param}) \
76 OR EXISTS (SELECT 1 FROM graph_edges WHERE id = {id_param} AND deleted_at IS NULL)"
77 )
78}
79
80pub fn edge_upsert_statement(edge: &Edge) -> SqlStatement {
84 let (source_id, target_id) =
85 canonical_edge_endpoints(edge.relation, edge.source_id, edge.target_id);
86 let metadata_str = edge
87 .metadata
88 .as_ref()
89 .map(|v| serde_json::to_string(v).unwrap_or_default());
90 SqlStatement {
91 sql: format!(
92 "INSERT INTO graph_edges \
93 (namespace, id, source_id, target_id, relation, weight, \
94 created_at, updated_at, deleted_at, metadata, target_backend) \
95 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11) \
96 ON CONFLICT(namespace, id) DO UPDATE SET \
97 source_id = excluded.source_id, \
98 target_id = excluded.target_id, \
99 relation = excluded.relation, \
100 {EDGE_NATURAL_KEY_CONFLICT_SET} \
101 ON CONFLICT(namespace, source_id, target_id, relation) DO UPDATE SET \
102 {EDGE_NATURAL_KEY_CONFLICT_SET}"
103 ),
104 params: vec![
105 SqlValue::Text(edge.namespace.clone()),
106 SqlValue::Text(Uuid::from(edge.id).to_string()),
107 SqlValue::Text(source_id.to_string()),
108 SqlValue::Text(target_id.to_string()),
109 SqlValue::Text(edge.relation.to_string()),
110 SqlValue::Float(edge.weight),
111 SqlValue::Integer(edge.created_at.timestamp_micros()),
112 SqlValue::Integer(edge.updated_at.timestamp_micros()),
113 match edge.deleted_at {
114 Some(t) => SqlValue::Integer(t.timestamp_micros()),
115 None => SqlValue::Null,
116 },
117 match metadata_str {
118 Some(m) => SqlValue::Text(m),
119 None => SqlValue::Null,
120 },
121 match &edge.target_backend {
122 Some(b) => SqlValue::Text(b.clone()),
123 None => SqlValue::Null,
124 },
125 ],
126 label: Some("edge-upsert".to_string()),
127 }
128}
129
130pub fn edge_replace_if_unchanged_statement(
141 edge: &Edge,
142 expected_updated_at: DateTime<Utc>,
143 expected_deleted_at: Option<DateTime<Utc>>,
144) -> SqlStatement {
145 let (source_id, target_id) =
146 canonical_edge_endpoints(edge.relation, edge.source_id, edge.target_id);
147 let metadata_str = edge
148 .metadata
149 .as_ref()
150 .map(|v| serde_json::to_string(v).unwrap_or_default());
151 SqlStatement {
152 sql: "UPDATE graph_edges SET \
153 namespace = ?1, source_id = ?2, target_id = ?3, relation = ?4, weight = ?5, \
154 updated_at = ?6, deleted_at = ?7, metadata = ?8, target_backend = ?9 \
155 WHERE id = ?10 AND updated_at = ?11 AND deleted_at IS ?12 \
156 AND ?6 > updated_at"
157 .to_string(),
158 params: vec![
159 SqlValue::Text(edge.namespace.clone()),
160 SqlValue::Text(source_id.to_string()),
161 SqlValue::Text(target_id.to_string()),
162 SqlValue::Text(edge.relation.to_string()),
163 SqlValue::Float(edge.weight),
164 SqlValue::Integer(edge.updated_at.timestamp_micros()),
165 match edge.deleted_at {
166 Some(t) => SqlValue::Integer(t.timestamp_micros()),
167 None => SqlValue::Null,
168 },
169 match metadata_str {
170 Some(m) => SqlValue::Text(m),
171 None => SqlValue::Null,
172 },
173 match &edge.target_backend {
174 Some(b) => SqlValue::Text(b.clone()),
175 None => SqlValue::Null,
176 },
177 SqlValue::Text(Uuid::from(edge.id).to_string()),
178 SqlValue::Integer(expected_updated_at.timestamp_micros()),
179 match expected_deleted_at {
180 Some(value) => SqlValue::Integer(value.timestamp_micros()),
181 None => SqlValue::Null,
182 },
183 ],
184 label: Some("edge-replace-if-unchanged".to_string()),
185 }
186}
187
188#[allow(clippy::too_many_arguments)]
208pub fn edge_insert_guarded_by_endpoints_statement(
209 namespace: &str,
210 edge_id: Uuid,
211 source_id: Uuid,
212 target_id: Uuid,
213 relation: EdgeRelation,
214 weight: f64,
215 now: i64,
216 metadata: Option<&str>,
217) -> SqlStatement {
218 let src_exists = endpoint_exists_clause("?3");
219 let tgt_exists = endpoint_exists_clause("?4");
220 SqlStatement {
221 sql: format!(
222 "INSERT INTO graph_edges \
223 (namespace, id, source_id, target_id, relation, weight, \
224 created_at, updated_at, metadata) \
225 SELECT ?1, ?2, ?3, ?4, ?5, ?6, ?7, ?7, ?8 \
226 WHERE ({src_exists}) AND ({tgt_exists}) \
227 ON CONFLICT(namespace, source_id, target_id, relation) DO UPDATE SET \
228 {EDGE_NATURAL_KEY_CONFLICT_SET}"
229 ),
230 params: vec![
231 SqlValue::Text(namespace.to_string()),
232 SqlValue::Text(edge_id.to_string()),
233 SqlValue::Text(source_id.to_string()),
234 SqlValue::Text(target_id.to_string()),
235 SqlValue::Text(relation.as_str().to_string()),
236 SqlValue::Float(weight),
237 SqlValue::Integer(now),
238 match metadata {
239 Some(m) => SqlValue::Text(m.to_string()),
240 None => SqlValue::Null,
241 },
242 ],
243 label: Some("atomic-link-insert-edge-where-exists".to_string()),
244 }
245}
246
247pub fn edge_soft_delete_statement(id: Uuid, now: i64) -> SqlStatement {
249 SqlStatement {
250 sql: "UPDATE graph_edges SET deleted_at = ?2, updated_at = ?2 \
251 WHERE id = ?1 AND deleted_at IS NULL"
252 .to_string(),
253 params: vec![SqlValue::Text(id.to_string()), SqlValue::Integer(now)],
254 label: Some("edge-delete-soft".to_string()),
255 }
256}
257
258pub fn edge_hard_delete_statement(id: Uuid) -> SqlStatement {
260 SqlStatement {
261 sql: "DELETE FROM graph_edges WHERE id = ?1".to_string(),
262 params: vec![SqlValue::Text(id.to_string())],
263 label: Some("edge-delete-hard".to_string()),
264 }
265}
266
267pub fn purge_incident_edges_statement(node_id: Uuid) -> SqlStatement {
269 SqlStatement {
270 sql: "DELETE FROM graph_edges WHERE source_id = ?1 OR target_id = ?1".to_string(),
271 params: vec![SqlValue::Text(node_id.to_string())],
272 label: Some("edge-purge-incident".to_string()),
273 }
274}
275
276pub const EDGE_SYMMETRIC_CONFLICT_PROBE_SQL: &str = "SELECT id FROM graph_edges \
297 WHERE namespace = ?1 AND source_id = ?2 AND target_id = ?3 \
298 AND relation = ?4 AND id != ?5";
299
300pub const EDGE_SYMMETRIC_DELETE_NONCANONICAL_SQL: &str =
301 "DELETE FROM graph_edges WHERE namespace = ?1 AND id = ?2";
302
303pub const EDGE_SYMMETRIC_DELETE_NONCANONICAL_GUARDED_SQL: &str =
316 "DELETE FROM graph_edges WHERE namespace = ?1 AND id = ?2 \
317 AND updated_at = ?3 AND deleted_at IS ?4";
318
319pub const EDGE_SYMMETRIC_UPDATE_INPLACE_SQL: &str = "UPDATE graph_edges SET \
326 source_id = ?1, target_id = ?2, relation = ?3, \
327 weight = ?4, updated_at = ?5, metadata = ?6 \
328 WHERE namespace = ?7 AND id = ?8 \
329 AND updated_at = ?9 AND deleted_at IS ?10 \
330 AND ?5 > updated_at";
331
332pub fn edge_symmetric_conflict_probe_statement(
335 namespace: &str,
336 canon_src: Uuid,
337 canon_tgt: Uuid,
338 relation: EdgeRelation,
339 exclude_id: Uuid,
340) -> SqlStatement {
341 SqlStatement {
342 sql: EDGE_SYMMETRIC_CONFLICT_PROBE_SQL.to_string(),
343 params: vec![
344 SqlValue::Text(namespace.to_string()),
345 SqlValue::Text(canon_src.to_string()),
346 SqlValue::Text(canon_tgt.to_string()),
347 SqlValue::Text(relation.to_string()),
348 SqlValue::Text(exclude_id.to_string()),
349 ],
350 label: Some("edge-symmetric-conflict-probe".to_string()),
351 }
352}
353
354pub fn edge_symmetric_delete_noncanonical_statement(namespace: &str, id: Uuid) -> SqlStatement {
357 SqlStatement {
358 sql: EDGE_SYMMETRIC_DELETE_NONCANONICAL_SQL.to_string(),
359 params: vec![
360 SqlValue::Text(namespace.to_string()),
361 SqlValue::Text(id.to_string()),
362 ],
363 label: Some("edge-symmetric-delete-noncanonical".to_string()),
364 }
365}
366
367#[allow(clippy::too_many_arguments)]
371pub fn edge_symmetric_update_inplace_statement(
372 namespace: &str,
373 id: Uuid,
374 canon_src: Uuid,
375 canon_tgt: Uuid,
376 relation: EdgeRelation,
377 weight: f64,
378 updated_at_micros: i64,
379 metadata: Option<&str>,
380 expected_updated_at_micros: i64,
381 expected_deleted_at_micros: Option<i64>,
382) -> SqlStatement {
383 SqlStatement {
384 sql: EDGE_SYMMETRIC_UPDATE_INPLACE_SQL.to_string(),
385 params: vec![
386 SqlValue::Text(canon_src.to_string()),
387 SqlValue::Text(canon_tgt.to_string()),
388 SqlValue::Text(relation.to_string()),
389 SqlValue::Float(weight),
390 SqlValue::Integer(updated_at_micros),
391 match metadata {
392 Some(m) => SqlValue::Text(m.to_string()),
393 None => SqlValue::Null,
394 },
395 SqlValue::Text(namespace.to_string()),
396 SqlValue::Text(id.to_string()),
397 SqlValue::Integer(expected_updated_at_micros),
398 match expected_deleted_at_micros {
399 Some(value) => SqlValue::Integer(value),
400 None => SqlValue::Null,
401 },
402 ],
403 label: Some("edge-symmetric-update-inplace".to_string()),
404 }
405}
406
407#[allow(clippy::too_many_arguments)]
501pub fn edge_symmetric_delete_if_conflict_statement(
502 namespace: &str,
503 id: Uuid,
504 canon_src: Uuid,
505 canon_tgt: Uuid,
506 relation: EdgeRelation,
507 expected_updated_at_micros: i64,
508 expected_deleted_at_micros: Option<i64>,
509) -> SqlStatement {
510 SqlStatement {
511 sql: "DELETE FROM graph_edges \
512 WHERE namespace = ?1 AND id = ?2 \
513 AND updated_at = ?6 AND deleted_at IS ?7 \
514 AND EXISTS ( \
515 SELECT 1 FROM graph_edges \
516 WHERE namespace = ?1 AND source_id = ?3 AND target_id = ?4 \
517 AND relation = ?5 AND id != ?2 \
518 )"
519 .to_string(),
520 params: vec![
521 SqlValue::Text(namespace.to_string()),
522 SqlValue::Text(id.to_string()),
523 SqlValue::Text(canon_src.to_string()),
524 SqlValue::Text(canon_tgt.to_string()),
525 SqlValue::Text(relation.to_string()),
526 SqlValue::Integer(expected_updated_at_micros),
527 match expected_deleted_at_micros {
528 Some(value) => SqlValue::Integer(value),
529 None => SqlValue::Null,
530 },
531 ],
532 label: Some("edge-symmetric-delete-if-conflict".to_string()),
533 }
534}
535
536#[allow(clippy::too_many_arguments)]
545pub fn edge_symmetric_absorb_or_update_inplace_statement(
546 namespace: &str,
547 id: Uuid,
548 canon_src: Uuid,
549 canon_tgt: Uuid,
550 relation: EdgeRelation,
551 weight: f64,
552 updated_at_micros: i64,
553 metadata: Option<&str>,
554 target_backend: Option<&str>,
555 expected_updated_at_micros: i64,
556 expected_deleted_at_micros: Option<i64>,
557) -> SqlStatement {
558 SqlStatement {
559 sql: "UPDATE graph_edges SET \
560 source_id = CASE WHEN id = ?2 THEN ?3 ELSE source_id END, \
561 target_id = CASE WHEN id = ?2 THEN ?4 ELSE target_id END, \
562 relation = CASE WHEN id = ?2 THEN ?5 ELSE relation END, \
563 weight = CASE WHEN id = ?2 THEN ?6 ELSE weight END, \
564 updated_at = CASE WHEN id = ?2 THEN ?7 ELSE updated_at END, \
565 deleted_at = CASE WHEN id = ?2 THEN NULL ELSE deleted_at END, \
566 metadata = CASE WHEN id = ?2 THEN ?8 ELSE metadata END, \
567 target_backend = CASE WHEN id = ?2 THEN ?9 ELSE target_backend END \
568 WHERE namespace = ?1 \
569 AND ( \
570 (id = ?2 AND changes() = 0 AND updated_at = ?10 AND deleted_at IS ?11 \
571 AND ?7 > updated_at) \
572 OR (source_id = ?3 AND target_id = ?4 AND relation = ?5 \
573 AND id != ?2 AND changes() = 1) \
574 )"
575 .to_string(),
576 params: vec![
577 SqlValue::Text(namespace.to_string()),
578 SqlValue::Text(id.to_string()),
579 SqlValue::Text(canon_src.to_string()),
580 SqlValue::Text(canon_tgt.to_string()),
581 SqlValue::Text(relation.to_string()),
582 SqlValue::Float(weight),
583 SqlValue::Integer(updated_at_micros),
584 match metadata {
585 Some(m) => SqlValue::Text(m.to_string()),
586 None => SqlValue::Null,
587 },
588 match target_backend {
589 Some(b) => SqlValue::Text(b.to_string()),
590 None => SqlValue::Null,
591 },
592 SqlValue::Integer(expected_updated_at_micros),
593 match expected_deleted_at_micros {
594 Some(value) => SqlValue::Integer(value),
595 None => SqlValue::Null,
596 },
597 ],
598 label: Some("edge-symmetric-absorb-or-update-inplace".to_string()),
599 }
600}
601
602pub struct SqlGraphStore {
604 pool: Arc<ConnectionPool>,
605 is_file_backed: bool,
606 namespace: String,
610 writer_task: Option<WriterTaskHandle>,
611}
612
613impl SqlGraphStore {
614 pub fn new_scoped(
622 pool: Arc<ConnectionPool>,
623 is_file_backed: bool,
624 namespace: impl Into<String>,
625 ) -> Self {
626 let writer_task = pool.writer_task_handle().ok().flatten();
632
633 Self {
634 pool,
635 is_file_backed,
636 namespace: namespace.into(),
637 writer_task,
638 }
639 }
640
641 fn open_standalone_writer(&self) -> Result<rusqlite::Connection, StorageError> {
642 self.pool
643 .open_standalone_writer()
644 .map_err(|e| map_sqlite_err(e, "open_graph_writer"))
645 }
646
647 fn current_writer_task(
648 &self,
649 operation: &'static str,
650 ) -> Result<Option<WriterTaskHandle>, StorageError> {
651 self.pool
652 .writer_task_for_write(self.writer_task.as_ref(), operation)
653 }
654
655 async fn with_writer<F, R>(&self, op: &'static str, f: F) -> Result<R, StorageError>
662 where
663 F: FnOnce(&rusqlite::Connection) -> Result<R, rusqlite::Error> + Send + 'static,
664 R: Send + 'static,
665 {
666 if let Some(writer_task) = self.current_writer_task(op)? {
667 return writer_task
668 .send_bounded(move |conn| f(conn).map_err(|e| map_err(e, op)))
669 .await;
670 }
671
672 self.pool
673 .record_direct_route(crate::timeout_sink::Site::DirectRouteGraphGeneralWrite);
674 if self.is_file_backed {
675 let conn = self.open_standalone_writer()?;
676 let db = crate::timeout_sink::db_label(&self.pool);
677 tokio::task::spawn_blocking(move || {
678 f(&conn).map_err(|e| {
679 crate::timeout_sink::maybe_emit_busy(
680 &db,
681 crate::timeout_sink::Site::StandaloneGraph,
682 &e,
683 );
684 map_err(e, op)
685 })
686 })
687 .await
688 .map_err(|e| StorageError::driver(StorageCapability::Graph, op, e))?
689 } else {
690 let pool = Arc::clone(&self.pool);
691 tokio::task::spawn_blocking(move || {
692 let guard = pool.try_writer().map_err(|e| map_sqlite_err(e, op))?;
693 f(guard.conn()).map_err(|e| map_err(e, op))
694 })
695 .await
696 .map_err(|e| StorageError::driver(StorageCapability::Graph, op, e))?
697 }
698 }
699
700 async fn with_reader<F, R>(&self, op: &'static str, f: F) -> Result<R, StorageError>
701 where
702 F: FnOnce(&rusqlite::Connection) -> Result<R, rusqlite::Error> + Send + 'static,
703 R: Send + 'static,
704 {
705 if self.is_file_backed {
706 let pool = Arc::clone(&self.pool);
707 crate::read_cancellation::run_declared_interruptible_read(
708 StorageCapability::Graph,
709 op,
710 move |scope| {
711 scope.ensure_active()?;
712 let conn = pool
713 .open_standalone_reader()
714 .map_err(|error| map_sqlite_err(error, op))?;
715 scope.run(&conn, || f(&conn).map_err(|e| map_err(e, op)))
716 },
717 )
718 .await
719 } else {
720 let pool = Arc::clone(&self.pool);
721 crate::read_cancellation::run_declared_interruptible_read(
722 StorageCapability::Graph,
723 op,
724 move |scope| {
725 let mut guard = pool.resolve_reader_checkout(
726 StorageCapability::Graph,
727 op,
728 pool.reader_until(|| scope.should_stop()),
729 )?;
730 scope.run_pooled_reader(&mut guard, |conn| f(conn).map_err(|e| map_err(e, op)))
731 },
732 )
733 .await
734 }
735 }
736}
737
738fn report_graph_usage(queries: &std::sync::atomic::AtomicU64, rows: &std::sync::atomic::AtomicU64) {
754 khive_storage::usage::count(
755 khive_storage::usage::UsageUnit::DbRoundTrips,
756 queries.load(std::sync::atomic::Ordering::Relaxed),
757 );
758 khive_storage::usage::count(
759 khive_storage::usage::UsageUnit::GraphHops,
760 rows.load(std::sync::atomic::Ordering::Relaxed),
761 );
762}
763
764fn read_edge(row: &rusqlite::Row<'_>) -> Result<Edge, rusqlite::Error> {
765 let namespace: String = row.get(0)?;
766 let id_str: String = row.get(1)?;
767 let source_str: String = row.get(2)?;
768 let target_str: String = row.get(3)?;
769 let relation_str: String = row.get(4)?;
770 let weight: f64 = row.get(5)?;
771 let created_micros: i64 = row.get(6)?;
772 let updated_micros: i64 = row.get(7)?;
773 let deleted_micros: Option<i64> = row.get(8)?;
774 let metadata_str: Option<String> = row.get(9)?;
775 let target_backend: Option<String> = row.get(10)?;
776
777 let id = parse_uuid(&id_str)?;
778 let source_id = parse_uuid(&source_str)?;
779 let target_id = parse_uuid(&target_str)?;
780 let created_at = micros_to_datetime(created_micros);
781 let relation = relation_str.parse::<EdgeRelation>().map_err(|e| {
782 rusqlite::Error::FromSqlConversionFailure(4, rusqlite::types::Type::Text, Box::new(e))
783 })?;
784 let metadata = match metadata_str {
785 Some(s) => {
786 let v = serde_json::from_str(&s).map_err(|e| {
787 rusqlite::Error::FromSqlConversionFailure(
788 9,
789 rusqlite::types::Type::Text,
790 Box::new(e),
791 )
792 })?;
793 Some(v)
794 }
795 None => None,
796 };
797
798 Ok(Edge {
799 id: id.into(),
800 namespace,
801 source_id,
802 target_id,
803 relation,
804 weight,
805 created_at,
806 updated_at: micros_to_datetime(updated_micros),
807 deleted_at: deleted_micros.map(micros_to_datetime),
808 metadata,
809 target_backend,
810 })
811}
812
813fn parse_uuid(s: &str) -> Result<Uuid, rusqlite::Error> {
814 Uuid::parse_str(s).map_err(|e| {
815 rusqlite::Error::FromSqlConversionFailure(0, rusqlite::types::Type::Text, Box::new(e))
816 })
817}
818
819fn neighbor_extra_clause(
825 query: &NeighborQuery,
826 start_param_idx: usize,
827) -> (String, String, Vec<Box<dyn rusqlite::types::ToSql>>) {
828 let mut conditions: Vec<String> = Vec::new();
829 let mut extra_params: Vec<Box<dyn rusqlite::types::ToSql>> = Vec::new();
830 let mut param_idx = start_param_idx;
831
832 if let Some(ref rels) = query.relations {
833 if !rels.is_empty() {
834 let placeholders: Vec<String> = rels
835 .iter()
836 .map(|r| {
837 extra_params.push(Box::new(r.to_string()));
838 let p = format!("?{}", param_idx);
839 param_idx += 1;
840 p
841 })
842 .collect();
843 conditions.push(format!("relation IN ({})", placeholders.join(",")));
844 }
845 }
846
847 if let Some(min_w) = query.min_weight {
848 extra_params.push(Box::new(min_w));
849 conditions.push(format!("weight >= ?{}", param_idx));
850 param_idx += 1;
851 }
852
853 let where_extra = if conditions.is_empty() {
854 String::new()
855 } else {
856 format!(" WHERE {}", conditions.join(" AND "))
857 };
858
859 let limit_clause = if let Some(lim) = query.limit {
860 extra_params.push(Box::new(lim as i64));
861 format!(" LIMIT ?{}", param_idx)
862 } else {
863 String::new()
864 };
865
866 (where_extra, limit_clause, extra_params)
867}
868
869#[cfg(test)]
877static NEIGHBOR_SELECT_COUNT: std::sync::atomic::AtomicUsize =
878 std::sync::atomic::AtomicUsize::new(0);
879
880#[cfg(test)]
881fn count_neighbor_select() {
882 NEIGHBOR_SELECT_COUNT.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
883}
884
885#[cfg(not(test))]
886fn count_neighbor_select() {}
887
888#[cfg(test)]
889pub(crate) fn reset_neighbor_select_count() {
890 NEIGHBOR_SELECT_COUNT.store(0, std::sync::atomic::Ordering::Relaxed);
891}
892
893#[cfg(test)]
894pub(crate) fn neighbor_select_count() -> usize {
895 NEIGHBOR_SELECT_COUNT.load(std::sync::atomic::Ordering::Relaxed)
896}
897
898fn micros_to_datetime(micros: i64) -> DateTime<Utc> {
899 Utc.timestamp_micros(micros)
900 .single()
901 .unwrap_or_else(Utc::now)
902}
903
904fn edge_order_clause(sort: &[SortOrder<EdgeSortField>]) -> String {
911 if sort.is_empty() {
912 return " ORDER BY created_at DESC, id DESC".to_string();
913 }
914 let mut parts: Vec<String> = sort
915 .iter()
916 .map(|s| {
917 let dir = match s.direction {
918 SortDirection::Asc => "ASC",
919 SortDirection::Desc => "DESC",
920 };
921 format!("{} {}", edge_sort_col(&s.field), dir)
922 })
923 .collect();
924 let dir = match sort.last().map(|s| &s.direction) {
925 Some(SortDirection::Asc) => "ASC",
926 _ => "DESC",
927 };
928 parts.push(format!("id {dir}"));
929 format!(" ORDER BY {}", parts.join(", "))
930}
931
932fn build_edge_filter_sql(
933 namespace: &str,
934 filter: &EdgeFilter,
935) -> (String, Vec<Box<dyn rusqlite::types::ToSql>>) {
936 build_edge_filter_sql_for_namespaces(&[namespace.to_string()], filter)
937}
938
939fn build_edge_filter_sql_for_namespaces(
940 namespaces: &[String],
941 filter: &EdgeFilter,
942) -> (String, Vec<Box<dyn rusqlite::types::ToSql>>) {
943 let params: Vec<Box<dyn rusqlite::types::ToSql>> = namespaces
944 .iter()
945 .map(|namespace| -> Box<dyn rusqlite::types::ToSql> { Box::new(namespace.clone()) })
946 .collect();
947 let namespace_condition = match namespaces.len() {
948 0 => "0".to_string(),
949 1 => "namespace = ?1".to_string(),
950 _ => {
951 let placeholders: Vec<String> =
952 (1..=namespaces.len()).map(|i| format!("?{i}")).collect();
953 format!("namespace IN ({})", placeholders.join(", "))
954 }
955 };
956 build_edge_filter_conditions(namespace_condition, params, filter)
957}
958
959fn build_edge_filter_sql_for_namespaces_json(
974 namespaces_json: &str,
975 filter: &EdgeFilter,
976) -> (String, Vec<Box<dyn rusqlite::types::ToSql>>) {
977 let params: Vec<Box<dyn rusqlite::types::ToSql>> = vec![Box::new(namespaces_json.to_string())];
978 let namespace_condition = "namespace IN (SELECT value FROM json_each(?1))".to_string();
979 build_edge_filter_conditions(namespace_condition, params, filter)
980}
981
982fn build_edge_filter_conditions(
983 namespace_condition: String,
984 mut params: Vec<Box<dyn rusqlite::types::ToSql>>,
985 filter: &EdgeFilter,
986) -> (String, Vec<Box<dyn rusqlite::types::ToSql>>) {
987 let mut conditions = vec![namespace_condition, "deleted_at IS NULL".to_string()];
988
989 if !filter.ids.is_empty() {
990 let placeholders: Vec<String> = filter
991 .ids
992 .iter()
993 .map(|id| {
994 params.push(Box::new(id.to_string()));
995 format!("?{}", params.len())
996 })
997 .collect();
998 conditions.push(format!("id IN ({})", placeholders.join(",")));
999 }
1000
1001 if !filter.source_ids.is_empty() {
1002 let placeholders: Vec<String> = filter
1003 .source_ids
1004 .iter()
1005 .map(|id| {
1006 params.push(Box::new(id.to_string()));
1007 format!("?{}", params.len())
1008 })
1009 .collect();
1010 conditions.push(format!("source_id IN ({})", placeholders.join(",")));
1011 }
1012
1013 if !filter.target_ids.is_empty() {
1014 let placeholders: Vec<String> = filter
1015 .target_ids
1016 .iter()
1017 .map(|id| {
1018 params.push(Box::new(id.to_string()));
1019 format!("?{}", params.len())
1020 })
1021 .collect();
1022 conditions.push(format!("target_id IN ({})", placeholders.join(",")));
1023 }
1024
1025 if !filter.relations.is_empty() {
1026 let placeholders: Vec<String> = filter
1027 .relations
1028 .iter()
1029 .map(|r| {
1030 params.push(Box::new(r.to_string()));
1031 format!("?{}", params.len())
1032 })
1033 .collect();
1034 conditions.push(format!("relation IN ({})", placeholders.join(",")));
1035 }
1036
1037 if let Some(min_w) = filter.min_weight {
1038 params.push(Box::new(min_w));
1039 conditions.push(format!("weight >= ?{}", params.len()));
1040 }
1041
1042 if let Some(max_w) = filter.max_weight {
1043 params.push(Box::new(max_w));
1044 conditions.push(format!("weight <= ?{}", params.len()));
1045 }
1046
1047 if let Some(ref time_range) = filter.created_at {
1048 if let Some(start) = time_range.start {
1049 params.push(Box::new(start.timestamp_micros()));
1050 conditions.push(format!("created_at >= ?{}", params.len()));
1051 }
1052 if let Some(end) = time_range.end {
1053 params.push(Box::new(end.timestamp_micros()));
1054 conditions.push(format!("created_at < ?{}", params.len()));
1055 }
1056 }
1057
1058 let clause = format!(" WHERE {}", conditions.join(" AND "));
1059 (clause, params)
1060}
1061
1062fn edge_sort_col(field: &EdgeSortField) -> &'static str {
1063 match field {
1064 EdgeSortField::CreatedAt => "created_at",
1065 EdgeSortField::Weight => "weight",
1066 EdgeSortField::Relation => "relation",
1067 }
1068}
1069
1070fn canonical_edge_endpoints(
1079 relation: EdgeRelation,
1080 source_id: Uuid,
1081 target_id: Uuid,
1082) -> (Uuid, Uuid) {
1083 if relation.is_symmetric() && target_id < source_id {
1084 (target_id, source_id)
1085 } else {
1086 (source_id, target_id)
1087 }
1088}
1089
1090fn batch_upsert_edges(
1098 conn: &rusqlite::Connection,
1099 edges: &[Edge],
1100 attempted: u64,
1101) -> Result<BatchWriteSummary, rusqlite::Error> {
1102 let mut affected = 0u64;
1103
1104 for edge in edges {
1105 let statement = edge_upsert_statement(edge);
1106 let mut stmt = conn.prepare(&statement.sql)?;
1107 bind_params(&mut stmt, &statement.params)?;
1108 stmt.raw_execute()?;
1109 affected += 1;
1110 }
1111
1112 Ok(BatchWriteSummary {
1113 attempted,
1114 affected,
1115 failed: 0,
1116 first_error: String::new(),
1117 })
1118}
1119
1120fn edge_endpoints_exist(
1128 conn: &rusqlite::Connection,
1129 source_id: Uuid,
1130 target_id: Uuid,
1131) -> Result<MissingEndpoints, rusqlite::Error> {
1132 let src_exists = endpoint_exists_clause("?1");
1133 let tgt_exists = endpoint_exists_clause("?2");
1134 let sql = format!("SELECT ({src_exists}), ({tgt_exists})");
1135 conn.query_row(
1136 &sql,
1137 rusqlite::params![source_id.to_string(), target_id.to_string()],
1138 |row| {
1139 let src_exists: bool = row.get(0)?;
1140 let tgt_exists: bool = row.get(1)?;
1141 Ok(MissingEndpoints {
1142 source: !src_exists,
1143 target: !tgt_exists,
1144 })
1145 },
1146 )
1147}
1148
1149fn edge_insert_guarded(
1160 conn: &rusqlite::Connection,
1161 statement: &SqlStatement,
1162 source_id: Uuid,
1163 target_id: Uuid,
1164) -> Result<GuardedWriteOutcome, rusqlite::Error> {
1165 let mut stmt = conn.prepare(&statement.sql)?;
1166 bind_params(&mut stmt, &statement.params)?;
1167 if stmt.raw_execute()? > 0 {
1168 return Ok(GuardedWriteOutcome::Written);
1169 }
1170 #[cfg(test)]
1178 tests::insert_probe_seam::hook((source_id, target_id));
1179 let missing = edge_endpoints_exist(conn, source_id, target_id)?;
1180 Ok(GuardedWriteOutcome::Refused(missing))
1181}
1182
1183fn batch_upsert_edges_guarded(
1191 conn: &rusqlite::Connection,
1192 edges: &[Edge],
1193 attempted: u64,
1194) -> Result<GuardedBatchOutcome, rusqlite::Error> {
1195 for (index, edge) in edges.iter().enumerate() {
1196 let (source_id, target_id) =
1197 canonical_edge_endpoints(edge.relation, edge.source_id, edge.target_id);
1198 let missing = edge_endpoints_exist(conn, source_id, target_id)?;
1199 if missing.any() {
1200 return Ok(GuardedBatchOutcome {
1201 summary: BatchWriteSummary {
1202 attempted,
1203 affected: 0,
1204 failed: attempted,
1205 first_error: format!(
1206 "batch entry {index}: edge endpoint no longer exists at write time: source {source_id} or target {target_id}"
1207 ),
1208 },
1209 refused: Some(GuardedBatchRefusal {
1210 entry_index: index,
1211 missing,
1212 }),
1213 });
1214 }
1215 }
1216
1217 let mut affected = 0u64;
1218 for edge in edges {
1219 let statement = edge_upsert_statement(edge);
1220 let mut stmt = conn.prepare(&statement.sql)?;
1221 bind_params(&mut stmt, &statement.params)?;
1222 stmt.raw_execute()?;
1223 affected += 1;
1224 }
1225
1226 Ok(GuardedBatchOutcome {
1227 summary: BatchWriteSummary {
1228 attempted,
1229 affected,
1230 failed: 0,
1231 first_error: String::new(),
1232 },
1233 refused: None,
1234 })
1235}
1236
1237fn traversal_neighbor_sql(
1242 direction: Direction,
1243 relation_count: usize,
1244 has_min_weight: bool,
1245) -> String {
1246 let (node_column, endpoint_column, index) = match direction {
1247 Direction::Out => ("target_id", "source_id", "idx_graph_edges_ns_src_rel"),
1248 Direction::In => ("source_id", "target_id", "idx_graph_edges_ns_tgt_rel"),
1249 Direction::Both => unreachable!("Direction::Both is split into indexed Out/In seeks"),
1250 };
1251 let mut sql = format!(
1252 "SELECT {node_column}, id, weight \
1253 FROM graph_edges INDEXED BY {index} \
1254 WHERE namespace = ?1 AND {endpoint_column} = ?2 AND deleted_at IS NULL"
1255 );
1256 if relation_count > 0 {
1257 let placeholders = (0..relation_count)
1258 .map(|offset| format!("?{}", 4 + offset))
1259 .collect::<Vec<_>>()
1260 .join(",");
1261 sql.push_str(&format!(" AND relation IN ({placeholders})"));
1262 }
1263 if has_min_weight {
1264 sql.push_str(&format!(" AND weight >= ?{}", 4 + relation_count));
1265 }
1266 sql.push_str(" LIMIT ?3");
1267 sql
1268}
1269
1270fn traversal_timeout_error(budget: &TraversalExecutionBudget) -> StorageError {
1271 StorageError::Timeout {
1272 operation: format!(
1273 "traverse ({}ms execution budget)",
1274 budget.max_duration().as_millis()
1275 )
1276 .into(),
1277 }
1278}
1279
1280fn traversal_work_error(budget: &TraversalExecutionBudget) -> StorageError {
1281 StorageError::InvalidInput {
1282 capability: StorageCapability::Graph,
1283 operation: "traverse".into(),
1284 message: format!(
1285 "traversal work budget exceeded after {} adjacency rows; \
1286 narrow roots, depth, relations, or result limit",
1287 budget.work_limit()
1288 ),
1289 }
1290}
1291
1292#[derive(Clone, Copy)]
1293struct TraversalFrontierNode {
1294 node_id: Uuid,
1295 depth: usize,
1296 total_weight: f64,
1297}
1298
1299#[allow(clippy::too_many_arguments)]
1300fn run_bounded_traversal(
1301 conn: &rusqlite::Connection,
1302 roots: Vec<Uuid>,
1303 opts: TraversalOptions,
1304 include_roots: bool,
1305 namespace: String,
1306 origin: khive_storage::tx_registry::TxOrigin,
1307 budget: TraversalExecutionBudget,
1308 counted_rows: &std::sync::atomic::AtomicU64,
1309 counted_queries: &std::sync::atomic::AtomicU64,
1310) -> Result<Vec<GraphPath>, StorageError> {
1311 let progress_timed_out = Arc::new(std::sync::atomic::AtomicBool::new(false));
1312 let callback_timed_out = Arc::clone(&progress_timed_out);
1313 let callback_budget = budget.clone();
1314 #[cfg(test)]
1315 let progress_seam_root = roots.first().copied();
1316 conn.progress_handler(
1317 1_000,
1318 Some(move || {
1319 if crate::read_cancellation::current_read_should_interrupt() {
1320 return true;
1321 }
1322 #[cfg(test)]
1323 if tests::traverse_progress_seam::hook(progress_seam_root) {
1324 callback_timed_out.store(true, std::sync::atomic::Ordering::Relaxed);
1325 return true;
1326 }
1327 let expired = callback_budget.is_expired();
1328 if expired {
1329 callback_timed_out.store(true, std::sync::atomic::Ordering::Relaxed);
1330 }
1331 expired
1332 }),
1333 )
1334 .map_err(|e| map_err(e, "traverse_progress_handler"))?;
1335
1336 let result = (|| {
1337 let result_limit = opts.effective_limit() as usize;
1338 let relation_count = opts.relations.as_ref().map_or(0, Vec::len);
1339 let directions = match opts.direction {
1340 Direction::Out => vec![Direction::Out],
1341 Direction::In => vec![Direction::In],
1342 Direction::Both => vec![Direction::Out, Direction::In],
1343 };
1344 let statements = directions
1345 .into_iter()
1346 .map(|direction| {
1347 traversal_neighbor_sql(direction, relation_count, opts.min_weight.is_some())
1348 })
1349 .collect::<Vec<_>>();
1350 let map_sql_error = |error| {
1351 if progress_timed_out.load(std::sync::atomic::Ordering::Relaxed) {
1352 traversal_timeout_error(&budget)
1353 } else {
1354 map_err(error, "traverse")
1355 }
1356 };
1357
1358 let mut all_paths = Vec::with_capacity(roots.len());
1359 for root_id in roots {
1360 let mut seen = HashSet::new();
1361 seen.insert(root_id);
1362 let mut frontier = VecDeque::from([TraversalFrontierNode {
1363 node_id: root_id,
1364 depth: 0,
1365 total_weight: 0.0,
1366 }]);
1367 let mut nodes = Vec::with_capacity(result_limit + usize::from(include_roots));
1368 if include_roots {
1369 nodes.push(PathNode {
1370 node_id: root_id,
1371 via_edge: None,
1372 depth: 0,
1373 name: None,
1374 kind: None,
1375 properties: None,
1376 weight: 0.0,
1377 });
1378 }
1379 let mut non_root_count = 0usize;
1380
1381 'root_walk: while non_root_count < result_limit {
1382 let Some(current) = frontier.pop_front() else {
1383 break;
1384 };
1385 if current.depth >= opts.max_depth {
1386 continue;
1387 }
1388 if budget.is_expired() {
1389 return Err(traversal_timeout_error(&budget));
1390 }
1391
1392 for sql in &statements {
1393 let row_cap = budget.remaining_work().saturating_add(1);
1394 let mut params: Vec<Box<dyn rusqlite::types::ToSql>> = vec![
1395 Box::new(namespace.clone()),
1396 Box::new(current.node_id.to_string()),
1397 Box::new(row_cap as i64),
1398 ];
1399 if let Some(relations) = &opts.relations {
1400 params.extend(relations.iter().map(|relation| {
1401 Box::new(relation.to_string()) as Box<dyn rusqlite::types::ToSql>
1402 }));
1403 }
1404 if let Some(min_weight) = opts.min_weight {
1405 params.push(Box::new(min_weight));
1406 }
1407 let param_refs = params
1408 .iter()
1409 .map(|param| param.as_ref())
1410 .collect::<Vec<&dyn rusqlite::types::ToSql>>();
1411
1412 let _snapshot = khive_storage::tx_registry::register_scoped(
1413 Some("graph_traverse_read".to_string()),
1414 origin.clone(),
1415 );
1416 let mut stmt = conn.prepare(sql).map_err(&map_sql_error)?;
1417 counted_queries.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1418 let mut rows = stmt.query(param_refs.as_slice()).map_err(&map_sql_error)?;
1419 while let Some(row) = rows.next().map_err(&map_sql_error)? {
1420 #[cfg(test)]
1423 tests::traverse_snapshot_seam::hook(current.node_id);
1424 counted_rows.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1425 if budget.is_expired() {
1426 return Err(traversal_timeout_error(&budget));
1427 }
1428 if !budget.try_consume_row() {
1429 return Err(traversal_work_error(&budget));
1430 }
1431 let node_str: String = row.get(0).map_err(&map_sql_error)?;
1432 let edge_str: String = row.get(1).map_err(&map_sql_error)?;
1433 let edge_weight: f64 = row.get(2).map_err(&map_sql_error)?;
1434 let node_id = parse_uuid(&node_str).map_err(&map_sql_error)?;
1435 if !seen.insert(node_id) {
1436 continue;
1437 }
1438 let via_edge = parse_uuid(&edge_str).map_err(&map_sql_error)?;
1439 let depth = current.depth + 1;
1440 let total_weight = current.total_weight + edge_weight;
1441 nodes.push(PathNode {
1442 node_id,
1443 via_edge: Some(via_edge),
1444 depth,
1445 name: None,
1446 kind: None,
1447 properties: None,
1448 weight: total_weight,
1449 });
1450 non_root_count += 1;
1451 if depth < opts.max_depth {
1452 frontier.push_back(TraversalFrontierNode {
1453 node_id,
1454 depth,
1455 total_weight,
1456 });
1457 }
1458 if non_root_count == result_limit {
1459 break 'root_walk;
1460 }
1461 }
1462 }
1463 }
1464
1465 if !nodes.is_empty() {
1466 let total_weight = nodes.iter().map(|node| node.weight).fold(0.0_f64, f64::max);
1467 all_paths.push(GraphPath {
1468 root_id,
1469 nodes,
1470 total_weight,
1471 });
1472 }
1473 }
1474 Ok(all_paths)
1475 })();
1476
1477 conn.progress_handler(0, None::<fn() -> bool>)
1478 .map_err(|e| map_err(e, "traverse_progress_handler_clear"))?;
1479 result
1480}
1481
1482#[async_trait]
1483impl GraphStore for SqlGraphStore {
1484 async fn upsert_edge(&self, edge: Edge) -> Result<(), StorageError> {
1485 let statement = edge_upsert_statement(&edge);
1486 self.with_writer("upsert_edge", move |conn| {
1487 let mut stmt = conn.prepare(&statement.sql)?;
1488 bind_params(&mut stmt, &statement.params)?;
1489 stmt.raw_execute()?;
1490 Ok(())
1491 })
1492 .await
1493 }
1494
1495 async fn replace_edge_if_unchanged(
1496 &self,
1497 edge: Edge,
1498 expected_updated_at: DateTime<Utc>,
1499 expected_deleted_at: Option<DateTime<Utc>>,
1500 ) -> Result<bool, StorageError> {
1501 let statement =
1502 edge_replace_if_unchanged_statement(&edge, expected_updated_at, expected_deleted_at);
1503 self.with_writer("replace_edge_if_unchanged", move |conn| {
1504 let mut stmt = conn.prepare(&statement.sql)?;
1505 bind_params(&mut stmt, &statement.params)?;
1506 Ok(stmt.raw_execute()? > 0)
1507 })
1508 .await
1509 }
1510
1511 async fn upsert_edges(&self, edges: Vec<Edge>) -> Result<BatchWriteSummary, StorageError> {
1512 let attempted = edges.len() as u64;
1513
1514 if let Some(writer_task) = self.current_writer_task("upsert_edges")? {
1520 return writer_task
1521 .send_bounded(move |conn| {
1522 batch_upsert_edges(conn, &edges, attempted)
1523 .map_err(|e| map_err(e, "upsert_edges"))
1524 })
1525 .await;
1526 }
1527
1528 let origin = self.pool.origin();
1532 self.with_writer("upsert_edges", move |conn| {
1533 conn.execute_batch("BEGIN IMMEDIATE")?;
1534 let _tx_handle = khive_storage::tx_registry::register_scoped(
1535 Some("graph_upsert_edges".to_string()),
1536 origin,
1537 );
1538
1539 let summary = match batch_upsert_edges(conn, &edges, attempted) {
1540 Ok(summary) => summary,
1541 Err(e) => {
1542 let _ = conn.execute_batch("ROLLBACK");
1543 return Err(e);
1544 }
1545 };
1546
1547 if let Err(e) = conn.execute_batch("COMMIT") {
1548 let _ = conn.execute_batch("ROLLBACK");
1549 return Err(e);
1550 }
1551 Ok(summary)
1552 })
1553 .await
1554 }
1555
1556 async fn upsert_edge_guarded(&self, edge: Edge) -> Result<GuardedWriteOutcome, StorageError> {
1557 let (source_id, target_id) =
1558 canonical_edge_endpoints(edge.relation, edge.source_id, edge.target_id);
1559 let metadata_str = edge
1560 .metadata
1561 .as_ref()
1562 .map(|v| serde_json::to_string(v).unwrap_or_default());
1563 let statement = edge_insert_guarded_by_endpoints_statement(
1564 &edge.namespace,
1565 Uuid::from(edge.id),
1566 source_id,
1567 target_id,
1568 edge.relation,
1569 edge.weight,
1570 edge.created_at.timestamp_micros(),
1571 metadata_str.as_deref(),
1572 );
1573
1574 if let Some(writer_task) = self.current_writer_task("upsert_edge_guarded")? {
1580 return writer_task
1581 .send_bounded(move |conn| {
1582 edge_insert_guarded(conn, &statement, source_id, target_id)
1583 .map_err(|e| map_err(e, "upsert_edge_guarded"))
1584 })
1585 .await;
1586 }
1587
1588 let origin = self.pool.origin();
1594 self.with_writer("upsert_edge_guarded", move |conn| {
1595 conn.execute_batch("BEGIN IMMEDIATE")?;
1596 let _tx_handle = khive_storage::tx_registry::register_scoped(
1597 Some("graph_upsert_edge_guarded".to_string()),
1598 origin,
1599 );
1600
1601 let outcome = match edge_insert_guarded(conn, &statement, source_id, target_id) {
1602 Ok(outcome) => outcome,
1603 Err(e) => {
1604 let _ = conn.execute_batch("ROLLBACK");
1605 return Err(e);
1606 }
1607 };
1608
1609 if let Err(e) = conn.execute_batch("COMMIT") {
1610 let _ = conn.execute_batch("ROLLBACK");
1611 return Err(e);
1612 }
1613 Ok(outcome)
1614 })
1615 .await
1616 }
1617
1618 async fn upsert_edges_guarded(
1619 &self,
1620 edges: Vec<Edge>,
1621 ) -> Result<GuardedBatchOutcome, StorageError> {
1622 let attempted = edges.len() as u64;
1623
1624 if let Some(writer_task) = self.current_writer_task("upsert_edges_guarded")? {
1628 return writer_task
1629 .send_bounded(move |conn| {
1630 batch_upsert_edges_guarded(conn, &edges, attempted)
1631 .map_err(|e| map_err(e, "upsert_edges_guarded"))
1632 })
1633 .await;
1634 }
1635
1636 let origin = self.pool.origin();
1637 self.with_writer("upsert_edges_guarded", move |conn| {
1638 conn.execute_batch("BEGIN IMMEDIATE")?;
1639 let _tx_handle = khive_storage::tx_registry::register_scoped(
1640 Some("graph_upsert_edges_guarded".to_string()),
1641 origin,
1642 );
1643
1644 let summary = match batch_upsert_edges_guarded(conn, &edges, attempted) {
1645 Ok(summary) => summary,
1646 Err(e) => {
1647 let _ = conn.execute_batch("ROLLBACK");
1648 return Err(e);
1649 }
1650 };
1651
1652 if let Err(e) = conn.execute_batch("COMMIT") {
1653 let _ = conn.execute_batch("ROLLBACK");
1654 return Err(e);
1655 }
1656 Ok(summary)
1657 })
1658 .await
1659 }
1660
1661 async fn get_edge(&self, id: LinkId) -> Result<Option<Edge>, StorageError> {
1662 let id_str = Uuid::from(id).to_string();
1663
1664 self.with_reader("get_edge", move |conn| {
1665 let mut stmt = conn.prepare(
1666 "SELECT namespace, id, source_id, target_id, relation, weight, \
1667 created_at, updated_at, deleted_at, metadata, target_backend \
1668 FROM graph_edges WHERE id = ?1 AND deleted_at IS NULL",
1669 )?;
1670 let mut rows = stmt.query(rusqlite::params![id_str])?;
1671 match rows.next()? {
1672 Some(row) => Ok(Some(read_edge(row)?)),
1673 None => Ok(None),
1674 }
1675 })
1676 .await
1677 }
1678
1679 async fn get_edge_including_deleted(&self, id: LinkId) -> Result<Option<Edge>, StorageError> {
1680 let id_str = Uuid::from(id).to_string();
1681
1682 self.with_reader("get_edge_including_deleted", move |conn| {
1683 let mut stmt = conn.prepare(
1684 "SELECT namespace, id, source_id, target_id, relation, weight, \
1685 created_at, updated_at, deleted_at, metadata, target_backend \
1686 FROM graph_edges WHERE id = ?1",
1687 )?;
1688 let mut rows = stmt.query(rusqlite::params![id_str])?;
1689 match rows.next()? {
1690 Some(row) => Ok(Some(read_edge(row)?)),
1691 None => Ok(None),
1692 }
1693 })
1694 .await
1695 }
1696
1697 async fn edge_sequence(&self, id: Uuid) -> Result<Option<i64>, StorageError> {
1698 let id = id.to_string();
1699 self.with_reader("edge_sequence", move |conn| {
1700 conn.query_row(
1701 "SELECT seq FROM graph_edges_seq WHERE edge_id = ?1",
1702 rusqlite::params![id],
1703 |row| row.get(0),
1704 )
1705 .optional()
1706 })
1707 .await
1708 }
1709
1710 async fn edge_sequences(&self, ids: &[Uuid]) -> Result<Vec<(Uuid, i64)>, StorageError> {
1711 if ids.is_empty() {
1712 return Ok(Vec::new());
1713 }
1714 let ids = ids.to_vec();
1715 self.with_reader("edge_sequences", move |conn| {
1716 const CHUNK: usize = 900;
1717 let mut resolved = Vec::with_capacity(ids.len());
1718 for chunk in ids.chunks(CHUNK) {
1719 let placeholders = (1..=chunk.len())
1720 .map(|index| format!("?{index}"))
1721 .collect::<Vec<_>>()
1722 .join(", ");
1723 let sql = format!(
1724 "SELECT edge_id, seq FROM graph_edges_seq WHERE edge_id IN ({placeholders})"
1725 );
1726 let strings = chunk.iter().map(Uuid::to_string).collect::<Vec<_>>();
1727 let params = strings
1728 .iter()
1729 .map(|id| id as &dyn rusqlite::types::ToSql)
1730 .collect::<Vec<_>>();
1731 let mut stmt = conn.prepare(&sql)?;
1732 let rows = stmt.query_map(params.as_slice(), |row| {
1733 let id: String = row.get(0)?;
1734 Ok((parse_uuid(&id)?, row.get::<_, i64>(1)?))
1735 })?;
1736 resolved.extend(rows.collect::<Result<Vec<_>, _>>()?);
1737 }
1738 Ok(resolved)
1739 })
1740 .await
1741 }
1742
1743 async fn get_edge_by_natural_key_including_deleted(
1744 &self,
1745 namespace: &str,
1746 source_id: Uuid,
1747 target_id: Uuid,
1748 relation: EdgeRelation,
1749 ) -> Result<Option<Edge>, StorageError> {
1750 let namespace = namespace.to_string();
1751 let source_str = source_id.to_string();
1752 let target_str = target_id.to_string();
1753 let relation_str = relation.to_string();
1754
1755 self.with_reader("get_edge_by_natural_key_including_deleted", move |conn| {
1756 let mut stmt = conn.prepare(
1757 "SELECT namespace, id, source_id, target_id, relation, weight, \
1758 created_at, updated_at, deleted_at, metadata, target_backend \
1759 FROM graph_edges \
1760 WHERE namespace = ?1 AND source_id = ?2 AND target_id = ?3 AND relation = ?4",
1761 )?;
1762 let mut rows = stmt.query(rusqlite::params![
1763 namespace,
1764 source_str,
1765 target_str,
1766 relation_str
1767 ])?;
1768 match rows.next()? {
1769 Some(row) => Ok(Some(read_edge(row)?)),
1770 None => Ok(None),
1771 }
1772 })
1773 .await
1774 }
1775
1776 async fn get_edges(&self, ids: &[LinkId]) -> Result<Vec<Edge>, StorageError> {
1777 if ids.is_empty() {
1778 return Ok(Vec::new());
1779 }
1780 const CHUNK: usize = 900;
1782 let id_strs: Vec<String> = ids.iter().map(|id| Uuid::from(*id).to_string()).collect();
1783
1784 let mut result: Vec<Edge> = Vec::with_capacity(ids.len());
1785 for chunk in id_strs.chunks(CHUNK) {
1786 let chunk_owned: Vec<String> = chunk.to_vec();
1787 let edges = self
1788 .with_reader("get_edges", move |conn| {
1789 let placeholders: Vec<String> =
1790 (1..=chunk_owned.len()).map(|i| format!("?{}", i)).collect();
1791 let sql = format!(
1792 "SELECT namespace, id, source_id, target_id, relation, weight, \
1793 created_at, updated_at, deleted_at, metadata, target_backend \
1794 FROM graph_edges WHERE id IN ({}) AND deleted_at IS NULL",
1795 placeholders.join(",")
1796 );
1797 let mut stmt = conn.prepare(&sql)?;
1798 let params: Vec<&dyn rusqlite::types::ToSql> = chunk_owned
1799 .iter()
1800 .map(|s| s as &dyn rusqlite::types::ToSql)
1801 .collect();
1802 let rows = stmt.query_map(params.as_slice(), read_edge)?;
1803 let mut edges = Vec::new();
1804 for row in rows {
1805 edges.push(row?);
1806 }
1807 Ok(edges)
1808 })
1809 .await?;
1810 result.extend(edges);
1811 }
1812 Ok(result)
1813 }
1814
1815 async fn batch_neighbors(
1816 &self,
1817 sources: &[Uuid],
1818 query: NeighborQuery,
1819 ) -> Result<Vec<(Uuid, NeighborHit)>, StorageError> {
1820 use khive_storage::types::Direction;
1821
1822 if sources.is_empty() {
1823 return Ok(Vec::new());
1824 }
1825 let mut seen_sources = HashSet::with_capacity(sources.len());
1826 let unique_sources: Vec<Uuid> = sources
1827 .iter()
1828 .copied()
1829 .filter(|source| seen_sources.insert(*source))
1830 .collect();
1831 const CHUNK_SIZE: usize = 880;
1832
1833 let namespace = self.namespace.clone();
1834 let mut result: Vec<(Uuid, NeighborHit)> = Vec::new();
1835 let counted_queries = Arc::new(std::sync::atomic::AtomicU64::new(0));
1838 let counted_rows = Arc::new(std::sync::atomic::AtomicU64::new(0));
1839
1840 for chunk in unique_sources.chunks(CHUNK_SIZE) {
1841 let chunk_owned: Vec<Uuid> = chunk.to_vec();
1842 let query_clone = query.clone();
1843 let ns = namespace.clone();
1844 let closure_queries = Arc::clone(&counted_queries);
1845 let closure_rows = Arc::clone(&counted_rows);
1846
1847 let chunk_result = self
1848 .with_reader("batch_neighbors", move |conn| {
1849 let src_strs: Vec<String> = chunk_owned.iter().map(|u| u.to_string()).collect();
1850
1851 let sources_json = serde_json::to_string(&src_strs).map_err(|error| {
1852 rusqlite::Error::ToSqlConversionFailure(Box::new(error))
1853 })?;
1854
1855 let build_inner_sql =
1856 |direction_out: bool,
1857 q: &NeighborQuery|
1858 -> (String, Vec<String>, Option<f64>) {
1859 let (filter_col, node_col) = if direction_out {
1860 ("source_id", "target_id")
1861 } else {
1862 ("target_id", "source_id")
1863 };
1864
1865 let mut rel_params: Vec<String> = Vec::new();
1866 let mut conditions: Vec<String> = Vec::new();
1867 let mut param_idx = 3;
1868
1869 if let Some(ref rels) = q.relations {
1870 if !rels.is_empty() {
1871 let ps: Vec<String> = rels
1872 .iter()
1873 .map(|r| {
1874 rel_params.push(r.to_string());
1875 let p = format!("?{param_idx}");
1876 param_idx += 1;
1877 p
1878 })
1879 .collect();
1880 conditions
1881 .push(format!("edges.relation IN ({})", ps.join(",")));
1882 }
1883 }
1884
1885 let min_weight_val = if let Some(min_w) = q.min_weight {
1888 conditions.push(format!("edges.weight >= ?{param_idx}"));
1889 Some(min_w)
1890 } else {
1891 None
1892 };
1893
1894 let where_extra = if conditions.is_empty() {
1895 String::new()
1896 } else {
1897 format!(" AND {}", conditions.join(" AND "))
1898 };
1899
1900 let sql = format!(
1901 "SELECT requested.origin_id, edges.{node_col} AS node_id, \
1902 edges.id AS edge_id, edges.relation, edges.weight \
1903 FROM requested CROSS JOIN graph_edges AS edges \
1904 ON edges.{filter_col} = requested.origin_id \
1905 WHERE edges.namespace = ?1 \
1906 AND edges.deleted_at IS NULL{where_extra}",
1907 );
1908 (sql, rel_params, min_weight_val)
1909 };
1910
1911 let mut all_params: Vec<Box<dyn rusqlite::types::ToSql>> = Vec::new();
1912 all_params.push(Box::new(ns.to_string()));
1913 all_params.push(Box::new(sources_json));
1914
1915 let (combined_inner, rel_params, min_weight_val) = match query_clone.direction {
1916 Direction::Out => build_inner_sql(true, &query_clone),
1917 Direction::In => build_inner_sql(false, &query_clone),
1918 Direction::Both => {
1919 let (out_sql, rel_params, min_weight_val) =
1920 build_inner_sql(true, &query_clone);
1921 let (in_sql, _, _) = build_inner_sql(false, &query_clone);
1922 (
1923 format!("{out_sql} UNION ALL {in_sql}"),
1924 rel_params,
1925 min_weight_val,
1926 )
1927 }
1928 };
1929
1930 for relation in rel_params {
1931 all_params.push(Box::new(relation));
1932 }
1933 if let Some(min_weight) = min_weight_val {
1934 all_params.push(Box::new(min_weight));
1935 }
1936 let limit_param_idx = all_params.len() + 1;
1937
1938 let full_sql = if let Some(lim) = query_clone.limit {
1946 all_params.push(Box::new(lim as i64));
1947 format!(
1948 "WITH requested(origin_id) AS (\
1949 SELECT value FROM json_each(?2)\
1950 ) SELECT origin_id, node_id, edge_id, relation, weight \
1951 FROM (SELECT *, ROW_NUMBER() OVER (PARTITION BY origin_id \
1952 ORDER BY weight DESC, node_id ASC) AS rn \
1953 FROM ({combined_inner})) WHERE rn <= ?{limit_param_idx}",
1954 )
1955 } else {
1956 format!(
1957 "WITH requested(origin_id) AS (\
1958 SELECT value FROM json_each(?2)\
1959 ) SELECT origin_id, node_id, edge_id, relation, weight \
1960 FROM ({combined_inner})",
1961 )
1962 };
1963
1964 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
1965 all_params.iter().map(|p| p.as_ref()).collect();
1966
1967 let mut stmt = conn.prepare(&full_sql)?;
1968 closure_queries.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1969 let rows = stmt.query_map(param_refs.as_slice(), |row| {
1970 let origin_str: String = row.get(0)?;
1971 let nid_str: String = row.get(1)?;
1972 let eid_str: String = row.get(2)?;
1973 let relation_str: String = row.get(3)?;
1974 let weight: f64 = row.get(4)?;
1975 Ok((origin_str, nid_str, eid_str, relation_str, weight))
1976 })?;
1977
1978 let mut pairs = Vec::new();
1979 for row in rows {
1980 let (origin_str, nid_str, eid_str, relation_str, weight) = row?;
1981 closure_rows.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1982 let origin = parse_uuid(&origin_str)?;
1983 let node_id = parse_uuid(&nid_str)?;
1984 let edge_id = parse_uuid(&eid_str)?;
1985 let relation = relation_str.parse::<EdgeRelation>().map_err(|e| {
1986 rusqlite::Error::FromSqlConversionFailure(
1987 3,
1988 rusqlite::types::Type::Text,
1989 Box::new(e),
1990 )
1991 })?;
1992 pairs.push((
1993 origin,
1994 NeighborHit {
1995 node_id,
1996 edge_id,
1997 relation,
1998 weight,
1999 name: None,
2000 kind: None,
2001 entity_type: None,
2002 },
2003 ));
2004 }
2005 Ok(pairs)
2006 })
2007 .await;
2008 let pairs = match chunk_result {
2009 Ok(pairs) => pairs,
2010 Err(e) => {
2011 report_graph_usage(&counted_queries, &counted_rows);
2012 return Err(e);
2013 }
2014 };
2015 result.extend(pairs);
2016 }
2017 report_graph_usage(&counted_queries, &counted_rows);
2018
2019 let requested: HashSet<Uuid> = unique_sources.iter().copied().collect();
2020 let mut grouped: HashMap<Uuid, Vec<NeighborHit>> =
2021 HashMap::with_capacity(unique_sources.len());
2022 for (origin, hit) in result {
2023 if !requested.contains(&origin) {
2024 return Err(StorageError::Internal(format!(
2025 "batch_neighbors returned unrequested origin {origin}"
2026 )));
2027 }
2028 grouped.entry(origin).or_default().push(hit);
2029 }
2030
2031 for hits in grouped.values_mut() {
2032 hits.sort_by(|a, b| {
2033 b.weight
2034 .partial_cmp(&a.weight)
2035 .unwrap_or(std::cmp::Ordering::Equal)
2036 .then(a.node_id.cmp(&b.node_id))
2037 .then(a.edge_id.cmp(&b.edge_id))
2038 });
2039 }
2040
2041 let mut ordered = Vec::new();
2042 for &source in sources {
2043 if let Some(hits) = grouped.get(&source) {
2044 ordered.extend(hits.iter().cloned().map(|hit| (source, hit)));
2045 }
2046 }
2047 Ok(ordered)
2048 }
2049
2050 async fn delete_edge(&self, id: LinkId, mode: DeleteMode) -> Result<bool, StorageError> {
2051 let id = Uuid::from(id);
2052 let statement = match mode {
2053 DeleteMode::Soft => {
2054 edge_soft_delete_statement(id, chrono::Utc::now().timestamp_micros())
2055 }
2056 DeleteMode::Hard => edge_hard_delete_statement(id),
2057 };
2058 self.with_writer("delete_edge", move |conn| {
2059 let mut stmt = conn.prepare(&statement.sql)?;
2060 bind_params(&mut stmt, &statement.params)?;
2061 Ok(stmt.raw_execute()? > 0)
2062 })
2063 .await
2064 }
2065
2066 async fn query_edges(
2067 &self,
2068 filter: EdgeFilter,
2069 sort: Vec<SortOrder<EdgeSortField>>,
2070 page: PageRequest,
2071 ) -> Result<Page<Edge>, StorageError> {
2072 let namespace = self.namespace.clone();
2073 let limit_i64 = i64::from(page.limit);
2074 let offset_i64 = i64::try_from(page.offset).map_err(|_| StorageError::InvalidInput {
2075 capability: StorageCapability::Graph,
2076 operation: "query_edges".into(),
2077 message: format!(
2078 "PageRequest: offset must be <= i64::MAX, got {}",
2079 page.offset
2080 ),
2081 })?;
2082 self.with_reader("query_edges", move |conn| {
2083 let (where_clause, filter_params) = build_edge_filter_sql(&namespace, &filter);
2084
2085 let count_sql = format!("SELECT COUNT(*) FROM graph_edges{}", where_clause);
2086 let total: i64 = {
2087 let mut stmt = conn.prepare(&count_sql)?;
2088 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
2089 filter_params.iter().map(|p| p.as_ref()).collect();
2090 stmt.query_row(param_refs.as_slice(), |row| row.get(0))?
2091 };
2092
2093 let order_clause = edge_order_clause(&sort);
2094
2095 let (_, data_filter_params) = build_edge_filter_sql(&namespace, &filter);
2096 let mut all_params: Vec<Box<dyn rusqlite::types::ToSql>> = data_filter_params;
2097 all_params.push(Box::new(limit_i64));
2098 all_params.push(Box::new(offset_i64));
2099
2100 let limit_idx = all_params.len() - 1;
2101 let offset_idx = all_params.len();
2102
2103 let data_sql = format!(
2104 "SELECT namespace, id, source_id, target_id, relation, weight, \
2105 created_at, updated_at, deleted_at, metadata, target_backend \
2106 FROM graph_edges{}{} LIMIT ?{} OFFSET ?{}",
2107 where_clause, order_clause, limit_idx, offset_idx,
2108 );
2109
2110 let mut stmt = conn.prepare(&data_sql)?;
2111 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
2112 all_params.iter().map(|p| p.as_ref()).collect();
2113 let rows = stmt.query_map(param_refs.as_slice(), read_edge)?;
2114
2115 let mut items = Vec::new();
2116 for row in rows {
2117 items.push(row?);
2118 }
2119
2120 Ok(Page {
2121 items,
2122 total: Some(total as u64),
2123 })
2124 })
2125 .await
2126 }
2127
2128 async fn count_edges(&self, filter: EdgeFilter) -> Result<u64, StorageError> {
2129 let namespace = self.namespace.clone();
2130 self.with_reader("count_edges", move |conn| {
2131 let (where_clause, params) = build_edge_filter_sql(&namespace, &filter);
2132 let sql = format!("SELECT COUNT(*) FROM graph_edges{}", where_clause);
2133 let mut stmt = conn.prepare(&sql)?;
2134 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
2135 params.iter().map(|p| p.as_ref()).collect();
2136 let count: i64 = stmt.query_row(param_refs.as_slice(), |row| row.get(0))?;
2137 Ok(count as u64)
2138 })
2139 .await
2140 }
2141
2142 async fn count_edges_in_namespaces(
2143 &self,
2144 namespaces: &[String],
2145 filter: EdgeFilter,
2146 ) -> Result<u64, StorageError> {
2147 let namespaces: Vec<String> = namespaces
2148 .iter()
2149 .cloned()
2150 .collect::<HashSet<_>>()
2151 .into_iter()
2152 .collect();
2153 self.with_reader("count_edges_in_namespaces", move |conn| {
2154 let mut total = 0;
2155 for chunk in namespaces.chunks(NAMESPACE_COUNT_CHUNK_SIZE) {
2156 let (where_clause, params) = build_edge_filter_sql_for_namespaces(chunk, &filter);
2157 let sql = format!("SELECT COUNT(*) FROM graph_edges{where_clause}");
2158 let mut stmt = conn.prepare(&sql)?;
2159 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
2160 params.iter().map(|p| p.as_ref()).collect();
2161 let count: i64 = stmt.query_row(param_refs.as_slice(), |row| row.get(0))?;
2162 total += count as u64;
2163 }
2164 Ok(total)
2165 })
2166 .await
2167 }
2168
2169 async fn query_edges_in_namespaces(
2170 &self,
2171 namespaces: &[String],
2172 filter: EdgeFilter,
2173 sort: Vec<SortOrder<EdgeSortField>>,
2174 page: PageRequest,
2175 ) -> Result<Page<Edge>, StorageError> {
2176 let namespaces: Vec<String> = {
2191 let mut seen = HashSet::new();
2192 namespaces
2193 .iter()
2194 .filter(|ns| seen.insert((*ns).clone()))
2195 .cloned()
2196 .collect()
2197 };
2198 let limit_i64 = i64::from(page.limit);
2199 let offset_i64 = i64::try_from(page.offset).map_err(|_| StorageError::InvalidInput {
2200 capability: StorageCapability::Graph,
2201 operation: "query_edges_in_namespaces".into(),
2202 message: format!(
2203 "PageRequest: offset must be <= i64::MAX, got {}",
2204 page.offset
2205 ),
2206 })?;
2207 self.with_reader("query_edges_in_namespaces", move |conn| {
2208 let namespaces_json = serde_json::to_string(&namespaces)
2209 .map_err(|error| rusqlite::Error::ToSqlConversionFailure(Box::new(error)))?;
2210
2211 let (where_clause, filter_params) =
2212 build_edge_filter_sql_for_namespaces_json(&namespaces_json, &filter);
2213
2214 let count_sql = format!("SELECT COUNT(*) FROM graph_edges{}", where_clause);
2215 let total: i64 = {
2216 let mut stmt = conn.prepare(&count_sql)?;
2217 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
2218 filter_params.iter().map(|p| p.as_ref()).collect();
2219 stmt.query_row(param_refs.as_slice(), |row| row.get(0))?
2220 };
2221
2222 let order_clause = edge_order_clause(&sort);
2223
2224 let (_, data_filter_params) =
2225 build_edge_filter_sql_for_namespaces_json(&namespaces_json, &filter);
2226 let mut all_params: Vec<Box<dyn rusqlite::types::ToSql>> = data_filter_params;
2227 all_params.push(Box::new(limit_i64));
2228 all_params.push(Box::new(offset_i64));
2229
2230 let limit_idx = all_params.len() - 1;
2231 let offset_idx = all_params.len();
2232
2233 let data_sql = format!(
2234 "SELECT namespace, id, source_id, target_id, relation, weight, \
2235 created_at, updated_at, deleted_at, metadata, target_backend \
2236 FROM graph_edges{}{} LIMIT ?{} OFFSET ?{}",
2237 where_clause, order_clause, limit_idx, offset_idx,
2238 );
2239
2240 let mut stmt = conn.prepare(&data_sql)?;
2241 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
2242 all_params.iter().map(|p| p.as_ref()).collect();
2243 let rows = stmt.query_map(param_refs.as_slice(), read_edge)?;
2244
2245 let mut items = Vec::new();
2246 for row in rows {
2247 items.push(row?);
2248 }
2249
2250 Ok(Page {
2251 items,
2252 total: Some(total as u64),
2253 })
2254 })
2255 .await
2256 }
2257
2258 async fn count_edges_by_relation(&self) -> Result<Vec<(EdgeRelation, u64)>, StorageError> {
2259 let namespace = self.namespace.clone();
2260 self.with_reader("count_edges_by_relation", move |conn| {
2261 let sql = "SELECT relation, COUNT(*) FROM graph_edges \
2262 WHERE namespace = ?1 AND deleted_at IS NULL \
2263 GROUP BY relation";
2264 let mut stmt = conn.prepare(sql)?;
2265 let rows = stmt.query_map([&namespace], |row| {
2266 let relation_str: String = row.get(0)?;
2267 let count: i64 = row.get(1)?;
2268 Ok((relation_str, count))
2269 })?;
2270 let mut out = Vec::new();
2271 for row in rows {
2272 let (relation_str, count) = row?;
2273 let relation = relation_str.parse::<EdgeRelation>().map_err(|e| {
2274 rusqlite::Error::FromSqlConversionFailure(
2275 0,
2276 rusqlite::types::Type::Text,
2277 Box::new(e),
2278 )
2279 })?;
2280 out.push((relation, count as u64));
2281 }
2282 Ok(out)
2283 })
2284 .await
2285 }
2286
2287 async fn count_edges_by_relation_in_namespaces(
2288 &self,
2289 namespaces: &[String],
2290 ) -> Result<Vec<(EdgeRelation, u64)>, StorageError> {
2291 let namespaces: Vec<String> = namespaces
2292 .iter()
2293 .cloned()
2294 .collect::<HashSet<_>>()
2295 .into_iter()
2296 .collect();
2297 self.with_reader("count_edges_by_relation_in_namespaces", move |conn| {
2298 let mut totals = HashMap::new();
2299 for chunk in namespaces.chunks(NAMESPACE_COUNT_CHUNK_SIZE) {
2300 let (where_clause, params) =
2301 build_edge_filter_sql_for_namespaces(chunk, &EdgeFilter::default());
2302 let sql = format!(
2303 "SELECT relation, COUNT(*) FROM graph_edges{where_clause} GROUP BY relation"
2304 );
2305 let mut stmt = conn.prepare(&sql)?;
2306 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
2307 params.iter().map(|p| p.as_ref()).collect();
2308 let rows = stmt.query_map(param_refs.as_slice(), |row| {
2309 let relation_str: String = row.get(0)?;
2310 let count: i64 = row.get(1)?;
2311 Ok((relation_str, count))
2312 })?;
2313 for row in rows {
2314 let (relation_str, count) = row?;
2315 let relation = relation_str.parse::<EdgeRelation>().map_err(|e| {
2316 rusqlite::Error::FromSqlConversionFailure(
2317 0,
2318 rusqlite::types::Type::Text,
2319 Box::new(e),
2320 )
2321 })?;
2322 *totals.entry(relation).or_insert(0) += count as u64;
2323 }
2324 }
2325 Ok(totals.into_iter().collect())
2326 })
2327 .await
2328 }
2329
2330 async fn query_edges_after(
2331 &self,
2332 filter: EdgeFilter,
2333 after: Option<Uuid>,
2334 limit: u32,
2335 ) -> Result<EdgeSeekPage, StorageError> {
2336 let namespace = self.namespace.clone();
2337 let limit_usize = limit as usize;
2338 let probe_limit_i64 = i64::from(limit) + 1;
2339 self.with_reader("query_edges_after", move |conn| {
2340 let (mut where_clause, mut params) = build_edge_filter_sql(&namespace, &filter);
2341 if let Some(cursor) = after {
2342 params.push(Box::new(cursor.to_string()));
2343 where_clause.push_str(&format!(" AND id > ?{}", params.len()));
2344 }
2345 params.push(Box::new(probe_limit_i64));
2346 let limit_idx = params.len();
2347
2348 let data_sql = format!(
2354 "SELECT namespace, id, source_id, target_id, relation, weight, \
2355 created_at, updated_at, deleted_at, metadata, target_backend \
2356 FROM graph_edges{} ORDER BY id ASC LIMIT ?{}",
2357 where_clause, limit_idx,
2358 );
2359
2360 let mut stmt = conn.prepare(&data_sql)?;
2361 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
2362 params.iter().map(|p| p.as_ref()).collect();
2363 let rows = stmt.query_map(param_refs.as_slice(), read_edge)?;
2364
2365 let mut items = Vec::new();
2366 for row in rows {
2367 items.push(row?);
2368 }
2369 let has_more = items.len() > limit_usize;
2370 if has_more {
2371 items.truncate(limit_usize);
2372 }
2373 let next_after = if has_more {
2374 items.last().map(|e| Uuid::from(e.id))
2375 } else {
2376 None
2377 };
2378
2379 Ok(EdgeSeekPage { items, next_after })
2380 })
2381 .await
2382 }
2383
2384 async fn query_edges_sequence_after(
2385 &self,
2386 filter: EdgeFilter,
2387 after: Option<SeekCursor>,
2388 limit: u32,
2389 ) -> Result<SeekPage<Edge>, StorageError> {
2390 if limit == 0 {
2391 return Ok(SeekPage::default());
2392 }
2393 let namespace = self.namespace.clone();
2394 let limit_usize = limit as usize;
2395 let probe_limit_i64 = i64::from(limit) + 1;
2396 self.with_reader("query_edges_sequence_after", move |conn| {
2397 let (mut where_clause, mut params) = build_edge_filter_sql(&namespace, &filter);
2398 if let Some(cursor) = after {
2399 params.push(Box::new(cursor.sequence));
2400 where_clause.push_str(&format!(" AND graph_edges_seq.seq > ?{}", params.len()));
2401 }
2402 params.push(Box::new(probe_limit_i64));
2403 let limit_idx = params.len();
2404 let sql = format!(
2407 "SELECT namespace, id, source_id, target_id, relation, weight, \
2408 created_at, updated_at, deleted_at, metadata, target_backend, \
2409 graph_edges_seq.seq \
2410 FROM graph_edges_seq CROSS JOIN graph_edges \
2411 ON graph_edges.id = graph_edges_seq.edge_id{where_clause} \
2412 ORDER BY graph_edges_seq.seq ASC LIMIT ?{limit_idx}"
2413 );
2414 let mut stmt = conn.prepare(&sql)?;
2415 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
2416 params.iter().map(|param| param.as_ref()).collect();
2417 let rows = stmt.query_map(param_refs.as_slice(), |row| {
2418 Ok((read_edge(row)?, row.get::<_, i64>(11)?))
2419 })?;
2420 let mut entries = rows.collect::<Result<Vec<_>, _>>()?;
2421 let has_more = entries.len() > limit_usize;
2422 if has_more {
2423 entries.truncate(limit_usize);
2424 }
2425 let next_after = if has_more {
2426 entries.last().map(|(edge, sequence)| SeekCursor {
2427 sequence: *sequence,
2428 id: Uuid::from(edge.id),
2429 })
2430 } else {
2431 None
2432 };
2433 let items = entries.into_iter().map(|(edge, _)| edge).collect();
2434 Ok(SeekPage { items, next_after })
2435 })
2436 .await
2437 }
2438
2439 async fn neighbors(
2440 &self,
2441 node_id: Uuid,
2442 query: NeighborQuery,
2443 ) -> Result<Vec<NeighborHit>, StorageError> {
2444 count_neighbor_select();
2445
2446 let namespace = self.namespace.clone();
2447 let node_str = node_id.to_string();
2448
2449 let counted_queries = Arc::new(std::sync::atomic::AtomicU64::new(0));
2450 let counted_rows = Arc::new(std::sync::atomic::AtomicU64::new(0));
2451 let closure_queries = Arc::clone(&counted_queries);
2452 let closure_rows = Arc::clone(&counted_rows);
2453 let result = self
2454 .with_reader("neighbors", move |conn| {
2455 let base_out = "SELECT target_id AS node_id, id AS edge_id, relation, weight \
2456 FROM graph_edges \
2457 WHERE namespace = ?1 AND source_id = ?2 AND deleted_at IS NULL";
2458 let base_in = "SELECT source_id AS node_id, id AS edge_id, relation, weight \
2459 FROM graph_edges \
2460 WHERE namespace = ?1 AND target_id = ?2 AND deleted_at IS NULL";
2461
2462 let sql = match query.direction {
2463 Direction::Out => base_out.to_string(),
2464 Direction::In => base_in.to_string(),
2465 Direction::Both => format!("{} UNION ALL {}", base_out, base_in),
2466 };
2467
2468 let (where_extra, limit_clause, extra_params) = neighbor_extra_clause(&query, 3);
2469
2470 let full_sql = format!(
2475 "SELECT node_id, edge_id, relation, weight FROM ({}){} \
2476 ORDER BY weight DESC, node_id ASC{}",
2477 sql, where_extra, limit_clause
2478 );
2479
2480 let mut stmt = conn.prepare(&full_sql)?;
2481
2482 let mut all_params: Vec<Box<dyn rusqlite::types::ToSql>> = Vec::new();
2483 all_params.push(Box::new(namespace.clone()));
2484 all_params.push(Box::new(node_str.clone()));
2485 all_params.extend(extra_params);
2486
2487 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
2488 all_params.iter().map(|p| p.as_ref()).collect();
2489
2490 closure_queries.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2491 let rows = stmt.query_map(param_refs.as_slice(), |row| {
2492 let nid_str: String = row.get(0)?;
2493 let eid_str: String = row.get(1)?;
2494 let relation_str: String = row.get(2)?;
2495 let weight: f64 = row.get(3)?;
2496 Ok((nid_str, eid_str, relation_str, weight))
2497 })?;
2498
2499 let mut hits = Vec::new();
2500 for row in rows {
2501 let (nid_str, eid_str, relation_str, weight) = row?;
2502 closure_rows.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2503 let relation = relation_str.parse::<EdgeRelation>().map_err(|e| {
2504 rusqlite::Error::FromSqlConversionFailure(
2505 2,
2506 rusqlite::types::Type::Text,
2507 Box::new(e),
2508 )
2509 })?;
2510 hits.push(NeighborHit {
2511 node_id: parse_uuid(&nid_str)?,
2512 edge_id: parse_uuid(&eid_str)?,
2513 relation,
2514 weight,
2515 name: None,
2516 kind: None,
2517 entity_type: None,
2518 });
2519 }
2520
2521 Ok(hits)
2522 })
2523 .await;
2524
2525 report_graph_usage(&counted_queries, &counted_rows);
2526 result
2527 }
2528
2529 async fn neighbors_both_directions(
2534 &self,
2535 node_id: Uuid,
2536 query: NeighborQuery,
2537 ) -> Result<Vec<DirectedNeighborHit>, StorageError> {
2538 count_neighbor_select();
2539
2540 let namespace = self.namespace.clone();
2541 let node_str = node_id.to_string();
2542
2543 let counted_queries = Arc::new(std::sync::atomic::AtomicU64::new(0));
2544 let counted_rows = Arc::new(std::sync::atomic::AtomicU64::new(0));
2545 let closure_queries = Arc::clone(&counted_queries);
2546 let closure_rows = Arc::clone(&counted_rows);
2547 let result = self
2548 .with_reader("neighbors_both_directions", move |conn| {
2549 let base_out = "SELECT target_id AS node_id, id AS edge_id, relation, weight, \
2550 'out' AS dir \
2551 FROM graph_edges \
2552 WHERE namespace = ?1 AND source_id = ?2 AND deleted_at IS NULL";
2553 let base_in = "SELECT source_id AS node_id, id AS edge_id, relation, weight, \
2554 'in' AS dir \
2555 FROM graph_edges \
2556 WHERE namespace = ?1 AND target_id = ?2 AND deleted_at IS NULL";
2557 let sql = format!("{} UNION ALL {}", base_out, base_in);
2558
2559 let (where_extra, limit_clause, extra_params) = neighbor_extra_clause(&query, 3);
2560
2561 let full_sql = format!(
2570 "SELECT node_id, edge_id, relation, weight, dir FROM ({}){} \
2571 ORDER BY weight DESC, node_id ASC, \
2572 CASE dir WHEN 'out' THEN 0 ELSE 1 END ASC, edge_id ASC{}",
2573 sql, where_extra, limit_clause
2574 );
2575
2576 let mut stmt = conn.prepare(&full_sql)?;
2577
2578 let mut all_params: Vec<Box<dyn rusqlite::types::ToSql>> = Vec::new();
2579 all_params.push(Box::new(namespace.clone()));
2580 all_params.push(Box::new(node_str.clone()));
2581 all_params.extend(extra_params);
2582
2583 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
2584 all_params.iter().map(|p| p.as_ref()).collect();
2585
2586 closure_queries.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2587 let rows = stmt.query_map(param_refs.as_slice(), |row| {
2588 let nid_str: String = row.get(0)?;
2589 let eid_str: String = row.get(1)?;
2590 let relation_str: String = row.get(2)?;
2591 let weight: f64 = row.get(3)?;
2592 let dir_str: String = row.get(4)?;
2593 Ok((nid_str, eid_str, relation_str, weight, dir_str))
2594 })?;
2595
2596 let mut hits = Vec::new();
2597 for row in rows {
2598 let (nid_str, eid_str, relation_str, weight, dir_str) = row?;
2599 closure_rows.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2600 let relation = relation_str.parse::<EdgeRelation>().map_err(|e| {
2601 rusqlite::Error::FromSqlConversionFailure(
2602 2,
2603 rusqlite::types::Type::Text,
2604 Box::new(e),
2605 )
2606 })?;
2607 let direction = if dir_str == "out" {
2608 Direction::Out
2609 } else {
2610 Direction::In
2611 };
2612 hits.push(DirectedNeighborHit {
2613 hit: NeighborHit {
2614 node_id: parse_uuid(&nid_str)?,
2615 edge_id: parse_uuid(&eid_str)?,
2616 relation,
2617 weight,
2618 name: None,
2619 kind: None,
2620 entity_type: None,
2621 },
2622 direction,
2623 });
2624 }
2625
2626 Ok(hits)
2627 })
2628 .await;
2629
2630 report_graph_usage(&counted_queries, &counted_rows);
2631 result
2632 }
2633
2634 async fn traverse(&self, request: TraversalRequest) -> Result<Vec<GraphPath>, StorageError> {
2635 request
2636 .validate()
2637 .map_err(|message| StorageError::InvalidInput {
2638 capability: StorageCapability::Graph,
2639 operation: "traverse".into(),
2640 message,
2641 })?;
2642 if request.roots.is_empty() {
2643 return Ok(Vec::new());
2644 }
2645
2646 let mut distinct_roots = HashSet::with_capacity(request.roots.len());
2647 let roots = request
2648 .roots
2649 .iter()
2650 .copied()
2651 .filter(|root| distinct_roots.insert(*root))
2652 .collect::<Vec<_>>();
2653 let opts = request.options;
2654 let include_roots = request.include_roots;
2655 let namespace = self.namespace.clone();
2656 let origin = self.pool.origin();
2657 let budget = request.execution_budget;
2658 let counted_rows = Arc::new(std::sync::atomic::AtomicU64::new(0));
2664 let counted_queries = Arc::new(std::sync::atomic::AtomicU64::new(0));
2665 let closure_rows = Arc::clone(&counted_rows);
2666 let closure_queries = Arc::clone(&counted_queries);
2667 let result = self
2668 .with_reader("traverse", move |conn| {
2669 Ok(run_bounded_traversal(
2670 conn,
2671 roots,
2672 opts,
2673 include_roots,
2674 namespace,
2675 origin,
2676 budget,
2677 closure_rows.as_ref(),
2678 closure_queries.as_ref(),
2679 ))
2680 })
2681 .await
2682 .and_then(|inner| inner);
2683
2684 khive_storage::usage::count(
2693 khive_storage::usage::UsageUnit::DbRoundTrips,
2694 counted_queries.load(std::sync::atomic::Ordering::Relaxed),
2695 );
2696 khive_storage::usage::count(
2697 khive_storage::usage::UsageUnit::GraphHops,
2698 counted_rows.load(std::sync::atomic::Ordering::Relaxed),
2699 );
2700
2701 result
2702 }
2703
2704 async fn purge_incident_edges(&self, node_id: Uuid) -> Result<u64, StorageError> {
2705 let statement = purge_incident_edges_statement(node_id);
2709 self.with_writer("purge_incident_edges", move |conn| {
2710 let mut stmt = conn.prepare(&statement.sql)?;
2711 bind_params(&mut stmt, &statement.params)?;
2712 Ok(stmt.raw_execute()? as u64)
2713 })
2714 .await
2715 }
2716}
2717
2718const GRAPH_DDL: &str = include_str!("../../sql/graph-ddl.sql");
2723
2724pub(crate) fn ensure_graph_schema(conn: &rusqlite::Connection) -> Result<(), rusqlite::Error> {
2725 conn.execute_batch(GRAPH_DDL)
2726}
2727
2728#[cfg(test)]
2729#[path = "graph_tests.rs"]
2730mod tests;