1use std::collections::{HashMap, HashSet};
4use std::sync::Arc;
5
6use async_trait::async_trait;
7use chrono::{DateTime, TimeZone, Utc};
8use uuid::Uuid;
9
10use khive_storage::error::StorageError;
11use khive_storage::types::{
12 BatchWriteSummary, DeleteMode, DirectedNeighborHit, Direction, Edge, EdgeFilter, EdgeSeekPage,
13 EdgeSortField, GraphPath, GuardedBatchOutcome, GuardedBatchRefusal, GuardedWriteOutcome,
14 MissingEndpoints, NeighborHit, NeighborQuery, Page, PageRequest, PathNode, SortDirection,
15 SortOrder, SqlStatement, SqlValue, TraversalRequest,
16};
17use khive_storage::GraphStore;
18use khive_storage::LinkId;
19use khive_storage::StorageCapability;
20use khive_types::EdgeRelation;
21
22use crate::error::SqliteError;
23use crate::pool::ConnectionPool;
24use crate::sql_bridge::bind_params;
25use crate::writer_task::WriterTaskHandle;
26
27fn map_err(e: rusqlite::Error, op: &'static str) -> StorageError {
29 StorageError::driver(StorageCapability::Graph, op, e)
30}
31
32fn map_sqlite_err(e: SqliteError, op: &'static str) -> StorageError {
33 StorageError::driver(StorageCapability::Graph, op, e)
34}
35
36const EDGE_NATURAL_KEY_CONFLICT_SET: &str = "weight = excluded.weight, \
52 updated_at = excluded.updated_at, \
53 deleted_at = NULL, \
54 metadata = excluded.metadata, \
55 target_backend = excluded.target_backend";
56
57fn endpoint_exists_clause(id_param: &str) -> String {
68 format!(
69 "EXISTS (SELECT 1 FROM entities WHERE id = {id_param} AND deleted_at IS NULL) \
70 OR EXISTS (SELECT 1 FROM notes WHERE id = {id_param} AND deleted_at IS NULL) \
71 OR EXISTS (SELECT 1 FROM events WHERE id = {id_param}) \
72 OR EXISTS (SELECT 1 FROM graph_edges WHERE id = {id_param} AND deleted_at IS NULL)"
73 )
74}
75
76pub fn edge_upsert_statement(edge: &Edge) -> SqlStatement {
80 let (source_id, target_id) =
81 canonical_edge_endpoints(edge.relation, edge.source_id, edge.target_id);
82 let metadata_str = edge
83 .metadata
84 .as_ref()
85 .map(|v| serde_json::to_string(v).unwrap_or_default());
86 SqlStatement {
87 sql: format!(
88 "INSERT INTO graph_edges \
89 (namespace, id, source_id, target_id, relation, weight, \
90 created_at, updated_at, deleted_at, metadata, target_backend) \
91 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11) \
92 ON CONFLICT(namespace, id) DO UPDATE SET \
93 source_id = excluded.source_id, \
94 target_id = excluded.target_id, \
95 relation = excluded.relation, \
96 {EDGE_NATURAL_KEY_CONFLICT_SET} \
97 ON CONFLICT(namespace, source_id, target_id, relation) DO UPDATE SET \
98 {EDGE_NATURAL_KEY_CONFLICT_SET}"
99 ),
100 params: vec![
101 SqlValue::Text(edge.namespace.clone()),
102 SqlValue::Text(Uuid::from(edge.id).to_string()),
103 SqlValue::Text(source_id.to_string()),
104 SqlValue::Text(target_id.to_string()),
105 SqlValue::Text(edge.relation.to_string()),
106 SqlValue::Float(edge.weight),
107 SqlValue::Integer(edge.created_at.timestamp_micros()),
108 SqlValue::Integer(edge.updated_at.timestamp_micros()),
109 match edge.deleted_at {
110 Some(t) => SqlValue::Integer(t.timestamp_micros()),
111 None => SqlValue::Null,
112 },
113 match metadata_str {
114 Some(m) => SqlValue::Text(m),
115 None => SqlValue::Null,
116 },
117 match &edge.target_backend {
118 Some(b) => SqlValue::Text(b.clone()),
119 None => SqlValue::Null,
120 },
121 ],
122 label: Some("edge-upsert".to_string()),
123 }
124}
125
126#[allow(clippy::too_many_arguments)]
146pub fn edge_insert_guarded_by_endpoints_statement(
147 namespace: &str,
148 edge_id: Uuid,
149 source_id: Uuid,
150 target_id: Uuid,
151 relation: EdgeRelation,
152 weight: f64,
153 now: i64,
154 metadata: Option<&str>,
155) -> SqlStatement {
156 let src_exists = endpoint_exists_clause("?3");
157 let tgt_exists = endpoint_exists_clause("?4");
158 SqlStatement {
159 sql: format!(
160 "INSERT INTO graph_edges \
161 (namespace, id, source_id, target_id, relation, weight, \
162 created_at, updated_at, metadata) \
163 SELECT ?1, ?2, ?3, ?4, ?5, ?6, ?7, ?7, ?8 \
164 WHERE ({src_exists}) AND ({tgt_exists}) \
165 ON CONFLICT(namespace, source_id, target_id, relation) DO UPDATE SET \
166 {EDGE_NATURAL_KEY_CONFLICT_SET}"
167 ),
168 params: vec![
169 SqlValue::Text(namespace.to_string()),
170 SqlValue::Text(edge_id.to_string()),
171 SqlValue::Text(source_id.to_string()),
172 SqlValue::Text(target_id.to_string()),
173 SqlValue::Text(relation.as_str().to_string()),
174 SqlValue::Float(weight),
175 SqlValue::Integer(now),
176 match metadata {
177 Some(m) => SqlValue::Text(m.to_string()),
178 None => SqlValue::Null,
179 },
180 ],
181 label: Some("atomic-link-insert-edge-where-exists".to_string()),
182 }
183}
184
185pub fn edge_soft_delete_statement(id: Uuid, now: i64) -> SqlStatement {
187 SqlStatement {
188 sql: "UPDATE graph_edges SET deleted_at = ?2, updated_at = ?2 \
189 WHERE id = ?1 AND deleted_at IS NULL"
190 .to_string(),
191 params: vec![SqlValue::Text(id.to_string()), SqlValue::Integer(now)],
192 label: Some("edge-delete-soft".to_string()),
193 }
194}
195
196pub fn edge_hard_delete_statement(id: Uuid) -> SqlStatement {
198 SqlStatement {
199 sql: "DELETE FROM graph_edges WHERE id = ?1".to_string(),
200 params: vec![SqlValue::Text(id.to_string())],
201 label: Some("edge-delete-hard".to_string()),
202 }
203}
204
205pub fn purge_incident_edges_statement(node_id: Uuid) -> SqlStatement {
207 SqlStatement {
208 sql: "DELETE FROM graph_edges WHERE source_id = ?1 OR target_id = ?1".to_string(),
209 params: vec![SqlValue::Text(node_id.to_string())],
210 label: Some("edge-purge-incident".to_string()),
211 }
212}
213
214pub const EDGE_SYMMETRIC_CONFLICT_PROBE_SQL: &str = "SELECT id FROM graph_edges \
235 WHERE namespace = ?1 AND source_id = ?2 AND target_id = ?3 \
236 AND relation = ?4 AND id != ?5";
237
238pub const EDGE_SYMMETRIC_DELETE_NONCANONICAL_SQL: &str =
239 "DELETE FROM graph_edges WHERE namespace = ?1 AND id = ?2";
240
241pub const EDGE_SYMMETRIC_REFRESH_CANONICAL_SQL: &str = "UPDATE graph_edges SET \
242 weight = ?1, updated_at = ?2, deleted_at = NULL, \
243 target_backend = ?3, metadata = ?4 \
244 WHERE namespace = ?5 AND id = ?6";
245
246pub const EDGE_SYMMETRIC_UPDATE_INPLACE_SQL: &str = "UPDATE graph_edges SET \
247 source_id = ?1, target_id = ?2, relation = ?3, \
248 weight = ?4, updated_at = ?5, metadata = ?6 \
249 WHERE namespace = ?7 AND id = ?8";
250
251pub fn edge_symmetric_conflict_probe_statement(
254 namespace: &str,
255 canon_src: Uuid,
256 canon_tgt: Uuid,
257 relation: EdgeRelation,
258 exclude_id: Uuid,
259) -> SqlStatement {
260 SqlStatement {
261 sql: EDGE_SYMMETRIC_CONFLICT_PROBE_SQL.to_string(),
262 params: vec![
263 SqlValue::Text(namespace.to_string()),
264 SqlValue::Text(canon_src.to_string()),
265 SqlValue::Text(canon_tgt.to_string()),
266 SqlValue::Text(relation.to_string()),
267 SqlValue::Text(exclude_id.to_string()),
268 ],
269 label: Some("edge-symmetric-conflict-probe".to_string()),
270 }
271}
272
273pub fn edge_symmetric_delete_noncanonical_statement(namespace: &str, id: Uuid) -> SqlStatement {
276 SqlStatement {
277 sql: EDGE_SYMMETRIC_DELETE_NONCANONICAL_SQL.to_string(),
278 params: vec![
279 SqlValue::Text(namespace.to_string()),
280 SqlValue::Text(id.to_string()),
281 ],
282 label: Some("edge-symmetric-delete-noncanonical".to_string()),
283 }
284}
285
286#[allow(clippy::too_many_arguments)]
289pub fn edge_symmetric_refresh_canonical_statement(
290 namespace: &str,
291 existing_id: Uuid,
292 weight: f64,
293 updated_at_micros: i64,
294 target_backend: Option<&str>,
295 metadata: Option<&str>,
296) -> SqlStatement {
297 SqlStatement {
298 sql: EDGE_SYMMETRIC_REFRESH_CANONICAL_SQL.to_string(),
299 params: vec![
300 SqlValue::Float(weight),
301 SqlValue::Integer(updated_at_micros),
302 match target_backend {
303 Some(b) => SqlValue::Text(b.to_string()),
304 None => SqlValue::Null,
305 },
306 match metadata {
307 Some(m) => SqlValue::Text(m.to_string()),
308 None => SqlValue::Null,
309 },
310 SqlValue::Text(namespace.to_string()),
311 SqlValue::Text(existing_id.to_string()),
312 ],
313 label: Some("edge-symmetric-refresh-canonical".to_string()),
314 }
315}
316
317#[allow(clippy::too_many_arguments)]
320pub fn edge_symmetric_update_inplace_statement(
321 namespace: &str,
322 id: Uuid,
323 canon_src: Uuid,
324 canon_tgt: Uuid,
325 relation: EdgeRelation,
326 weight: f64,
327 updated_at_micros: i64,
328 metadata: Option<&str>,
329) -> SqlStatement {
330 SqlStatement {
331 sql: EDGE_SYMMETRIC_UPDATE_INPLACE_SQL.to_string(),
332 params: vec![
333 SqlValue::Text(canon_src.to_string()),
334 SqlValue::Text(canon_tgt.to_string()),
335 SqlValue::Text(relation.to_string()),
336 SqlValue::Float(weight),
337 SqlValue::Integer(updated_at_micros),
338 match metadata {
339 Some(m) => SqlValue::Text(m.to_string()),
340 None => SqlValue::Null,
341 },
342 SqlValue::Text(namespace.to_string()),
343 SqlValue::Text(id.to_string()),
344 ],
345 label: Some("edge-symmetric-update-inplace".to_string()),
346 }
347}
348
349pub fn edge_symmetric_delete_if_conflict_statement(
422 namespace: &str,
423 id: Uuid,
424 canon_src: Uuid,
425 canon_tgt: Uuid,
426 relation: EdgeRelation,
427) -> SqlStatement {
428 SqlStatement {
429 sql: "DELETE FROM graph_edges \
430 WHERE namespace = ?1 AND id = ?2 \
431 AND EXISTS ( \
432 SELECT 1 FROM graph_edges \
433 WHERE namespace = ?1 AND source_id = ?3 AND target_id = ?4 \
434 AND relation = ?5 AND id != ?2 \
435 )"
436 .to_string(),
437 params: vec![
438 SqlValue::Text(namespace.to_string()),
439 SqlValue::Text(id.to_string()),
440 SqlValue::Text(canon_src.to_string()),
441 SqlValue::Text(canon_tgt.to_string()),
442 SqlValue::Text(relation.to_string()),
443 ],
444 label: Some("edge-symmetric-delete-if-conflict".to_string()),
445 }
446}
447
448#[allow(clippy::too_many_arguments)]
449pub fn edge_symmetric_refresh_or_update_inplace_statement(
450 namespace: &str,
451 id: Uuid,
452 canon_src: Uuid,
453 canon_tgt: Uuid,
454 relation: EdgeRelation,
455 weight: f64,
456 updated_at_micros: i64,
457 metadata: Option<&str>,
458 target_backend: Option<&str>,
459) -> SqlStatement {
460 SqlStatement {
461 sql: "UPDATE graph_edges SET \
462 source_id = ?3, target_id = ?4, relation = ?5, \
463 weight = ?6, updated_at = ?7, deleted_at = NULL, metadata = ?8, \
464 target_backend = CASE WHEN changes() = 1 THEN ?9 ELSE target_backend END \
465 WHERE namespace = ?1 \
466 AND ( \
467 (id = ?2 AND changes() = 0) \
468 OR (source_id = ?3 AND target_id = ?4 AND relation = ?5 \
469 AND id != ?2 AND changes() = 1) \
470 )"
471 .to_string(),
472 params: vec![
473 SqlValue::Text(namespace.to_string()),
474 SqlValue::Text(id.to_string()),
475 SqlValue::Text(canon_src.to_string()),
476 SqlValue::Text(canon_tgt.to_string()),
477 SqlValue::Text(relation.to_string()),
478 SqlValue::Float(weight),
479 SqlValue::Integer(updated_at_micros),
480 match metadata {
481 Some(m) => SqlValue::Text(m.to_string()),
482 None => SqlValue::Null,
483 },
484 match target_backend {
485 Some(b) => SqlValue::Text(b.to_string()),
486 None => SqlValue::Null,
487 },
488 ],
489 label: Some("edge-symmetric-refresh-or-update-inplace".to_string()),
490 }
491}
492
493pub struct SqlGraphStore {
495 pool: Arc<ConnectionPool>,
496 is_file_backed: bool,
497 namespace: String,
501 writer_task: Option<WriterTaskHandle>,
502}
503
504impl SqlGraphStore {
505 pub fn new_scoped(
513 pool: Arc<ConnectionPool>,
514 is_file_backed: bool,
515 namespace: impl Into<String>,
516 ) -> Self {
517 let writer_task = pool.writer_task_handle().ok().flatten();
521
522 Self {
523 pool,
524 is_file_backed,
525 namespace: namespace.into(),
526 writer_task,
527 }
528 }
529
530 fn open_standalone_writer(&self) -> Result<rusqlite::Connection, StorageError> {
531 self.pool
532 .open_standalone_writer()
533 .map_err(|e| map_sqlite_err(e, "open_graph_writer"))
534 }
535
536 fn open_standalone_reader(&self) -> Result<rusqlite::Connection, StorageError> {
537 self.pool
538 .open_standalone_reader()
539 .map_err(|e| map_sqlite_err(e, "open_graph_reader"))
540 }
541
542 async fn with_writer<F, R>(&self, op: &'static str, f: F) -> Result<R, StorageError>
548 where
549 F: FnOnce(&rusqlite::Connection) -> Result<R, rusqlite::Error> + Send + 'static,
550 R: Send + 'static,
551 {
552 if let Some(writer_task) = &self.writer_task {
553 return writer_task
554 .send(move |conn| f(conn).map_err(|e| map_err(e, op)))
555 .await;
556 }
557
558 if self.is_file_backed {
559 let conn = self.open_standalone_writer()?;
560 tokio::task::spawn_blocking(move || f(&conn).map_err(|e| map_err(e, op)))
561 .await
562 .map_err(|e| StorageError::driver(StorageCapability::Graph, op, e))?
563 } else {
564 let pool = Arc::clone(&self.pool);
565 tokio::task::spawn_blocking(move || {
566 let guard = pool.try_writer().map_err(|e| map_sqlite_err(e, op))?;
567 f(guard.conn()).map_err(|e| map_err(e, op))
568 })
569 .await
570 .map_err(|e| StorageError::driver(StorageCapability::Graph, op, e))?
571 }
572 }
573
574 async fn with_reader<F, R>(&self, op: &'static str, f: F) -> Result<R, StorageError>
575 where
576 F: FnOnce(&rusqlite::Connection) -> Result<R, rusqlite::Error> + Send + 'static,
577 R: Send + 'static,
578 {
579 if self.is_file_backed {
580 let conn = self.open_standalone_reader()?;
581 tokio::task::spawn_blocking(move || f(&conn).map_err(|e| map_err(e, op)))
582 .await
583 .map_err(|e| StorageError::driver(StorageCapability::Graph, op, e))?
584 } else {
585 let pool = Arc::clone(&self.pool);
586 tokio::task::spawn_blocking(move || {
587 let guard = pool.reader().map_err(|e| map_sqlite_err(e, op))?;
588 f(guard.conn()).map_err(|e| map_err(e, op))
589 })
590 .await
591 .map_err(|e| StorageError::driver(StorageCapability::Graph, op, e))?
592 }
593 }
594}
595
596fn read_edge(row: &rusqlite::Row<'_>) -> Result<Edge, rusqlite::Error> {
601 let namespace: String = row.get(0)?;
602 let id_str: String = row.get(1)?;
603 let source_str: String = row.get(2)?;
604 let target_str: String = row.get(3)?;
605 let relation_str: String = row.get(4)?;
606 let weight: f64 = row.get(5)?;
607 let created_micros: i64 = row.get(6)?;
608 let updated_micros: i64 = row.get(7)?;
609 let deleted_micros: Option<i64> = row.get(8)?;
610 let metadata_str: Option<String> = row.get(9)?;
611 let target_backend: Option<String> = row.get(10)?;
612
613 let id = parse_uuid(&id_str)?;
614 let source_id = parse_uuid(&source_str)?;
615 let target_id = parse_uuid(&target_str)?;
616 let created_at = micros_to_datetime(created_micros);
617 let relation = relation_str.parse::<EdgeRelation>().map_err(|e| {
618 rusqlite::Error::FromSqlConversionFailure(4, rusqlite::types::Type::Text, Box::new(e))
619 })?;
620 let metadata = match metadata_str {
621 Some(s) => {
622 let v = serde_json::from_str(&s).map_err(|e| {
623 rusqlite::Error::FromSqlConversionFailure(
624 9,
625 rusqlite::types::Type::Text,
626 Box::new(e),
627 )
628 })?;
629 Some(v)
630 }
631 None => None,
632 };
633
634 Ok(Edge {
635 id: id.into(),
636 namespace,
637 source_id,
638 target_id,
639 relation,
640 weight,
641 created_at,
642 updated_at: micros_to_datetime(updated_micros),
643 deleted_at: deleted_micros.map(micros_to_datetime),
644 metadata,
645 target_backend,
646 })
647}
648
649fn parse_uuid(s: &str) -> Result<Uuid, rusqlite::Error> {
650 Uuid::parse_str(s).map_err(|e| {
651 rusqlite::Error::FromSqlConversionFailure(0, rusqlite::types::Type::Text, Box::new(e))
652 })
653}
654
655fn neighbor_extra_clause(
661 query: &NeighborQuery,
662 start_param_idx: usize,
663) -> (String, String, Vec<Box<dyn rusqlite::types::ToSql>>) {
664 let mut conditions: Vec<String> = Vec::new();
665 let mut extra_params: Vec<Box<dyn rusqlite::types::ToSql>> = Vec::new();
666 let mut param_idx = start_param_idx;
667
668 if let Some(ref rels) = query.relations {
669 if !rels.is_empty() {
670 let placeholders: Vec<String> = rels
671 .iter()
672 .map(|r| {
673 extra_params.push(Box::new(r.to_string()));
674 let p = format!("?{}", param_idx);
675 param_idx += 1;
676 p
677 })
678 .collect();
679 conditions.push(format!("relation IN ({})", placeholders.join(",")));
680 }
681 }
682
683 if let Some(min_w) = query.min_weight {
684 extra_params.push(Box::new(min_w));
685 conditions.push(format!("weight >= ?{}", param_idx));
686 param_idx += 1;
687 }
688
689 let where_extra = if conditions.is_empty() {
690 String::new()
691 } else {
692 format!(" WHERE {}", conditions.join(" AND "))
693 };
694
695 let limit_clause = if let Some(lim) = query.limit {
696 extra_params.push(Box::new(lim as i64));
697 format!(" LIMIT ?{}", param_idx)
698 } else {
699 String::new()
700 };
701
702 (where_extra, limit_clause, extra_params)
703}
704
705#[cfg(test)]
713static NEIGHBOR_SELECT_COUNT: std::sync::atomic::AtomicUsize =
714 std::sync::atomic::AtomicUsize::new(0);
715
716#[cfg(test)]
717fn count_neighbor_select() {
718 NEIGHBOR_SELECT_COUNT.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
719}
720
721#[cfg(not(test))]
722fn count_neighbor_select() {}
723
724#[cfg(test)]
725pub(crate) fn reset_neighbor_select_count() {
726 NEIGHBOR_SELECT_COUNT.store(0, std::sync::atomic::Ordering::Relaxed);
727}
728
729#[cfg(test)]
730pub(crate) fn neighbor_select_count() -> usize {
731 NEIGHBOR_SELECT_COUNT.load(std::sync::atomic::Ordering::Relaxed)
732}
733
734fn micros_to_datetime(micros: i64) -> DateTime<Utc> {
735 Utc.timestamp_micros(micros)
736 .single()
737 .unwrap_or_else(Utc::now)
738}
739
740fn build_edge_filter_sql(
741 namespace: &str,
742 filter: &EdgeFilter,
743) -> (String, Vec<Box<dyn rusqlite::types::ToSql>>) {
744 let mut conditions: Vec<String> = vec![
745 "namespace = ?1".to_string(),
746 "deleted_at IS NULL".to_string(),
747 ];
748 let mut params: Vec<Box<dyn rusqlite::types::ToSql>> = vec![Box::new(namespace.to_string())];
749
750 if !filter.ids.is_empty() {
751 let placeholders: Vec<String> = filter
752 .ids
753 .iter()
754 .map(|id| {
755 params.push(Box::new(id.to_string()));
756 format!("?{}", params.len())
757 })
758 .collect();
759 conditions.push(format!("id IN ({})", placeholders.join(",")));
760 }
761
762 if !filter.source_ids.is_empty() {
763 let placeholders: Vec<String> = filter
764 .source_ids
765 .iter()
766 .map(|id| {
767 params.push(Box::new(id.to_string()));
768 format!("?{}", params.len())
769 })
770 .collect();
771 conditions.push(format!("source_id IN ({})", placeholders.join(",")));
772 }
773
774 if !filter.target_ids.is_empty() {
775 let placeholders: Vec<String> = filter
776 .target_ids
777 .iter()
778 .map(|id| {
779 params.push(Box::new(id.to_string()));
780 format!("?{}", params.len())
781 })
782 .collect();
783 conditions.push(format!("target_id IN ({})", placeholders.join(",")));
784 }
785
786 if !filter.relations.is_empty() {
787 let placeholders: Vec<String> = filter
788 .relations
789 .iter()
790 .map(|r| {
791 params.push(Box::new(r.to_string()));
792 format!("?{}", params.len())
793 })
794 .collect();
795 conditions.push(format!("relation IN ({})", placeholders.join(",")));
796 }
797
798 if let Some(min_w) = filter.min_weight {
799 params.push(Box::new(min_w));
800 conditions.push(format!("weight >= ?{}", params.len()));
801 }
802
803 if let Some(max_w) = filter.max_weight {
804 params.push(Box::new(max_w));
805 conditions.push(format!("weight <= ?{}", params.len()));
806 }
807
808 if let Some(ref time_range) = filter.created_at {
809 if let Some(start) = time_range.start {
810 params.push(Box::new(start.timestamp_micros()));
811 conditions.push(format!("created_at >= ?{}", params.len()));
812 }
813 if let Some(end) = time_range.end {
814 params.push(Box::new(end.timestamp_micros()));
815 conditions.push(format!("created_at < ?{}", params.len()));
816 }
817 }
818
819 let clause = format!(" WHERE {}", conditions.join(" AND "));
820 (clause, params)
821}
822
823fn edge_sort_col(field: &EdgeSortField) -> &'static str {
824 match field {
825 EdgeSortField::CreatedAt => "created_at",
826 EdgeSortField::Weight => "weight",
827 EdgeSortField::Relation => "relation",
828 }
829}
830
831fn canonical_edge_endpoints(
840 relation: EdgeRelation,
841 source_id: Uuid,
842 target_id: Uuid,
843) -> (Uuid, Uuid) {
844 if relation.is_symmetric() && target_id < source_id {
845 (target_id, source_id)
846 } else {
847 (source_id, target_id)
848 }
849}
850
851fn batch_upsert_edges(
859 conn: &rusqlite::Connection,
860 edges: &[Edge],
861 attempted: u64,
862) -> Result<BatchWriteSummary, rusqlite::Error> {
863 let mut affected = 0u64;
864
865 for edge in edges {
866 let statement = edge_upsert_statement(edge);
867 let mut stmt = conn.prepare(&statement.sql)?;
868 bind_params(&mut stmt, &statement.params)?;
869 stmt.raw_execute()?;
870 affected += 1;
871 }
872
873 Ok(BatchWriteSummary {
874 attempted,
875 affected,
876 failed: 0,
877 first_error: String::new(),
878 })
879}
880
881fn edge_endpoints_exist(
889 conn: &rusqlite::Connection,
890 source_id: Uuid,
891 target_id: Uuid,
892) -> Result<MissingEndpoints, rusqlite::Error> {
893 let src_exists = endpoint_exists_clause("?1");
894 let tgt_exists = endpoint_exists_clause("?2");
895 let sql = format!("SELECT ({src_exists}), ({tgt_exists})");
896 conn.query_row(
897 &sql,
898 rusqlite::params![source_id.to_string(), target_id.to_string()],
899 |row| {
900 let src_exists: bool = row.get(0)?;
901 let tgt_exists: bool = row.get(1)?;
902 Ok(MissingEndpoints {
903 source: !src_exists,
904 target: !tgt_exists,
905 })
906 },
907 )
908}
909
910fn edge_insert_guarded(
921 conn: &rusqlite::Connection,
922 statement: &SqlStatement,
923 source_id: Uuid,
924 target_id: Uuid,
925) -> Result<GuardedWriteOutcome, rusqlite::Error> {
926 let mut stmt = conn.prepare(&statement.sql)?;
927 bind_params(&mut stmt, &statement.params)?;
928 if stmt.raw_execute()? > 0 {
929 return Ok(GuardedWriteOutcome::Written);
930 }
931 #[cfg(test)]
939 tests::insert_probe_seam::hook((source_id, target_id));
940 let missing = edge_endpoints_exist(conn, source_id, target_id)?;
941 Ok(GuardedWriteOutcome::Refused(missing))
942}
943
944fn batch_upsert_edges_guarded(
952 conn: &rusqlite::Connection,
953 edges: &[Edge],
954 attempted: u64,
955) -> Result<GuardedBatchOutcome, rusqlite::Error> {
956 for (index, edge) in edges.iter().enumerate() {
957 let (source_id, target_id) =
958 canonical_edge_endpoints(edge.relation, edge.source_id, edge.target_id);
959 let missing = edge_endpoints_exist(conn, source_id, target_id)?;
960 if missing.any() {
961 return Ok(GuardedBatchOutcome {
962 summary: BatchWriteSummary {
963 attempted,
964 affected: 0,
965 failed: attempted,
966 first_error: format!(
967 "batch entry {index}: edge endpoint no longer exists at write time: source {source_id} or target {target_id}"
968 ),
969 },
970 refused: Some(GuardedBatchRefusal {
971 entry_index: index,
972 missing,
973 }),
974 });
975 }
976 }
977
978 let mut affected = 0u64;
979 for edge in edges {
980 let statement = edge_upsert_statement(edge);
981 let mut stmt = conn.prepare(&statement.sql)?;
982 bind_params(&mut stmt, &statement.params)?;
983 stmt.raw_execute()?;
984 affected += 1;
985 }
986
987 Ok(GuardedBatchOutcome {
988 summary: BatchWriteSummary {
989 attempted,
990 affected,
991 failed: 0,
992 first_error: String::new(),
993 },
994 refused: None,
995 })
996}
997
998#[async_trait]
999impl GraphStore for SqlGraphStore {
1000 async fn upsert_edge(&self, edge: Edge) -> Result<(), StorageError> {
1001 let statement = edge_upsert_statement(&edge);
1002 self.with_writer("upsert_edge", move |conn| {
1003 let mut stmt = conn.prepare(&statement.sql)?;
1004 bind_params(&mut stmt, &statement.params)?;
1005 stmt.raw_execute()?;
1006 Ok(())
1007 })
1008 .await
1009 }
1010
1011 async fn upsert_edges(&self, edges: Vec<Edge>) -> Result<BatchWriteSummary, StorageError> {
1012 let attempted = edges.len() as u64;
1013
1014 if let Some(writer_task) = &self.writer_task {
1020 return writer_task
1021 .send(move |conn| {
1022 batch_upsert_edges(conn, &edges, attempted)
1023 .map_err(|e| map_err(e, "upsert_edges"))
1024 })
1025 .await;
1026 }
1027
1028 self.with_writer("upsert_edges", move |conn| {
1032 conn.execute_batch("BEGIN IMMEDIATE")?;
1033 let _tx_handle =
1034 khive_storage::tx_registry::register(Some("graph_upsert_edges".to_string()));
1035
1036 let summary = match batch_upsert_edges(conn, &edges, attempted) {
1037 Ok(summary) => summary,
1038 Err(e) => {
1039 let _ = conn.execute_batch("ROLLBACK");
1040 return Err(e);
1041 }
1042 };
1043
1044 if let Err(e) = conn.execute_batch("COMMIT") {
1045 let _ = conn.execute_batch("ROLLBACK");
1046 return Err(e);
1047 }
1048 Ok(summary)
1049 })
1050 .await
1051 }
1052
1053 async fn upsert_edge_guarded(&self, edge: Edge) -> Result<GuardedWriteOutcome, StorageError> {
1054 let (source_id, target_id) =
1055 canonical_edge_endpoints(edge.relation, edge.source_id, edge.target_id);
1056 let metadata_str = edge
1057 .metadata
1058 .as_ref()
1059 .map(|v| serde_json::to_string(v).unwrap_or_default());
1060 let statement = edge_insert_guarded_by_endpoints_statement(
1061 &edge.namespace,
1062 Uuid::from(edge.id),
1063 source_id,
1064 target_id,
1065 edge.relation,
1066 edge.weight,
1067 edge.created_at.timestamp_micros(),
1068 metadata_str.as_deref(),
1069 );
1070
1071 if let Some(writer_task) = &self.writer_task {
1077 return writer_task
1078 .send(move |conn| {
1079 edge_insert_guarded(conn, &statement, source_id, target_id)
1080 .map_err(|e| map_err(e, "upsert_edge_guarded"))
1081 })
1082 .await;
1083 }
1084
1085 self.with_writer("upsert_edge_guarded", move |conn| {
1091 conn.execute_batch("BEGIN IMMEDIATE")?;
1092 let _tx_handle =
1093 khive_storage::tx_registry::register(Some("graph_upsert_edge_guarded".to_string()));
1094
1095 let outcome = match edge_insert_guarded(conn, &statement, source_id, target_id) {
1096 Ok(outcome) => outcome,
1097 Err(e) => {
1098 let _ = conn.execute_batch("ROLLBACK");
1099 return Err(e);
1100 }
1101 };
1102
1103 if let Err(e) = conn.execute_batch("COMMIT") {
1104 let _ = conn.execute_batch("ROLLBACK");
1105 return Err(e);
1106 }
1107 Ok(outcome)
1108 })
1109 .await
1110 }
1111
1112 async fn upsert_edges_guarded(
1113 &self,
1114 edges: Vec<Edge>,
1115 ) -> Result<GuardedBatchOutcome, StorageError> {
1116 let attempted = edges.len() as u64;
1117
1118 if let Some(writer_task) = &self.writer_task {
1122 return writer_task
1123 .send(move |conn| {
1124 batch_upsert_edges_guarded(conn, &edges, attempted)
1125 .map_err(|e| map_err(e, "upsert_edges_guarded"))
1126 })
1127 .await;
1128 }
1129
1130 self.with_writer("upsert_edges_guarded", move |conn| {
1131 conn.execute_batch("BEGIN IMMEDIATE")?;
1132 let _tx_handle = khive_storage::tx_registry::register(Some(
1133 "graph_upsert_edges_guarded".to_string(),
1134 ));
1135
1136 let summary = match batch_upsert_edges_guarded(conn, &edges, attempted) {
1137 Ok(summary) => summary,
1138 Err(e) => {
1139 let _ = conn.execute_batch("ROLLBACK");
1140 return Err(e);
1141 }
1142 };
1143
1144 if let Err(e) = conn.execute_batch("COMMIT") {
1145 let _ = conn.execute_batch("ROLLBACK");
1146 return Err(e);
1147 }
1148 Ok(summary)
1149 })
1150 .await
1151 }
1152
1153 async fn get_edge(&self, id: LinkId) -> Result<Option<Edge>, StorageError> {
1154 let id_str = Uuid::from(id).to_string();
1155
1156 self.with_reader("get_edge", move |conn| {
1157 let mut stmt = conn.prepare(
1158 "SELECT namespace, id, source_id, target_id, relation, weight, \
1159 created_at, updated_at, deleted_at, metadata, target_backend \
1160 FROM graph_edges WHERE id = ?1 AND deleted_at IS NULL",
1161 )?;
1162 let mut rows = stmt.query(rusqlite::params![id_str])?;
1163 match rows.next()? {
1164 Some(row) => Ok(Some(read_edge(row)?)),
1165 None => Ok(None),
1166 }
1167 })
1168 .await
1169 }
1170
1171 async fn get_edge_including_deleted(&self, id: LinkId) -> Result<Option<Edge>, StorageError> {
1172 let id_str = Uuid::from(id).to_string();
1173
1174 self.with_reader("get_edge_including_deleted", move |conn| {
1175 let mut stmt = conn.prepare(
1176 "SELECT namespace, id, source_id, target_id, relation, weight, \
1177 created_at, updated_at, deleted_at, metadata, target_backend \
1178 FROM graph_edges WHERE id = ?1",
1179 )?;
1180 let mut rows = stmt.query(rusqlite::params![id_str])?;
1181 match rows.next()? {
1182 Some(row) => Ok(Some(read_edge(row)?)),
1183 None => Ok(None),
1184 }
1185 })
1186 .await
1187 }
1188
1189 async fn get_edges(&self, ids: &[LinkId]) -> Result<Vec<Edge>, StorageError> {
1190 if ids.is_empty() {
1191 return Ok(Vec::new());
1192 }
1193 const CHUNK: usize = 900;
1195 let id_strs: Vec<String> = ids.iter().map(|id| Uuid::from(*id).to_string()).collect();
1196
1197 let mut result: Vec<Edge> = Vec::with_capacity(ids.len());
1198 for chunk in id_strs.chunks(CHUNK) {
1199 let chunk_owned: Vec<String> = chunk.to_vec();
1200 let edges = self
1201 .with_reader("get_edges", move |conn| {
1202 let placeholders: Vec<String> =
1203 (1..=chunk_owned.len()).map(|i| format!("?{}", i)).collect();
1204 let sql = format!(
1205 "SELECT namespace, id, source_id, target_id, relation, weight, \
1206 created_at, updated_at, deleted_at, metadata, target_backend \
1207 FROM graph_edges WHERE id IN ({}) AND deleted_at IS NULL",
1208 placeholders.join(",")
1209 );
1210 let mut stmt = conn.prepare(&sql)?;
1211 let params: Vec<&dyn rusqlite::types::ToSql> = chunk_owned
1212 .iter()
1213 .map(|s| s as &dyn rusqlite::types::ToSql)
1214 .collect();
1215 let rows = stmt.query_map(params.as_slice(), read_edge)?;
1216 let mut edges = Vec::new();
1217 for row in rows {
1218 edges.push(row?);
1219 }
1220 Ok(edges)
1221 })
1222 .await?;
1223 result.extend(edges);
1224 }
1225 Ok(result)
1226 }
1227
1228 async fn batch_neighbors(
1229 &self,
1230 sources: &[Uuid],
1231 query: NeighborQuery,
1232 ) -> Result<Vec<(Uuid, NeighborHit)>, StorageError> {
1233 use khive_storage::types::Direction;
1234
1235 if sources.is_empty() {
1236 return Ok(Vec::new());
1237 }
1238 let mut seen_sources = HashSet::with_capacity(sources.len());
1239 let unique_sources: Vec<Uuid> = sources
1240 .iter()
1241 .copied()
1242 .filter(|source| seen_sources.insert(*source))
1243 .collect();
1244 const CHUNK_SIZE: usize = 880;
1245
1246 let namespace = self.namespace.clone();
1247 let mut result: Vec<(Uuid, NeighborHit)> = Vec::new();
1248
1249 for chunk in unique_sources.chunks(CHUNK_SIZE) {
1250 let chunk_owned: Vec<Uuid> = chunk.to_vec();
1251 let query_clone = query.clone();
1252 let ns = namespace.clone();
1253
1254 let pairs = self
1255 .with_reader("batch_neighbors", move |conn| {
1256 let src_strs: Vec<String> = chunk_owned.iter().map(|u| u.to_string()).collect();
1257
1258 let sources_json = serde_json::to_string(&src_strs).map_err(|error| {
1259 rusqlite::Error::ToSqlConversionFailure(Box::new(error))
1260 })?;
1261
1262 let build_inner_sql =
1263 |direction_out: bool,
1264 q: &NeighborQuery|
1265 -> (String, Vec<String>, Option<f64>) {
1266 let (filter_col, node_col) = if direction_out {
1267 ("source_id", "target_id")
1268 } else {
1269 ("target_id", "source_id")
1270 };
1271
1272 let mut rel_params: Vec<String> = Vec::new();
1273 let mut conditions: Vec<String> = Vec::new();
1274 let mut param_idx = 3;
1275
1276 if let Some(ref rels) = q.relations {
1277 if !rels.is_empty() {
1278 let ps: Vec<String> = rels
1279 .iter()
1280 .map(|r| {
1281 rel_params.push(r.to_string());
1282 let p = format!("?{param_idx}");
1283 param_idx += 1;
1284 p
1285 })
1286 .collect();
1287 conditions
1288 .push(format!("edges.relation IN ({})", ps.join(",")));
1289 }
1290 }
1291
1292 let min_weight_val = if let Some(min_w) = q.min_weight {
1295 conditions.push(format!("edges.weight >= ?{param_idx}"));
1296 Some(min_w)
1297 } else {
1298 None
1299 };
1300
1301 let where_extra = if conditions.is_empty() {
1302 String::new()
1303 } else {
1304 format!(" AND {}", conditions.join(" AND "))
1305 };
1306
1307 let sql = format!(
1308 "SELECT requested.origin_id, edges.{node_col} AS node_id, \
1309 edges.id AS edge_id, edges.relation, edges.weight \
1310 FROM requested CROSS JOIN graph_edges AS edges \
1311 ON edges.{filter_col} = requested.origin_id \
1312 WHERE edges.namespace = ?1 \
1313 AND edges.deleted_at IS NULL{where_extra}",
1314 );
1315 (sql, rel_params, min_weight_val)
1316 };
1317
1318 let mut all_params: Vec<Box<dyn rusqlite::types::ToSql>> = Vec::new();
1319 all_params.push(Box::new(ns.to_string()));
1320 all_params.push(Box::new(sources_json));
1321
1322 let (combined_inner, rel_params, min_weight_val) = match query_clone.direction {
1323 Direction::Out => build_inner_sql(true, &query_clone),
1324 Direction::In => build_inner_sql(false, &query_clone),
1325 Direction::Both => {
1326 let (out_sql, rel_params, min_weight_val) =
1327 build_inner_sql(true, &query_clone);
1328 let (in_sql, _, _) = build_inner_sql(false, &query_clone);
1329 (
1330 format!("{out_sql} UNION ALL {in_sql}"),
1331 rel_params,
1332 min_weight_val,
1333 )
1334 }
1335 };
1336
1337 for relation in rel_params {
1338 all_params.push(Box::new(relation));
1339 }
1340 if let Some(min_weight) = min_weight_val {
1341 all_params.push(Box::new(min_weight));
1342 }
1343 let limit_param_idx = all_params.len() + 1;
1344
1345 let full_sql = if let Some(lim) = query_clone.limit {
1353 all_params.push(Box::new(lim as i64));
1354 format!(
1355 "WITH requested(origin_id) AS (\
1356 SELECT value FROM json_each(?2)\
1357 ) SELECT origin_id, node_id, edge_id, relation, weight \
1358 FROM (SELECT *, ROW_NUMBER() OVER (PARTITION BY origin_id \
1359 ORDER BY weight DESC, node_id ASC) AS rn \
1360 FROM ({combined_inner})) WHERE rn <= ?{limit_param_idx}",
1361 )
1362 } else {
1363 format!(
1364 "WITH requested(origin_id) AS (\
1365 SELECT value FROM json_each(?2)\
1366 ) SELECT origin_id, node_id, edge_id, relation, weight \
1367 FROM ({combined_inner})",
1368 )
1369 };
1370
1371 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
1372 all_params.iter().map(|p| p.as_ref()).collect();
1373
1374 let mut stmt = conn.prepare(&full_sql)?;
1375 let rows = stmt.query_map(param_refs.as_slice(), |row| {
1376 let origin_str: String = row.get(0)?;
1377 let nid_str: String = row.get(1)?;
1378 let eid_str: String = row.get(2)?;
1379 let relation_str: String = row.get(3)?;
1380 let weight: f64 = row.get(4)?;
1381 Ok((origin_str, nid_str, eid_str, relation_str, weight))
1382 })?;
1383
1384 let mut pairs = Vec::new();
1385 for row in rows {
1386 let (origin_str, nid_str, eid_str, relation_str, weight) = row?;
1387 let origin = parse_uuid(&origin_str)?;
1388 let node_id = parse_uuid(&nid_str)?;
1389 let edge_id = parse_uuid(&eid_str)?;
1390 let relation = relation_str.parse::<EdgeRelation>().map_err(|e| {
1391 rusqlite::Error::FromSqlConversionFailure(
1392 3,
1393 rusqlite::types::Type::Text,
1394 Box::new(e),
1395 )
1396 })?;
1397 pairs.push((
1398 origin,
1399 NeighborHit {
1400 node_id,
1401 edge_id,
1402 relation,
1403 weight,
1404 name: None,
1405 kind: None,
1406 entity_type: None,
1407 },
1408 ));
1409 }
1410 Ok(pairs)
1411 })
1412 .await?;
1413 result.extend(pairs);
1414 }
1415
1416 let requested: HashSet<Uuid> = unique_sources.iter().copied().collect();
1417 let mut grouped: HashMap<Uuid, Vec<NeighborHit>> =
1418 HashMap::with_capacity(unique_sources.len());
1419 for (origin, hit) in result {
1420 if !requested.contains(&origin) {
1421 return Err(StorageError::Internal(format!(
1422 "batch_neighbors returned unrequested origin {origin}"
1423 )));
1424 }
1425 grouped.entry(origin).or_default().push(hit);
1426 }
1427
1428 for hits in grouped.values_mut() {
1429 hits.sort_by(|a, b| {
1430 b.weight
1431 .partial_cmp(&a.weight)
1432 .unwrap_or(std::cmp::Ordering::Equal)
1433 .then(a.node_id.cmp(&b.node_id))
1434 .then(a.edge_id.cmp(&b.edge_id))
1435 });
1436 }
1437
1438 let mut ordered = Vec::new();
1439 for &source in sources {
1440 if let Some(hits) = grouped.get(&source) {
1441 ordered.extend(hits.iter().cloned().map(|hit| (source, hit)));
1442 }
1443 }
1444 Ok(ordered)
1445 }
1446
1447 async fn delete_edge(&self, id: LinkId, mode: DeleteMode) -> Result<bool, StorageError> {
1448 let id = Uuid::from(id);
1449 let statement = match mode {
1450 DeleteMode::Soft => {
1451 edge_soft_delete_statement(id, chrono::Utc::now().timestamp_micros())
1452 }
1453 DeleteMode::Hard => edge_hard_delete_statement(id),
1454 };
1455 self.with_writer("delete_edge", move |conn| {
1456 let mut stmt = conn.prepare(&statement.sql)?;
1457 bind_params(&mut stmt, &statement.params)?;
1458 Ok(stmt.raw_execute()? > 0)
1459 })
1460 .await
1461 }
1462
1463 async fn query_edges(
1464 &self,
1465 filter: EdgeFilter,
1466 sort: Vec<SortOrder<EdgeSortField>>,
1467 page: PageRequest,
1468 ) -> Result<Page<Edge>, StorageError> {
1469 let namespace = self.namespace.clone();
1470 let limit_i64 = i64::from(page.limit);
1471 let offset_i64 = i64::try_from(page.offset).map_err(|_| StorageError::InvalidInput {
1472 capability: StorageCapability::Graph,
1473 operation: "query_edges".into(),
1474 message: format!(
1475 "PageRequest: offset must be <= i64::MAX, got {}",
1476 page.offset
1477 ),
1478 })?;
1479 self.with_reader("query_edges", move |conn| {
1480 let (where_clause, filter_params) = build_edge_filter_sql(&namespace, &filter);
1481
1482 let count_sql = format!("SELECT COUNT(*) FROM graph_edges{}", where_clause);
1483 let total: i64 = {
1484 let mut stmt = conn.prepare(&count_sql)?;
1485 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
1486 filter_params.iter().map(|p| p.as_ref()).collect();
1487 stmt.query_row(param_refs.as_slice(), |row| row.get(0))?
1488 };
1489
1490 let order_clause = if sort.is_empty() {
1491 " ORDER BY created_at DESC".to_string()
1492 } else {
1493 let parts: Vec<String> = sort
1494 .iter()
1495 .map(|s| {
1496 let dir = match s.direction {
1497 SortDirection::Asc => "ASC",
1498 SortDirection::Desc => "DESC",
1499 };
1500 format!("{} {}", edge_sort_col(&s.field), dir)
1501 })
1502 .collect();
1503 format!(" ORDER BY {}", parts.join(", "))
1504 };
1505
1506 let (_, data_filter_params) = build_edge_filter_sql(&namespace, &filter);
1507 let mut all_params: Vec<Box<dyn rusqlite::types::ToSql>> = data_filter_params;
1508 all_params.push(Box::new(limit_i64));
1509 all_params.push(Box::new(offset_i64));
1510
1511 let limit_idx = all_params.len() - 1;
1512 let offset_idx = all_params.len();
1513
1514 let data_sql = format!(
1515 "SELECT namespace, id, source_id, target_id, relation, weight, \
1516 created_at, updated_at, deleted_at, metadata, target_backend \
1517 FROM graph_edges{}{} LIMIT ?{} OFFSET ?{}",
1518 where_clause, order_clause, limit_idx, offset_idx,
1519 );
1520
1521 let mut stmt = conn.prepare(&data_sql)?;
1522 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
1523 all_params.iter().map(|p| p.as_ref()).collect();
1524 let rows = stmt.query_map(param_refs.as_slice(), read_edge)?;
1525
1526 let mut items = Vec::new();
1527 for row in rows {
1528 items.push(row?);
1529 }
1530
1531 Ok(Page {
1532 items,
1533 total: Some(total as u64),
1534 })
1535 })
1536 .await
1537 }
1538
1539 async fn count_edges(&self, filter: EdgeFilter) -> Result<u64, StorageError> {
1540 let namespace = self.namespace.clone();
1541 self.with_reader("count_edges", move |conn| {
1542 let (where_clause, params) = build_edge_filter_sql(&namespace, &filter);
1543 let sql = format!("SELECT COUNT(*) FROM graph_edges{}", where_clause);
1544 let mut stmt = conn.prepare(&sql)?;
1545 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
1546 params.iter().map(|p| p.as_ref()).collect();
1547 let count: i64 = stmt.query_row(param_refs.as_slice(), |row| row.get(0))?;
1548 Ok(count as u64)
1549 })
1550 .await
1551 }
1552
1553 async fn count_edges_by_relation(&self) -> Result<Vec<(EdgeRelation, u64)>, StorageError> {
1554 let namespace = self.namespace.clone();
1555 self.with_reader("count_edges_by_relation", move |conn| {
1556 let sql = "SELECT relation, COUNT(*) FROM graph_edges \
1557 WHERE namespace = ?1 AND deleted_at IS NULL \
1558 GROUP BY relation";
1559 let mut stmt = conn.prepare(sql)?;
1560 let rows = stmt.query_map([&namespace], |row| {
1561 let relation_str: String = row.get(0)?;
1562 let count: i64 = row.get(1)?;
1563 Ok((relation_str, count))
1564 })?;
1565 let mut out = Vec::new();
1566 for row in rows {
1567 let (relation_str, count) = row?;
1568 let relation = relation_str.parse::<EdgeRelation>().map_err(|e| {
1569 rusqlite::Error::FromSqlConversionFailure(
1570 0,
1571 rusqlite::types::Type::Text,
1572 Box::new(e),
1573 )
1574 })?;
1575 out.push((relation, count as u64));
1576 }
1577 Ok(out)
1578 })
1579 .await
1580 }
1581
1582 async fn query_edges_after(
1583 &self,
1584 filter: EdgeFilter,
1585 after: Option<Uuid>,
1586 limit: u32,
1587 ) -> Result<EdgeSeekPage, StorageError> {
1588 let namespace = self.namespace.clone();
1589 let limit_usize = limit as usize;
1590 let probe_limit_i64 = i64::from(limit) + 1;
1591 self.with_reader("query_edges_after", move |conn| {
1592 let (mut where_clause, mut params) = build_edge_filter_sql(&namespace, &filter);
1593 if let Some(cursor) = after {
1594 params.push(Box::new(cursor.to_string()));
1595 where_clause.push_str(&format!(" AND id > ?{}", params.len()));
1596 }
1597 params.push(Box::new(probe_limit_i64));
1598 let limit_idx = params.len();
1599
1600 let data_sql = format!(
1606 "SELECT namespace, id, source_id, target_id, relation, weight, \
1607 created_at, updated_at, deleted_at, metadata, target_backend \
1608 FROM graph_edges{} ORDER BY id ASC LIMIT ?{}",
1609 where_clause, limit_idx,
1610 );
1611
1612 let mut stmt = conn.prepare(&data_sql)?;
1613 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
1614 params.iter().map(|p| p.as_ref()).collect();
1615 let rows = stmt.query_map(param_refs.as_slice(), read_edge)?;
1616
1617 let mut items = Vec::new();
1618 for row in rows {
1619 items.push(row?);
1620 }
1621 let has_more = items.len() > limit_usize;
1622 if has_more {
1623 items.truncate(limit_usize);
1624 }
1625 let next_after = if has_more {
1626 items.last().map(|e| Uuid::from(e.id))
1627 } else {
1628 None
1629 };
1630
1631 Ok(EdgeSeekPage { items, next_after })
1632 })
1633 .await
1634 }
1635
1636 async fn neighbors(
1637 &self,
1638 node_id: Uuid,
1639 query: NeighborQuery,
1640 ) -> Result<Vec<NeighborHit>, StorageError> {
1641 count_neighbor_select();
1642
1643 let namespace = self.namespace.clone();
1644 let node_str = node_id.to_string();
1645
1646 self.with_reader("neighbors", move |conn| {
1647 let base_out = "SELECT target_id AS node_id, id AS edge_id, relation, weight \
1648 FROM graph_edges \
1649 WHERE namespace = ?1 AND source_id = ?2 AND deleted_at IS NULL";
1650 let base_in = "SELECT source_id AS node_id, id AS edge_id, relation, weight \
1651 FROM graph_edges \
1652 WHERE namespace = ?1 AND target_id = ?2 AND deleted_at IS NULL";
1653
1654 let sql = match query.direction {
1655 Direction::Out => base_out.to_string(),
1656 Direction::In => base_in.to_string(),
1657 Direction::Both => format!("{} UNION ALL {}", base_out, base_in),
1658 };
1659
1660 let (where_extra, limit_clause, extra_params) = neighbor_extra_clause(&query, 3);
1661
1662 let full_sql = format!(
1667 "SELECT node_id, edge_id, relation, weight FROM ({}){} \
1668 ORDER BY weight DESC, node_id ASC{}",
1669 sql, where_extra, limit_clause
1670 );
1671
1672 let mut stmt = conn.prepare(&full_sql)?;
1673
1674 let mut all_params: Vec<Box<dyn rusqlite::types::ToSql>> = Vec::new();
1675 all_params.push(Box::new(namespace.clone()));
1676 all_params.push(Box::new(node_str.clone()));
1677 all_params.extend(extra_params);
1678
1679 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
1680 all_params.iter().map(|p| p.as_ref()).collect();
1681
1682 let rows = stmt.query_map(param_refs.as_slice(), |row| {
1683 let nid_str: String = row.get(0)?;
1684 let eid_str: String = row.get(1)?;
1685 let relation_str: String = row.get(2)?;
1686 let weight: f64 = row.get(3)?;
1687 Ok((nid_str, eid_str, relation_str, weight))
1688 })?;
1689
1690 let mut hits = Vec::new();
1691 for row in rows {
1692 let (nid_str, eid_str, relation_str, weight) = row?;
1693 let relation = relation_str.parse::<EdgeRelation>().map_err(|e| {
1694 rusqlite::Error::FromSqlConversionFailure(
1695 2,
1696 rusqlite::types::Type::Text,
1697 Box::new(e),
1698 )
1699 })?;
1700 hits.push(NeighborHit {
1701 node_id: parse_uuid(&nid_str)?,
1702 edge_id: parse_uuid(&eid_str)?,
1703 relation,
1704 weight,
1705 name: None,
1706 kind: None,
1707 entity_type: None,
1708 });
1709 }
1710
1711 Ok(hits)
1712 })
1713 .await
1714 }
1715
1716 async fn neighbors_both_directions(
1721 &self,
1722 node_id: Uuid,
1723 query: NeighborQuery,
1724 ) -> Result<Vec<DirectedNeighborHit>, StorageError> {
1725 count_neighbor_select();
1726
1727 let namespace = self.namespace.clone();
1728 let node_str = node_id.to_string();
1729
1730 self.with_reader("neighbors_both_directions", move |conn| {
1731 let base_out = "SELECT target_id AS node_id, id AS edge_id, relation, weight, \
1732 'out' AS dir \
1733 FROM graph_edges \
1734 WHERE namespace = ?1 AND source_id = ?2 AND deleted_at IS NULL";
1735 let base_in = "SELECT source_id AS node_id, id AS edge_id, relation, weight, \
1736 'in' AS dir \
1737 FROM graph_edges \
1738 WHERE namespace = ?1 AND target_id = ?2 AND deleted_at IS NULL";
1739 let sql = format!("{} UNION ALL {}", base_out, base_in);
1740
1741 let (where_extra, limit_clause, extra_params) = neighbor_extra_clause(&query, 3);
1742
1743 let full_sql = format!(
1752 "SELECT node_id, edge_id, relation, weight, dir FROM ({}){} \
1753 ORDER BY weight DESC, node_id ASC, \
1754 CASE dir WHEN 'out' THEN 0 ELSE 1 END ASC, edge_id ASC{}",
1755 sql, where_extra, limit_clause
1756 );
1757
1758 let mut stmt = conn.prepare(&full_sql)?;
1759
1760 let mut all_params: Vec<Box<dyn rusqlite::types::ToSql>> = Vec::new();
1761 all_params.push(Box::new(namespace.clone()));
1762 all_params.push(Box::new(node_str.clone()));
1763 all_params.extend(extra_params);
1764
1765 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
1766 all_params.iter().map(|p| p.as_ref()).collect();
1767
1768 let rows = stmt.query_map(param_refs.as_slice(), |row| {
1769 let nid_str: String = row.get(0)?;
1770 let eid_str: String = row.get(1)?;
1771 let relation_str: String = row.get(2)?;
1772 let weight: f64 = row.get(3)?;
1773 let dir_str: String = row.get(4)?;
1774 Ok((nid_str, eid_str, relation_str, weight, dir_str))
1775 })?;
1776
1777 let mut hits = Vec::new();
1778 for row in rows {
1779 let (nid_str, eid_str, relation_str, weight, dir_str) = row?;
1780 let relation = relation_str.parse::<EdgeRelation>().map_err(|e| {
1781 rusqlite::Error::FromSqlConversionFailure(
1782 2,
1783 rusqlite::types::Type::Text,
1784 Box::new(e),
1785 )
1786 })?;
1787 let direction = if dir_str == "out" {
1788 Direction::Out
1789 } else {
1790 Direction::In
1791 };
1792 hits.push(DirectedNeighborHit {
1793 hit: NeighborHit {
1794 node_id: parse_uuid(&nid_str)?,
1795 edge_id: parse_uuid(&eid_str)?,
1796 relation,
1797 weight,
1798 name: None,
1799 kind: None,
1800 entity_type: None,
1801 },
1802 direction,
1803 });
1804 }
1805
1806 Ok(hits)
1807 })
1808 .await
1809 }
1810
1811 async fn traverse(&self, request: TraversalRequest) -> Result<Vec<GraphPath>, StorageError> {
1812 use std::collections::{HashMap, HashSet};
1813
1814 use khive_storage::types::Direction;
1815
1816 if request.roots.is_empty() {
1817 return Ok(Vec::new());
1818 }
1819
1820 let roots = request.roots.clone();
1821 let opts = request.options.clone();
1822 let include_roots = request.include_roots;
1823 let namespace = self.namespace.clone();
1824 let max_depth_i64 =
1825 i64::try_from(opts.max_depth).map_err(|_| StorageError::InvalidInput {
1826 capability: StorageCapability::Graph,
1827 operation: "traverse".into(),
1828 message: format!(
1829 "TraversalOptions: max_depth must be <= i64::MAX, got {}",
1830 opts.max_depth
1831 ),
1832 })?;
1833
1834 self.with_reader("traverse", move |conn| {
1835 const CHUNK_ROOTS: usize = 400;
1848
1849 let (join_condition, next_node) = match opts.direction {
1851 Direction::Out => ("e.source_id = t.node_id", "e.target_id"),
1852 Direction::In => ("e.target_id = t.node_id", "e.source_id"),
1853 Direction::Both => (
1854 "(e.source_id = t.node_id OR e.target_id = t.node_id)",
1855 "CASE WHEN e.source_id = t.node_id THEN e.target_id ELSE e.source_id END",
1856 ),
1857 };
1858
1859 let _tx_handle =
1870 khive_storage::tx_registry::register(Some("graph_traverse_read".to_string()));
1871 let tx = conn.unchecked_transaction()?;
1872
1873 let mut root_data: HashMap<Uuid, (Vec<(PathNode, f64)>, HashSet<Uuid>)> =
1878 HashMap::with_capacity(roots.len());
1879
1880 for root_id in &roots {
1882 let (nodes, seen) = root_data.entry(*root_id).or_default();
1883 if include_roots {
1884 seen.insert(*root_id);
1885 nodes.push((
1886 PathNode {
1887 node_id: *root_id,
1888 via_edge: None,
1889 depth: 0,
1890 name: None,
1891 kind: None,
1892 properties: None,
1893 },
1894 0.0,
1895 ));
1896 }
1897 }
1898
1899 for chunk in roots.chunks(CHUNK_ROOTS) {
1900 let n_chunk = chunk.len();
1901
1902 let ns_param = n_chunk + 1;
1908 let depth_param = n_chunk + 2;
1909 let mut extra_param_idx = n_chunk + 3;
1910
1911 let mut relation_cond = String::new();
1912 let mut extra_params: Vec<Box<dyn rusqlite::types::ToSql>> = Vec::new();
1913
1914 if let Some(ref rels) = opts.relations {
1915 if !rels.is_empty() {
1916 let placeholders: Vec<String> = rels
1917 .iter()
1918 .map(|r| {
1919 extra_params.push(Box::new(r.to_string()));
1920 let p = format!("?{extra_param_idx}");
1921 extra_param_idx += 1;
1922 p
1923 })
1924 .collect();
1925 relation_cond = format!(" AND e.relation IN ({})", placeholders.join(","));
1926 }
1927 }
1928
1929 let mut weight_cond = String::new();
1930 if let Some(min_w) = opts.min_weight {
1931 extra_params.push(Box::new(min_w));
1932 weight_cond = format!(" AND e.weight >= ?{extra_param_idx}");
1933 }
1935
1936 let seed_rows: Vec<String> = (1..=n_chunk)
1939 .map(|i| format!("(?{i}, ?{i}, NULL, 0, ?{i}, 0.0)"))
1940 .collect();
1941 let seeds = seed_rows.join(", ");
1942
1943 let cte_sql = format!(
1947 "WITH RECURSIVE traversal(\
1948 root_id, node_id, edge_id, depth, path, total_weight\
1949 ) AS (\
1950 VALUES {seeds} \
1951 UNION ALL \
1952 SELECT t.root_id, {next_node}, e.id, t.depth + 1, \
1953 t.path || ',' || {next_node}, \
1954 t.total_weight + e.weight \
1955 FROM traversal t CROSS JOIN graph_edges e \
1956 ON {join_condition} \
1957 WHERE e.namespace = ?{ns} \
1958 AND e.deleted_at IS NULL \
1959 AND t.depth < ?{depth} \
1960 AND (',' || t.path || ',') NOT LIKE '%,' || {next_node} || ',%'\
1961 {rel_cond}{wt_cond} \
1962 ) \
1963 SELECT root_id, node_id, edge_id, depth, total_weight \
1964 FROM traversal WHERE depth > 0 \
1965 ORDER BY root_id, depth",
1966 seeds = seeds,
1967 next_node = next_node,
1968 join_condition = join_condition,
1969 ns = ns_param,
1970 depth = depth_param,
1971 rel_cond = relation_cond,
1972 wt_cond = weight_cond,
1973 );
1974
1975 let mut all_params: Vec<Box<dyn rusqlite::types::ToSql>> = Vec::new();
1976 for root_id in chunk {
1977 all_params.push(Box::new(root_id.to_string()));
1978 }
1979 all_params.push(Box::new(namespace.clone()));
1980 all_params.push(Box::new(max_depth_i64));
1981 all_params.extend(extra_params);
1982
1983 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
1984 all_params.iter().map(|p| p.as_ref()).collect();
1985
1986 let mut stmt = conn.prepare(&cte_sql)?;
1989 let rows_iter = stmt.query_map(param_refs.as_slice(), |row| {
1990 let root_str: String = row.get(0)?;
1991 let node_str: String = row.get(1)?;
1992 let edge_str: Option<String> = row.get(2)?;
1993 let depth: i64 = row.get(3)?;
1994 let total_weight: f64 = row.get(4)?;
1995 Ok((root_str, node_str, edge_str, depth, total_weight))
1996 })?;
1997
1998 for row in rows_iter {
2002 let (root_str, node_str, edge_str, depth, total_weight) = row?;
2003 let root_id = parse_uuid(&root_str)?;
2004 let node_id = parse_uuid(&node_str)?;
2005 let (nodes, seen) = root_data.entry(root_id).or_default();
2006 if !seen.insert(node_id) {
2007 continue;
2008 }
2009 let via_edge = edge_str.map(|s| parse_uuid(&s)).transpose()?;
2010 nodes.push((
2011 PathNode {
2012 node_id,
2013 via_edge,
2014 depth: depth as usize,
2015 name: None,
2016 kind: None,
2017 properties: None,
2018 },
2019 total_weight,
2020 ));
2021 }
2022 }
2023
2024 tx.commit()?;
2025
2026 let mut all_paths: Vec<GraphPath> = Vec::with_capacity(roots.len());
2032 for root_id in &roots {
2033 if let Some((mut nw, _)) = root_data.remove(root_id) {
2034 if nw.is_empty() {
2035 continue;
2036 }
2037 if let Some(lim) = opts.limit {
2038 let root_count = usize::from(include_roots);
2039 nw.truncate(root_count + lim as usize);
2040 }
2041 if nw.is_empty() {
2044 continue;
2045 }
2046 let max_weight = nw.iter().map(|(_, w)| *w).fold(0.0_f64, f64::max);
2047 let nodes: Vec<PathNode> = nw.into_iter().map(|(n, _)| n).collect();
2048 all_paths.push(GraphPath {
2049 root_id: *root_id,
2050 nodes,
2051 total_weight: max_weight,
2052 });
2053 }
2054 }
2055
2056 Ok(all_paths)
2057 })
2058 .await
2059 }
2060
2061 async fn purge_incident_edges(&self, node_id: Uuid) -> Result<u64, StorageError> {
2062 let statement = purge_incident_edges_statement(node_id);
2066 self.with_writer("purge_incident_edges", move |conn| {
2067 let mut stmt = conn.prepare(&statement.sql)?;
2068 bind_params(&mut stmt, &statement.params)?;
2069 Ok(stmt.raw_execute()? as u64)
2070 })
2071 .await
2072 }
2073}
2074
2075const GRAPH_DDL: &str = include_str!("../../sql/graph-ddl.sql");
2080
2081pub(crate) fn ensure_graph_schema(conn: &rusqlite::Connection) -> Result<(), rusqlite::Error> {
2082 conn.execute_batch(GRAPH_DDL)
2083}
2084
2085#[cfg(test)]
2086#[path = "graph_tests.rs"]
2087mod tests;