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 let origin = self.pool.origin();
1032 self.with_writer("upsert_edges", move |conn| {
1033 conn.execute_batch("BEGIN IMMEDIATE")?;
1034 let _tx_handle = khive_storage::tx_registry::register_scoped(
1035 Some("graph_upsert_edges".to_string()),
1036 origin,
1037 );
1038
1039 let summary = match batch_upsert_edges(conn, &edges, attempted) {
1040 Ok(summary) => summary,
1041 Err(e) => {
1042 let _ = conn.execute_batch("ROLLBACK");
1043 return Err(e);
1044 }
1045 };
1046
1047 if let Err(e) = conn.execute_batch("COMMIT") {
1048 let _ = conn.execute_batch("ROLLBACK");
1049 return Err(e);
1050 }
1051 Ok(summary)
1052 })
1053 .await
1054 }
1055
1056 async fn upsert_edge_guarded(&self, edge: Edge) -> Result<GuardedWriteOutcome, StorageError> {
1057 let (source_id, target_id) =
1058 canonical_edge_endpoints(edge.relation, edge.source_id, edge.target_id);
1059 let metadata_str = edge
1060 .metadata
1061 .as_ref()
1062 .map(|v| serde_json::to_string(v).unwrap_or_default());
1063 let statement = edge_insert_guarded_by_endpoints_statement(
1064 &edge.namespace,
1065 Uuid::from(edge.id),
1066 source_id,
1067 target_id,
1068 edge.relation,
1069 edge.weight,
1070 edge.created_at.timestamp_micros(),
1071 metadata_str.as_deref(),
1072 );
1073
1074 if let Some(writer_task) = &self.writer_task {
1080 return writer_task
1081 .send(move |conn| {
1082 edge_insert_guarded(conn, &statement, source_id, target_id)
1083 .map_err(|e| map_err(e, "upsert_edge_guarded"))
1084 })
1085 .await;
1086 }
1087
1088 let origin = self.pool.origin();
1094 self.with_writer("upsert_edge_guarded", move |conn| {
1095 conn.execute_batch("BEGIN IMMEDIATE")?;
1096 let _tx_handle = khive_storage::tx_registry::register_scoped(
1097 Some("graph_upsert_edge_guarded".to_string()),
1098 origin,
1099 );
1100
1101 let outcome = match edge_insert_guarded(conn, &statement, source_id, target_id) {
1102 Ok(outcome) => outcome,
1103 Err(e) => {
1104 let _ = conn.execute_batch("ROLLBACK");
1105 return Err(e);
1106 }
1107 };
1108
1109 if let Err(e) = conn.execute_batch("COMMIT") {
1110 let _ = conn.execute_batch("ROLLBACK");
1111 return Err(e);
1112 }
1113 Ok(outcome)
1114 })
1115 .await
1116 }
1117
1118 async fn upsert_edges_guarded(
1119 &self,
1120 edges: Vec<Edge>,
1121 ) -> Result<GuardedBatchOutcome, StorageError> {
1122 let attempted = edges.len() as u64;
1123
1124 if let Some(writer_task) = &self.writer_task {
1128 return writer_task
1129 .send(move |conn| {
1130 batch_upsert_edges_guarded(conn, &edges, attempted)
1131 .map_err(|e| map_err(e, "upsert_edges_guarded"))
1132 })
1133 .await;
1134 }
1135
1136 let origin = self.pool.origin();
1137 self.with_writer("upsert_edges_guarded", move |conn| {
1138 conn.execute_batch("BEGIN IMMEDIATE")?;
1139 let _tx_handle = khive_storage::tx_registry::register_scoped(
1140 Some("graph_upsert_edges_guarded".to_string()),
1141 origin,
1142 );
1143
1144 let summary = match batch_upsert_edges_guarded(conn, &edges, attempted) {
1145 Ok(summary) => summary,
1146 Err(e) => {
1147 let _ = conn.execute_batch("ROLLBACK");
1148 return Err(e);
1149 }
1150 };
1151
1152 if let Err(e) = conn.execute_batch("COMMIT") {
1153 let _ = conn.execute_batch("ROLLBACK");
1154 return Err(e);
1155 }
1156 Ok(summary)
1157 })
1158 .await
1159 }
1160
1161 async fn get_edge(&self, id: LinkId) -> Result<Option<Edge>, StorageError> {
1162 let id_str = Uuid::from(id).to_string();
1163
1164 self.with_reader("get_edge", move |conn| {
1165 let mut stmt = conn.prepare(
1166 "SELECT namespace, id, source_id, target_id, relation, weight, \
1167 created_at, updated_at, deleted_at, metadata, target_backend \
1168 FROM graph_edges WHERE id = ?1 AND deleted_at IS NULL",
1169 )?;
1170 let mut rows = stmt.query(rusqlite::params![id_str])?;
1171 match rows.next()? {
1172 Some(row) => Ok(Some(read_edge(row)?)),
1173 None => Ok(None),
1174 }
1175 })
1176 .await
1177 }
1178
1179 async fn get_edge_including_deleted(&self, id: LinkId) -> Result<Option<Edge>, StorageError> {
1180 let id_str = Uuid::from(id).to_string();
1181
1182 self.with_reader("get_edge_including_deleted", move |conn| {
1183 let mut stmt = conn.prepare(
1184 "SELECT namespace, id, source_id, target_id, relation, weight, \
1185 created_at, updated_at, deleted_at, metadata, target_backend \
1186 FROM graph_edges WHERE id = ?1",
1187 )?;
1188 let mut rows = stmt.query(rusqlite::params![id_str])?;
1189 match rows.next()? {
1190 Some(row) => Ok(Some(read_edge(row)?)),
1191 None => Ok(None),
1192 }
1193 })
1194 .await
1195 }
1196
1197 async fn get_edges(&self, ids: &[LinkId]) -> Result<Vec<Edge>, StorageError> {
1198 if ids.is_empty() {
1199 return Ok(Vec::new());
1200 }
1201 const CHUNK: usize = 900;
1203 let id_strs: Vec<String> = ids.iter().map(|id| Uuid::from(*id).to_string()).collect();
1204
1205 let mut result: Vec<Edge> = Vec::with_capacity(ids.len());
1206 for chunk in id_strs.chunks(CHUNK) {
1207 let chunk_owned: Vec<String> = chunk.to_vec();
1208 let edges = self
1209 .with_reader("get_edges", move |conn| {
1210 let placeholders: Vec<String> =
1211 (1..=chunk_owned.len()).map(|i| format!("?{}", i)).collect();
1212 let sql = format!(
1213 "SELECT namespace, id, source_id, target_id, relation, weight, \
1214 created_at, updated_at, deleted_at, metadata, target_backend \
1215 FROM graph_edges WHERE id IN ({}) AND deleted_at IS NULL",
1216 placeholders.join(",")
1217 );
1218 let mut stmt = conn.prepare(&sql)?;
1219 let params: Vec<&dyn rusqlite::types::ToSql> = chunk_owned
1220 .iter()
1221 .map(|s| s as &dyn rusqlite::types::ToSql)
1222 .collect();
1223 let rows = stmt.query_map(params.as_slice(), read_edge)?;
1224 let mut edges = Vec::new();
1225 for row in rows {
1226 edges.push(row?);
1227 }
1228 Ok(edges)
1229 })
1230 .await?;
1231 result.extend(edges);
1232 }
1233 Ok(result)
1234 }
1235
1236 async fn batch_neighbors(
1237 &self,
1238 sources: &[Uuid],
1239 query: NeighborQuery,
1240 ) -> Result<Vec<(Uuid, NeighborHit)>, StorageError> {
1241 use khive_storage::types::Direction;
1242
1243 if sources.is_empty() {
1244 return Ok(Vec::new());
1245 }
1246 let mut seen_sources = HashSet::with_capacity(sources.len());
1247 let unique_sources: Vec<Uuid> = sources
1248 .iter()
1249 .copied()
1250 .filter(|source| seen_sources.insert(*source))
1251 .collect();
1252 const CHUNK_SIZE: usize = 880;
1253
1254 let namespace = self.namespace.clone();
1255 let mut result: Vec<(Uuid, NeighborHit)> = Vec::new();
1256
1257 for chunk in unique_sources.chunks(CHUNK_SIZE) {
1258 let chunk_owned: Vec<Uuid> = chunk.to_vec();
1259 let query_clone = query.clone();
1260 let ns = namespace.clone();
1261
1262 let pairs = self
1263 .with_reader("batch_neighbors", move |conn| {
1264 let src_strs: Vec<String> = chunk_owned.iter().map(|u| u.to_string()).collect();
1265
1266 let sources_json = serde_json::to_string(&src_strs).map_err(|error| {
1267 rusqlite::Error::ToSqlConversionFailure(Box::new(error))
1268 })?;
1269
1270 let build_inner_sql =
1271 |direction_out: bool,
1272 q: &NeighborQuery|
1273 -> (String, Vec<String>, Option<f64>) {
1274 let (filter_col, node_col) = if direction_out {
1275 ("source_id", "target_id")
1276 } else {
1277 ("target_id", "source_id")
1278 };
1279
1280 let mut rel_params: Vec<String> = Vec::new();
1281 let mut conditions: Vec<String> = Vec::new();
1282 let mut param_idx = 3;
1283
1284 if let Some(ref rels) = q.relations {
1285 if !rels.is_empty() {
1286 let ps: Vec<String> = rels
1287 .iter()
1288 .map(|r| {
1289 rel_params.push(r.to_string());
1290 let p = format!("?{param_idx}");
1291 param_idx += 1;
1292 p
1293 })
1294 .collect();
1295 conditions
1296 .push(format!("edges.relation IN ({})", ps.join(",")));
1297 }
1298 }
1299
1300 let min_weight_val = if let Some(min_w) = q.min_weight {
1303 conditions.push(format!("edges.weight >= ?{param_idx}"));
1304 Some(min_w)
1305 } else {
1306 None
1307 };
1308
1309 let where_extra = if conditions.is_empty() {
1310 String::new()
1311 } else {
1312 format!(" AND {}", conditions.join(" AND "))
1313 };
1314
1315 let sql = format!(
1316 "SELECT requested.origin_id, edges.{node_col} AS node_id, \
1317 edges.id AS edge_id, edges.relation, edges.weight \
1318 FROM requested CROSS JOIN graph_edges AS edges \
1319 ON edges.{filter_col} = requested.origin_id \
1320 WHERE edges.namespace = ?1 \
1321 AND edges.deleted_at IS NULL{where_extra}",
1322 );
1323 (sql, rel_params, min_weight_val)
1324 };
1325
1326 let mut all_params: Vec<Box<dyn rusqlite::types::ToSql>> = Vec::new();
1327 all_params.push(Box::new(ns.to_string()));
1328 all_params.push(Box::new(sources_json));
1329
1330 let (combined_inner, rel_params, min_weight_val) = match query_clone.direction {
1331 Direction::Out => build_inner_sql(true, &query_clone),
1332 Direction::In => build_inner_sql(false, &query_clone),
1333 Direction::Both => {
1334 let (out_sql, rel_params, min_weight_val) =
1335 build_inner_sql(true, &query_clone);
1336 let (in_sql, _, _) = build_inner_sql(false, &query_clone);
1337 (
1338 format!("{out_sql} UNION ALL {in_sql}"),
1339 rel_params,
1340 min_weight_val,
1341 )
1342 }
1343 };
1344
1345 for relation in rel_params {
1346 all_params.push(Box::new(relation));
1347 }
1348 if let Some(min_weight) = min_weight_val {
1349 all_params.push(Box::new(min_weight));
1350 }
1351 let limit_param_idx = all_params.len() + 1;
1352
1353 let full_sql = if let Some(lim) = query_clone.limit {
1361 all_params.push(Box::new(lim as i64));
1362 format!(
1363 "WITH requested(origin_id) AS (\
1364 SELECT value FROM json_each(?2)\
1365 ) SELECT origin_id, node_id, edge_id, relation, weight \
1366 FROM (SELECT *, ROW_NUMBER() OVER (PARTITION BY origin_id \
1367 ORDER BY weight DESC, node_id ASC) AS rn \
1368 FROM ({combined_inner})) WHERE rn <= ?{limit_param_idx}",
1369 )
1370 } else {
1371 format!(
1372 "WITH requested(origin_id) AS (\
1373 SELECT value FROM json_each(?2)\
1374 ) SELECT origin_id, node_id, edge_id, relation, weight \
1375 FROM ({combined_inner})",
1376 )
1377 };
1378
1379 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
1380 all_params.iter().map(|p| p.as_ref()).collect();
1381
1382 let mut stmt = conn.prepare(&full_sql)?;
1383 let rows = stmt.query_map(param_refs.as_slice(), |row| {
1384 let origin_str: String = row.get(0)?;
1385 let nid_str: String = row.get(1)?;
1386 let eid_str: String = row.get(2)?;
1387 let relation_str: String = row.get(3)?;
1388 let weight: f64 = row.get(4)?;
1389 Ok((origin_str, nid_str, eid_str, relation_str, weight))
1390 })?;
1391
1392 let mut pairs = Vec::new();
1393 for row in rows {
1394 let (origin_str, nid_str, eid_str, relation_str, weight) = row?;
1395 let origin = parse_uuid(&origin_str)?;
1396 let node_id = parse_uuid(&nid_str)?;
1397 let edge_id = parse_uuid(&eid_str)?;
1398 let relation = relation_str.parse::<EdgeRelation>().map_err(|e| {
1399 rusqlite::Error::FromSqlConversionFailure(
1400 3,
1401 rusqlite::types::Type::Text,
1402 Box::new(e),
1403 )
1404 })?;
1405 pairs.push((
1406 origin,
1407 NeighborHit {
1408 node_id,
1409 edge_id,
1410 relation,
1411 weight,
1412 name: None,
1413 kind: None,
1414 entity_type: None,
1415 },
1416 ));
1417 }
1418 Ok(pairs)
1419 })
1420 .await?;
1421 result.extend(pairs);
1422 }
1423
1424 let requested: HashSet<Uuid> = unique_sources.iter().copied().collect();
1425 let mut grouped: HashMap<Uuid, Vec<NeighborHit>> =
1426 HashMap::with_capacity(unique_sources.len());
1427 for (origin, hit) in result {
1428 if !requested.contains(&origin) {
1429 return Err(StorageError::Internal(format!(
1430 "batch_neighbors returned unrequested origin {origin}"
1431 )));
1432 }
1433 grouped.entry(origin).or_default().push(hit);
1434 }
1435
1436 for hits in grouped.values_mut() {
1437 hits.sort_by(|a, b| {
1438 b.weight
1439 .partial_cmp(&a.weight)
1440 .unwrap_or(std::cmp::Ordering::Equal)
1441 .then(a.node_id.cmp(&b.node_id))
1442 .then(a.edge_id.cmp(&b.edge_id))
1443 });
1444 }
1445
1446 let mut ordered = Vec::new();
1447 for &source in sources {
1448 if let Some(hits) = grouped.get(&source) {
1449 ordered.extend(hits.iter().cloned().map(|hit| (source, hit)));
1450 }
1451 }
1452 Ok(ordered)
1453 }
1454
1455 async fn delete_edge(&self, id: LinkId, mode: DeleteMode) -> Result<bool, StorageError> {
1456 let id = Uuid::from(id);
1457 let statement = match mode {
1458 DeleteMode::Soft => {
1459 edge_soft_delete_statement(id, chrono::Utc::now().timestamp_micros())
1460 }
1461 DeleteMode::Hard => edge_hard_delete_statement(id),
1462 };
1463 self.with_writer("delete_edge", move |conn| {
1464 let mut stmt = conn.prepare(&statement.sql)?;
1465 bind_params(&mut stmt, &statement.params)?;
1466 Ok(stmt.raw_execute()? > 0)
1467 })
1468 .await
1469 }
1470
1471 async fn query_edges(
1472 &self,
1473 filter: EdgeFilter,
1474 sort: Vec<SortOrder<EdgeSortField>>,
1475 page: PageRequest,
1476 ) -> Result<Page<Edge>, StorageError> {
1477 let namespace = self.namespace.clone();
1478 let limit_i64 = i64::from(page.limit);
1479 let offset_i64 = i64::try_from(page.offset).map_err(|_| StorageError::InvalidInput {
1480 capability: StorageCapability::Graph,
1481 operation: "query_edges".into(),
1482 message: format!(
1483 "PageRequest: offset must be <= i64::MAX, got {}",
1484 page.offset
1485 ),
1486 })?;
1487 self.with_reader("query_edges", move |conn| {
1488 let (where_clause, filter_params) = build_edge_filter_sql(&namespace, &filter);
1489
1490 let count_sql = format!("SELECT COUNT(*) FROM graph_edges{}", where_clause);
1491 let total: i64 = {
1492 let mut stmt = conn.prepare(&count_sql)?;
1493 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
1494 filter_params.iter().map(|p| p.as_ref()).collect();
1495 stmt.query_row(param_refs.as_slice(), |row| row.get(0))?
1496 };
1497
1498 let order_clause = if sort.is_empty() {
1499 " ORDER BY created_at DESC".to_string()
1500 } else {
1501 let parts: Vec<String> = sort
1502 .iter()
1503 .map(|s| {
1504 let dir = match s.direction {
1505 SortDirection::Asc => "ASC",
1506 SortDirection::Desc => "DESC",
1507 };
1508 format!("{} {}", edge_sort_col(&s.field), dir)
1509 })
1510 .collect();
1511 format!(" ORDER BY {}", parts.join(", "))
1512 };
1513
1514 let (_, data_filter_params) = build_edge_filter_sql(&namespace, &filter);
1515 let mut all_params: Vec<Box<dyn rusqlite::types::ToSql>> = data_filter_params;
1516 all_params.push(Box::new(limit_i64));
1517 all_params.push(Box::new(offset_i64));
1518
1519 let limit_idx = all_params.len() - 1;
1520 let offset_idx = all_params.len();
1521
1522 let data_sql = format!(
1523 "SELECT namespace, id, source_id, target_id, relation, weight, \
1524 created_at, updated_at, deleted_at, metadata, target_backend \
1525 FROM graph_edges{}{} LIMIT ?{} OFFSET ?{}",
1526 where_clause, order_clause, limit_idx, offset_idx,
1527 );
1528
1529 let mut stmt = conn.prepare(&data_sql)?;
1530 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
1531 all_params.iter().map(|p| p.as_ref()).collect();
1532 let rows = stmt.query_map(param_refs.as_slice(), read_edge)?;
1533
1534 let mut items = Vec::new();
1535 for row in rows {
1536 items.push(row?);
1537 }
1538
1539 Ok(Page {
1540 items,
1541 total: Some(total as u64),
1542 })
1543 })
1544 .await
1545 }
1546
1547 async fn count_edges(&self, filter: EdgeFilter) -> Result<u64, StorageError> {
1548 let namespace = self.namespace.clone();
1549 self.with_reader("count_edges", move |conn| {
1550 let (where_clause, params) = build_edge_filter_sql(&namespace, &filter);
1551 let sql = format!("SELECT COUNT(*) FROM graph_edges{}", where_clause);
1552 let mut stmt = conn.prepare(&sql)?;
1553 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
1554 params.iter().map(|p| p.as_ref()).collect();
1555 let count: i64 = stmt.query_row(param_refs.as_slice(), |row| row.get(0))?;
1556 Ok(count as u64)
1557 })
1558 .await
1559 }
1560
1561 async fn count_edges_by_relation(&self) -> Result<Vec<(EdgeRelation, u64)>, StorageError> {
1562 let namespace = self.namespace.clone();
1563 self.with_reader("count_edges_by_relation", move |conn| {
1564 let sql = "SELECT relation, COUNT(*) FROM graph_edges \
1565 WHERE namespace = ?1 AND deleted_at IS NULL \
1566 GROUP BY relation";
1567 let mut stmt = conn.prepare(sql)?;
1568 let rows = stmt.query_map([&namespace], |row| {
1569 let relation_str: String = row.get(0)?;
1570 let count: i64 = row.get(1)?;
1571 Ok((relation_str, count))
1572 })?;
1573 let mut out = Vec::new();
1574 for row in rows {
1575 let (relation_str, count) = row?;
1576 let relation = relation_str.parse::<EdgeRelation>().map_err(|e| {
1577 rusqlite::Error::FromSqlConversionFailure(
1578 0,
1579 rusqlite::types::Type::Text,
1580 Box::new(e),
1581 )
1582 })?;
1583 out.push((relation, count as u64));
1584 }
1585 Ok(out)
1586 })
1587 .await
1588 }
1589
1590 async fn query_edges_after(
1591 &self,
1592 filter: EdgeFilter,
1593 after: Option<Uuid>,
1594 limit: u32,
1595 ) -> Result<EdgeSeekPage, StorageError> {
1596 let namespace = self.namespace.clone();
1597 let limit_usize = limit as usize;
1598 let probe_limit_i64 = i64::from(limit) + 1;
1599 self.with_reader("query_edges_after", move |conn| {
1600 let (mut where_clause, mut params) = build_edge_filter_sql(&namespace, &filter);
1601 if let Some(cursor) = after {
1602 params.push(Box::new(cursor.to_string()));
1603 where_clause.push_str(&format!(" AND id > ?{}", params.len()));
1604 }
1605 params.push(Box::new(probe_limit_i64));
1606 let limit_idx = params.len();
1607
1608 let data_sql = format!(
1614 "SELECT namespace, id, source_id, target_id, relation, weight, \
1615 created_at, updated_at, deleted_at, metadata, target_backend \
1616 FROM graph_edges{} ORDER BY id ASC LIMIT ?{}",
1617 where_clause, limit_idx,
1618 );
1619
1620 let mut stmt = conn.prepare(&data_sql)?;
1621 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
1622 params.iter().map(|p| p.as_ref()).collect();
1623 let rows = stmt.query_map(param_refs.as_slice(), read_edge)?;
1624
1625 let mut items = Vec::new();
1626 for row in rows {
1627 items.push(row?);
1628 }
1629 let has_more = items.len() > limit_usize;
1630 if has_more {
1631 items.truncate(limit_usize);
1632 }
1633 let next_after = if has_more {
1634 items.last().map(|e| Uuid::from(e.id))
1635 } else {
1636 None
1637 };
1638
1639 Ok(EdgeSeekPage { items, next_after })
1640 })
1641 .await
1642 }
1643
1644 async fn neighbors(
1645 &self,
1646 node_id: Uuid,
1647 query: NeighborQuery,
1648 ) -> Result<Vec<NeighborHit>, StorageError> {
1649 count_neighbor_select();
1650
1651 let namespace = self.namespace.clone();
1652 let node_str = node_id.to_string();
1653
1654 self.with_reader("neighbors", move |conn| {
1655 let base_out = "SELECT target_id AS node_id, id AS edge_id, relation, weight \
1656 FROM graph_edges \
1657 WHERE namespace = ?1 AND source_id = ?2 AND deleted_at IS NULL";
1658 let base_in = "SELECT source_id AS node_id, id AS edge_id, relation, weight \
1659 FROM graph_edges \
1660 WHERE namespace = ?1 AND target_id = ?2 AND deleted_at IS NULL";
1661
1662 let sql = match query.direction {
1663 Direction::Out => base_out.to_string(),
1664 Direction::In => base_in.to_string(),
1665 Direction::Both => format!("{} UNION ALL {}", base_out, base_in),
1666 };
1667
1668 let (where_extra, limit_clause, extra_params) = neighbor_extra_clause(&query, 3);
1669
1670 let full_sql = format!(
1675 "SELECT node_id, edge_id, relation, weight FROM ({}){} \
1676 ORDER BY weight DESC, node_id ASC{}",
1677 sql, where_extra, limit_clause
1678 );
1679
1680 let mut stmt = conn.prepare(&full_sql)?;
1681
1682 let mut all_params: Vec<Box<dyn rusqlite::types::ToSql>> = Vec::new();
1683 all_params.push(Box::new(namespace.clone()));
1684 all_params.push(Box::new(node_str.clone()));
1685 all_params.extend(extra_params);
1686
1687 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
1688 all_params.iter().map(|p| p.as_ref()).collect();
1689
1690 let rows = stmt.query_map(param_refs.as_slice(), |row| {
1691 let nid_str: String = row.get(0)?;
1692 let eid_str: String = row.get(1)?;
1693 let relation_str: String = row.get(2)?;
1694 let weight: f64 = row.get(3)?;
1695 Ok((nid_str, eid_str, relation_str, weight))
1696 })?;
1697
1698 let mut hits = Vec::new();
1699 for row in rows {
1700 let (nid_str, eid_str, relation_str, weight) = row?;
1701 let relation = relation_str.parse::<EdgeRelation>().map_err(|e| {
1702 rusqlite::Error::FromSqlConversionFailure(
1703 2,
1704 rusqlite::types::Type::Text,
1705 Box::new(e),
1706 )
1707 })?;
1708 hits.push(NeighborHit {
1709 node_id: parse_uuid(&nid_str)?,
1710 edge_id: parse_uuid(&eid_str)?,
1711 relation,
1712 weight,
1713 name: None,
1714 kind: None,
1715 entity_type: None,
1716 });
1717 }
1718
1719 Ok(hits)
1720 })
1721 .await
1722 }
1723
1724 async fn neighbors_both_directions(
1729 &self,
1730 node_id: Uuid,
1731 query: NeighborQuery,
1732 ) -> Result<Vec<DirectedNeighborHit>, StorageError> {
1733 count_neighbor_select();
1734
1735 let namespace = self.namespace.clone();
1736 let node_str = node_id.to_string();
1737
1738 self.with_reader("neighbors_both_directions", move |conn| {
1739 let base_out = "SELECT target_id AS node_id, id AS edge_id, relation, weight, \
1740 'out' AS dir \
1741 FROM graph_edges \
1742 WHERE namespace = ?1 AND source_id = ?2 AND deleted_at IS NULL";
1743 let base_in = "SELECT source_id AS node_id, id AS edge_id, relation, weight, \
1744 'in' AS dir \
1745 FROM graph_edges \
1746 WHERE namespace = ?1 AND target_id = ?2 AND deleted_at IS NULL";
1747 let sql = format!("{} UNION ALL {}", base_out, base_in);
1748
1749 let (where_extra, limit_clause, extra_params) = neighbor_extra_clause(&query, 3);
1750
1751 let full_sql = format!(
1760 "SELECT node_id, edge_id, relation, weight, dir FROM ({}){} \
1761 ORDER BY weight DESC, node_id ASC, \
1762 CASE dir WHEN 'out' THEN 0 ELSE 1 END ASC, edge_id ASC{}",
1763 sql, where_extra, limit_clause
1764 );
1765
1766 let mut stmt = conn.prepare(&full_sql)?;
1767
1768 let mut all_params: Vec<Box<dyn rusqlite::types::ToSql>> = Vec::new();
1769 all_params.push(Box::new(namespace.clone()));
1770 all_params.push(Box::new(node_str.clone()));
1771 all_params.extend(extra_params);
1772
1773 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
1774 all_params.iter().map(|p| p.as_ref()).collect();
1775
1776 let rows = stmt.query_map(param_refs.as_slice(), |row| {
1777 let nid_str: String = row.get(0)?;
1778 let eid_str: String = row.get(1)?;
1779 let relation_str: String = row.get(2)?;
1780 let weight: f64 = row.get(3)?;
1781 let dir_str: String = row.get(4)?;
1782 Ok((nid_str, eid_str, relation_str, weight, dir_str))
1783 })?;
1784
1785 let mut hits = Vec::new();
1786 for row in rows {
1787 let (nid_str, eid_str, relation_str, weight, dir_str) = row?;
1788 let relation = relation_str.parse::<EdgeRelation>().map_err(|e| {
1789 rusqlite::Error::FromSqlConversionFailure(
1790 2,
1791 rusqlite::types::Type::Text,
1792 Box::new(e),
1793 )
1794 })?;
1795 let direction = if dir_str == "out" {
1796 Direction::Out
1797 } else {
1798 Direction::In
1799 };
1800 hits.push(DirectedNeighborHit {
1801 hit: NeighborHit {
1802 node_id: parse_uuid(&nid_str)?,
1803 edge_id: parse_uuid(&eid_str)?,
1804 relation,
1805 weight,
1806 name: None,
1807 kind: None,
1808 entity_type: None,
1809 },
1810 direction,
1811 });
1812 }
1813
1814 Ok(hits)
1815 })
1816 .await
1817 }
1818
1819 async fn traverse(&self, request: TraversalRequest) -> Result<Vec<GraphPath>, StorageError> {
1820 use std::collections::{HashMap, HashSet};
1821
1822 use khive_storage::types::Direction;
1823
1824 if request.roots.is_empty() {
1825 return Ok(Vec::new());
1826 }
1827
1828 let roots = request.roots.clone();
1829 let opts = request.options.clone();
1830 let include_roots = request.include_roots;
1831 let namespace = self.namespace.clone();
1832 let max_depth_i64 =
1833 i64::try_from(opts.max_depth).map_err(|_| StorageError::InvalidInput {
1834 capability: StorageCapability::Graph,
1835 operation: "traverse".into(),
1836 message: format!(
1837 "TraversalOptions: max_depth must be <= i64::MAX, got {}",
1838 opts.max_depth
1839 ),
1840 })?;
1841
1842 let origin = self.pool.origin();
1843 self.with_reader("traverse", move |conn| {
1844 const CHUNK_ROOTS: usize = 400;
1857
1858 let (join_condition, next_node) = match opts.direction {
1860 Direction::Out => ("e.source_id = t.node_id", "e.target_id"),
1861 Direction::In => ("e.target_id = t.node_id", "e.source_id"),
1862 Direction::Both => (
1863 "(e.source_id = t.node_id OR e.target_id = t.node_id)",
1864 "CASE WHEN e.source_id = t.node_id THEN e.target_id ELSE e.source_id END",
1865 ),
1866 };
1867
1868 let _tx_handle = khive_storage::tx_registry::register_scoped(
1879 Some("graph_traverse_read".to_string()),
1880 origin,
1881 );
1882 let tx = conn.unchecked_transaction()?;
1883
1884 let mut root_data: HashMap<Uuid, (Vec<(PathNode, f64)>, HashSet<Uuid>)> =
1889 HashMap::with_capacity(roots.len());
1890
1891 for root_id in &roots {
1893 let (nodes, seen) = root_data.entry(*root_id).or_default();
1894 if include_roots {
1895 seen.insert(*root_id);
1896 nodes.push((
1897 PathNode {
1898 node_id: *root_id,
1899 via_edge: None,
1900 depth: 0,
1901 name: None,
1902 kind: None,
1903 properties: None,
1904 },
1905 0.0,
1906 ));
1907 }
1908 }
1909
1910 for chunk in roots.chunks(CHUNK_ROOTS) {
1911 let n_chunk = chunk.len();
1912
1913 let ns_param = n_chunk + 1;
1919 let depth_param = n_chunk + 2;
1920 let mut extra_param_idx = n_chunk + 3;
1921
1922 let mut relation_cond = String::new();
1923 let mut extra_params: Vec<Box<dyn rusqlite::types::ToSql>> = Vec::new();
1924
1925 if let Some(ref rels) = opts.relations {
1926 if !rels.is_empty() {
1927 let placeholders: Vec<String> = rels
1928 .iter()
1929 .map(|r| {
1930 extra_params.push(Box::new(r.to_string()));
1931 let p = format!("?{extra_param_idx}");
1932 extra_param_idx += 1;
1933 p
1934 })
1935 .collect();
1936 relation_cond = format!(" AND e.relation IN ({})", placeholders.join(","));
1937 }
1938 }
1939
1940 let mut weight_cond = String::new();
1941 if let Some(min_w) = opts.min_weight {
1942 extra_params.push(Box::new(min_w));
1943 weight_cond = format!(" AND e.weight >= ?{extra_param_idx}");
1944 }
1946
1947 let seed_rows: Vec<String> = (1..=n_chunk)
1950 .map(|i| format!("(?{i}, ?{i}, NULL, 0, ?{i}, 0.0)"))
1951 .collect();
1952 let seeds = seed_rows.join(", ");
1953
1954 let cte_sql = format!(
1958 "WITH RECURSIVE traversal(\
1959 root_id, node_id, edge_id, depth, path, total_weight\
1960 ) AS (\
1961 VALUES {seeds} \
1962 UNION ALL \
1963 SELECT t.root_id, {next_node}, e.id, t.depth + 1, \
1964 t.path || ',' || {next_node}, \
1965 t.total_weight + e.weight \
1966 FROM traversal t CROSS JOIN graph_edges e \
1967 ON {join_condition} \
1968 WHERE e.namespace = ?{ns} \
1969 AND e.deleted_at IS NULL \
1970 AND t.depth < ?{depth} \
1971 AND (',' || t.path || ',') NOT LIKE '%,' || {next_node} || ',%'\
1972 {rel_cond}{wt_cond} \
1973 ) \
1974 SELECT root_id, node_id, edge_id, depth, total_weight \
1975 FROM traversal WHERE depth > 0 \
1976 ORDER BY root_id, depth",
1977 seeds = seeds,
1978 next_node = next_node,
1979 join_condition = join_condition,
1980 ns = ns_param,
1981 depth = depth_param,
1982 rel_cond = relation_cond,
1983 wt_cond = weight_cond,
1984 );
1985
1986 let mut all_params: Vec<Box<dyn rusqlite::types::ToSql>> = Vec::new();
1987 for root_id in chunk {
1988 all_params.push(Box::new(root_id.to_string()));
1989 }
1990 all_params.push(Box::new(namespace.clone()));
1991 all_params.push(Box::new(max_depth_i64));
1992 all_params.extend(extra_params);
1993
1994 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
1995 all_params.iter().map(|p| p.as_ref()).collect();
1996
1997 let mut stmt = conn.prepare(&cte_sql)?;
2000 let rows_iter = stmt.query_map(param_refs.as_slice(), |row| {
2001 let root_str: String = row.get(0)?;
2002 let node_str: String = row.get(1)?;
2003 let edge_str: Option<String> = row.get(2)?;
2004 let depth: i64 = row.get(3)?;
2005 let total_weight: f64 = row.get(4)?;
2006 Ok((root_str, node_str, edge_str, depth, total_weight))
2007 })?;
2008
2009 for row in rows_iter {
2013 let (root_str, node_str, edge_str, depth, total_weight) = row?;
2014 let root_id = parse_uuid(&root_str)?;
2015 let node_id = parse_uuid(&node_str)?;
2016 let (nodes, seen) = root_data.entry(root_id).or_default();
2017 if !seen.insert(node_id) {
2018 continue;
2019 }
2020 let via_edge = edge_str.map(|s| parse_uuid(&s)).transpose()?;
2021 nodes.push((
2022 PathNode {
2023 node_id,
2024 via_edge,
2025 depth: depth as usize,
2026 name: None,
2027 kind: None,
2028 properties: None,
2029 },
2030 total_weight,
2031 ));
2032 }
2033 }
2034
2035 tx.commit()?;
2036
2037 let mut all_paths: Vec<GraphPath> = Vec::with_capacity(roots.len());
2043 for root_id in &roots {
2044 if let Some((mut nw, _)) = root_data.remove(root_id) {
2045 if nw.is_empty() {
2046 continue;
2047 }
2048 if let Some(lim) = opts.limit {
2049 let root_count = usize::from(include_roots);
2050 nw.truncate(root_count + lim as usize);
2051 }
2052 if nw.is_empty() {
2055 continue;
2056 }
2057 let max_weight = nw.iter().map(|(_, w)| *w).fold(0.0_f64, f64::max);
2058 let nodes: Vec<PathNode> = nw.into_iter().map(|(n, _)| n).collect();
2059 all_paths.push(GraphPath {
2060 root_id: *root_id,
2061 nodes,
2062 total_weight: max_weight,
2063 });
2064 }
2065 }
2066
2067 Ok(all_paths)
2068 })
2069 .await
2070 }
2071
2072 async fn purge_incident_edges(&self, node_id: Uuid) -> Result<u64, StorageError> {
2073 let statement = purge_incident_edges_statement(node_id);
2077 self.with_writer("purge_incident_edges", move |conn| {
2078 let mut stmt = conn.prepare(&statement.sql)?;
2079 bind_params(&mut stmt, &statement.params)?;
2080 Ok(stmt.raw_execute()? as u64)
2081 })
2082 .await
2083 }
2084}
2085
2086const GRAPH_DDL: &str = include_str!("../../sql/graph-ddl.sql");
2091
2092pub(crate) fn ensure_graph_schema(conn: &rusqlite::Connection) -> Result<(), rusqlite::Error> {
2093 conn.execute_batch(GRAPH_DDL)
2094}
2095
2096#[cfg(test)]
2097#[path = "graph_tests.rs"]
2098mod tests;