1use std::collections::{HashMap, HashSet, VecDeque};
4use std::sync::Arc;
5
6use async_trait::async_trait;
7use chrono::{DateTime, TimeZone, Utc};
8use rusqlite::OptionalExtension;
9use uuid::Uuid;
10
11use khive_storage::error::StorageError;
12use khive_storage::types::{
13 BatchWriteSummary, DeleteMode, DirectedNeighborHit, Direction, Edge, EdgeEndpointBaseCounts,
14 EdgeFilter, EdgeSeekPage, EdgeSortField, EdgeUpsertDisposition, EdgeUpsertRefusal,
15 EdgeUpsertRequest, EdgeUpsertResult, GraphPath, GuardedBatchOutcome, GuardedBatchRefusal,
16 GuardedEdgeBatchRefusal, GuardedEdgeBatchUpsertOutcome, GuardedEdgeUpsertOutcome,
17 GuardedWriteOutcome, MissingEndpoints, NeighborCursor, NeighborHit, NeighborQuery, Page,
18 PageRequest, PathNode, SeekCursor, SeekPage, SortDirection, SortOrder, SqlStatement, SqlValue,
19 TraversalExecutionBudget, TraversalOptions, TraversalRequest,
20};
21use khive_storage::GraphStore;
22use khive_storage::LinkId;
23use khive_storage::StorageCapability;
24use khive_types::EdgeRelation;
25
26use crate::error::SqliteError;
27use crate::pool::ConnectionPool;
28use crate::sql_bridge::bind_params;
29use crate::writer_task::WriterTaskHandle;
30
31fn map_err(e: rusqlite::Error, op: &'static str) -> StorageError {
33 StorageError::driver(StorageCapability::Graph, op, e)
34}
35
36fn map_sqlite_err(e: SqliteError, op: &'static str) -> StorageError {
37 StorageError::driver(StorageCapability::Graph, op, e)
38}
39
40fn resurrection_required_error(operation: &'static str, edge: &Edge) -> StorageError {
41 StorageError::Conflict {
42 capability: StorageCapability::Graph,
43 operation: operation.into(),
44 message: format!(
45 "edge {} is soft-deleted; explicit resurrection is required",
46 edge.id
47 ),
48 }
49}
50
51const NAMESPACE_COUNT_CHUNK_SIZE: usize = 500;
52
53const LATEST_ANNOTATING_NOTE_SQL: &str = r#"WITH incident AS MATERIALIZED (
62 SELECT source_id, deleted_at
63 FROM graph_edges INDEXED BY idx_graph_edges_ns_tgt_rel
64 WHERE namespace = ?1 AND target_id = ?2 AND relation = 'annotates'
65 LIMIT 3
66)
67SELECT result.id, result.created_at
68FROM notes AS result
69WHERE result.id = CASE WHEN (SELECT count(*) FROM incident) <= 2 THEN (
70 SELECT n.id
71 FROM incident AS e CROSS JOIN notes AS n
72 WHERE n.id = e.source_id AND e.deleted_at IS NULL
73 AND n.deleted_at IS NULL AND n.kind = ?3
74 AND EXISTS (SELECT 1 FROM json_each(CASE
75 WHEN json_type(n.properties, '$.tags') = 'array'
76 THEN json_extract(n.properties, '$.tags') ELSE '[]' END) AS tag
77 WHERE tag.type = 'text' AND tag.value = ?4 COLLATE BINARY)
78 ORDER BY n.created_at DESC, n.id ASC LIMIT 1
79) ELSE (
80 SELECT n.id
81 FROM notes AS n INDEXED BY idx_notes_created
82 WHERE n.deleted_at IS NULL AND n.kind = ?3
83 AND EXISTS (SELECT 1 FROM json_each(CASE
84 WHEN json_type(n.properties, '$.tags') = 'array'
85 THEN json_extract(n.properties, '$.tags') ELSE '[]' END) AS tag
86 WHERE tag.type = 'text' AND tag.value = ?4 COLLATE BINARY)
87 AND EXISTS (SELECT 1 FROM graph_edges AS e INDEXED BY idx_graph_edges_unique_triple
88 WHERE e.namespace = ?1 AND e.source_id = n.id AND e.target_id = ?2
89 AND e.relation = 'annotates' AND e.deleted_at IS NULL)
90 ORDER BY n.created_at DESC, n.id ASC LIMIT 1
91) END"#;
92
93fn edge_conflict_clause(resurrect: bool) -> String {
105 let deleted_at = if resurrect {
106 "NULL"
107 } else {
108 "graph_edges.deleted_at"
109 };
110 let predicate = if resurrect {
111 ""
112 } else {
113 " WHERE graph_edges.deleted_at IS NULL"
114 };
115 format!(
116 "weight = excluded.weight, \
117 updated_at = excluded.updated_at, \
118 deleted_at = {deleted_at}, \
119 metadata = excluded.metadata, \
120 target_backend = excluded.target_backend{predicate}"
121 )
122}
123
124fn endpoint_exists_clause(id_param: &str) -> String {
135 format!(
136 "EXISTS (SELECT 1 FROM entities WHERE id = {id_param} AND deleted_at IS NULL) \
137 OR EXISTS (SELECT 1 FROM notes WHERE id = {id_param} AND deleted_at IS NULL) \
138 OR EXISTS (SELECT 1 FROM events WHERE id = {id_param}) \
139 OR EXISTS (SELECT 1 FROM graph_edges WHERE id = {id_param} AND deleted_at IS NULL)"
140 )
141}
142
143pub fn edge_snapshot_assertion_statement(edge: &Edge, require_endpoints: bool) -> SqlStatement {
147 let mut sql = "SELECT 1 FROM graph_edges WHERE id=?1 AND namespace=?2 \
148 AND source_id=?3 AND target_id=?4 AND relation=?5 \
149 AND updated_at=?6 AND deleted_at IS NULL"
150 .to_string();
151 if require_endpoints {
152 sql.push_str(&format!(
153 " AND ({}) AND ({})",
154 endpoint_exists_clause("?3"),
155 endpoint_exists_clause("?4")
156 ));
157 }
158 SqlStatement {
159 sql,
160 params: vec![
161 SqlValue::Text(Uuid::from(edge.id).to_string()),
162 SqlValue::Text(edge.namespace.clone()),
163 SqlValue::Text(edge.source_id.to_string()),
164 SqlValue::Text(edge.target_id.to_string()),
165 SqlValue::Text(edge.relation.to_string()),
166 SqlValue::Integer(edge.updated_at.timestamp_micros()),
167 ],
168 label: Some("edge-snapshot-assertion".into()),
169 }
170}
171
172pub fn edge_upsert_statement(edge: &Edge) -> SqlStatement {
176 edge_upsert_statement_with_resurrection(edge, false)
177}
178
179pub fn edge_upsert_statement_with_resurrection(edge: &Edge, resurrect: bool) -> SqlStatement {
181 let (source_id, target_id) =
182 canonical_edge_endpoints(edge.relation, edge.source_id, edge.target_id);
183 let metadata_str = edge
184 .metadata
185 .as_ref()
186 .map(|v| serde_json::to_string(v).unwrap_or_default());
187 let conflict_clause = edge_conflict_clause(resurrect);
188 SqlStatement {
189 sql: format!(
190 "INSERT INTO graph_edges \
191 (namespace, id, source_id, target_id, relation, weight, \
192 created_at, updated_at, deleted_at, metadata, target_backend) \
193 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11) \
194 ON CONFLICT(namespace, id) DO UPDATE SET \
195 source_id = excluded.source_id, \
196 target_id = excluded.target_id, \
197 relation = excluded.relation, \
198 {conflict_clause} \
199 ON CONFLICT(namespace, source_id, target_id, relation) DO UPDATE SET \
200 {conflict_clause}"
201 ),
202 params: vec![
203 SqlValue::Text(edge.namespace.clone()),
204 SqlValue::Text(Uuid::from(edge.id).to_string()),
205 SqlValue::Text(source_id.to_string()),
206 SqlValue::Text(target_id.to_string()),
207 SqlValue::Text(edge.relation.to_string()),
208 SqlValue::Float(edge.weight),
209 SqlValue::Integer(edge.created_at.timestamp_micros()),
210 SqlValue::Integer(edge.updated_at.timestamp_micros()),
211 match edge.deleted_at {
212 Some(t) => SqlValue::Integer(t.timestamp_micros()),
213 None => SqlValue::Null,
214 },
215 match metadata_str {
216 Some(m) => SqlValue::Text(m),
217 None => SqlValue::Null,
218 },
219 match &edge.target_backend {
220 Some(b) => SqlValue::Text(b.clone()),
221 None => SqlValue::Null,
222 },
223 ],
224 label: Some("edge-upsert".to_string()),
225 }
226}
227
228pub fn edge_insert_only_guarded_by_endpoints_statement(edge: &Edge) -> SqlStatement {
233 let mut statement = edge_upsert_statement(edge);
234 let src_exists = endpoint_exists_clause("?3");
235 let tgt_exists = endpoint_exists_clause("?4");
236 statement.sql = format!(
237 "INSERT INTO graph_edges \
238 (namespace, id, source_id, target_id, relation, weight, \
239 created_at, updated_at, deleted_at, metadata, target_backend) \
240 SELECT ?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11 \
241 WHERE ({src_exists}) AND ({tgt_exists})"
242 );
243 statement.label = Some("edge-insert-only-where-endpoints-exist".to_string());
244 statement
245}
246
247pub fn edge_insert_if_absent_statement(edge: &Edge) -> SqlStatement {
251 let mut statement = edge_upsert_statement(edge);
252 statement.sql = "INSERT INTO graph_edges \
253 (namespace, id, source_id, target_id, relation, weight, \
254 created_at, updated_at, deleted_at, metadata, target_backend) \
255 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11) \
256 ON CONFLICT DO NOTHING"
257 .to_string();
258 statement.label = Some("edge-insert-if-absent".to_string());
259 statement
260}
261
262pub fn edge_replace_if_unchanged_statement(
273 edge: &Edge,
274 expected_updated_at: DateTime<Utc>,
275 expected_deleted_at: Option<DateTime<Utc>>,
276) -> SqlStatement {
277 let (source_id, target_id) =
278 canonical_edge_endpoints(edge.relation, edge.source_id, edge.target_id);
279 let metadata_str = edge
280 .metadata
281 .as_ref()
282 .map(|v| serde_json::to_string(v).unwrap_or_default());
283 SqlStatement {
284 sql: "UPDATE graph_edges SET \
285 namespace = ?1, source_id = ?2, target_id = ?3, relation = ?4, weight = ?5, \
286 updated_at = ?6, deleted_at = ?7, metadata = ?8, target_backend = ?9 \
287 WHERE id = ?10 AND updated_at = ?11 AND deleted_at IS ?12 \
288 AND ?6 > updated_at"
289 .to_string(),
290 params: vec![
291 SqlValue::Text(edge.namespace.clone()),
292 SqlValue::Text(source_id.to_string()),
293 SqlValue::Text(target_id.to_string()),
294 SqlValue::Text(edge.relation.to_string()),
295 SqlValue::Float(edge.weight),
296 SqlValue::Integer(edge.updated_at.timestamp_micros()),
297 match edge.deleted_at {
298 Some(t) => SqlValue::Integer(t.timestamp_micros()),
299 None => SqlValue::Null,
300 },
301 match metadata_str {
302 Some(m) => SqlValue::Text(m),
303 None => SqlValue::Null,
304 },
305 match &edge.target_backend {
306 Some(b) => SqlValue::Text(b.clone()),
307 None => SqlValue::Null,
308 },
309 SqlValue::Text(Uuid::from(edge.id).to_string()),
310 SqlValue::Integer(expected_updated_at.timestamp_micros()),
311 match expected_deleted_at {
312 Some(value) => SqlValue::Integer(value.timestamp_micros()),
313 None => SqlValue::Null,
314 },
315 ],
316 label: Some("edge-replace-if-unchanged".to_string()),
317 }
318}
319
320#[allow(clippy::too_many_arguments)]
340pub fn edge_insert_guarded_by_endpoints_statement(
341 namespace: &str,
342 edge_id: Uuid,
343 source_id: Uuid,
344 target_id: Uuid,
345 relation: EdgeRelation,
346 weight: f64,
347 now: i64,
348 metadata: Option<&str>,
349) -> SqlStatement {
350 edge_insert_guarded_by_endpoints_with_resurrection_statement(
351 namespace, source_id, target_id, edge_id, relation, weight, now, metadata, false,
352 )
353}
354
355#[allow(clippy::too_many_arguments)]
357pub fn edge_insert_guarded_by_endpoints_with_resurrection_statement(
358 namespace: &str,
359 source_id: Uuid,
360 target_id: Uuid,
361 edge_id: Uuid,
362 relation: EdgeRelation,
363 weight: f64,
364 now: i64,
365 metadata: Option<&str>,
366 resurrect: bool,
367) -> SqlStatement {
368 let src_exists = endpoint_exists_clause("?3");
369 let tgt_exists = endpoint_exists_clause("?4");
370 let conflict_clause = edge_conflict_clause(resurrect);
371 SqlStatement {
372 sql: format!(
373 "INSERT INTO graph_edges \
374 (namespace, id, source_id, target_id, relation, weight, \
375 created_at, updated_at, metadata) \
376 SELECT ?1, ?2, ?3, ?4, ?5, ?6, ?7, ?7, ?8 \
377 WHERE ({src_exists}) AND ({tgt_exists}) \
378 ON CONFLICT(namespace, source_id, target_id, relation) DO UPDATE SET \
379 {conflict_clause}"
380 ),
381 params: vec![
382 SqlValue::Text(namespace.to_string()),
383 SqlValue::Text(edge_id.to_string()),
384 SqlValue::Text(source_id.to_string()),
385 SqlValue::Text(target_id.to_string()),
386 SqlValue::Text(relation.as_str().to_string()),
387 SqlValue::Float(weight),
388 SqlValue::Integer(now),
389 match metadata {
390 Some(m) => SqlValue::Text(m.to_string()),
391 None => SqlValue::Null,
392 },
393 ],
394 label: Some("atomic-link-insert-edge-where-exists".to_string()),
395 }
396}
397
398#[allow(clippy::too_many_arguments)]
402pub fn edge_insert_new_guarded_by_endpoints_statement(
403 namespace: &str,
404 edge_id: Uuid,
405 source_id: Uuid,
406 target_id: Uuid,
407 relation: EdgeRelation,
408 weight: f64,
409 now: i64,
410 metadata: Option<&str>,
411) -> SqlStatement {
412 let src_exists = endpoint_exists_clause("?3");
413 let tgt_exists = endpoint_exists_clause("?4");
414 SqlStatement {
415 sql: format!(
416 "INSERT INTO graph_edges \
417 (namespace, id, source_id, target_id, relation, weight, \
418 created_at, updated_at, metadata) \
419 SELECT ?1, ?2, ?3, ?4, ?5, ?6, ?7, ?7, ?8 \
420 WHERE ({src_exists}) AND ({tgt_exists}) \
421 ON CONFLICT DO NOTHING"
422 ),
423 params: vec![
424 SqlValue::Text(namespace.to_string()),
425 SqlValue::Text(edge_id.to_string()),
426 SqlValue::Text(source_id.to_string()),
427 SqlValue::Text(target_id.to_string()),
428 SqlValue::Text(relation.to_string()),
429 SqlValue::Float(weight),
430 SqlValue::Integer(now),
431 metadata
432 .map(|value| SqlValue::Text(value.to_string()))
433 .unwrap_or(SqlValue::Null),
434 ],
435 label: Some("edge-link-create-if-absent-and-endpoints-exist".to_string()),
436 }
437}
438
439pub fn edge_link_replace_if_unchanged_and_endpoints_exist_statement(
445 previous: &Edge,
446 weight: f64,
447 now: i64,
448 metadata: Option<&str>,
449) -> SqlStatement {
450 let src_exists = endpoint_exists_clause("?6");
451 let tgt_exists = endpoint_exists_clause("?7");
452 SqlStatement {
453 sql: format!(
454 "UPDATE graph_edges SET \
455 weight = ?1, updated_at = ?2, deleted_at = NULL, \
456 metadata = ?3, target_backend = NULL \
457 WHERE namespace = ?4 AND id = ?5 \
458 AND source_id = ?6 AND target_id = ?7 AND relation = ?8 \
459 AND updated_at = ?9 AND deleted_at IS ?10 AND ?2 > updated_at \
460 AND ({src_exists}) AND ({tgt_exists})"
461 ),
462 params: vec![
463 SqlValue::Float(weight),
464 SqlValue::Integer(now),
465 metadata
466 .map(|value| SqlValue::Text(value.to_string()))
467 .unwrap_or(SqlValue::Null),
468 SqlValue::Text(previous.namespace.clone()),
469 SqlValue::Text(Uuid::from(previous.id).to_string()),
470 SqlValue::Text(previous.source_id.to_string()),
471 SqlValue::Text(previous.target_id.to_string()),
472 SqlValue::Text(previous.relation.to_string()),
473 SqlValue::Integer(previous.updated_at.timestamp_micros()),
474 previous
475 .deleted_at
476 .map(|value| SqlValue::Integer(value.timestamp_micros()))
477 .unwrap_or(SqlValue::Null),
478 ],
479 label: Some("edge-link-replace-if-unchanged-and-endpoints-exist".to_string()),
480 }
481}
482
483pub fn edge_soft_delete_statement(id: Uuid, now: i64) -> SqlStatement {
485 SqlStatement {
486 sql: "UPDATE graph_edges SET deleted_at = ?2, updated_at = ?2 \
487 WHERE id = ?1 AND deleted_at IS NULL"
488 .to_string(),
489 params: vec![SqlValue::Text(id.to_string()), SqlValue::Integer(now)],
490 label: Some("edge-delete-soft".to_string()),
491 }
492}
493
494pub fn edge_hard_delete_statement(id: Uuid) -> SqlStatement {
496 SqlStatement {
497 sql: "DELETE FROM graph_edges WHERE id = ?1".to_string(),
498 params: vec![SqlValue::Text(id.to_string())],
499 label: Some("edge-delete-hard".to_string()),
500 }
501}
502
503pub fn purge_incident_edges_statement(node_id: Uuid) -> SqlStatement {
505 SqlStatement {
506 sql: "DELETE FROM graph_edges WHERE source_id = ?1 OR target_id = ?1".to_string(),
507 params: vec![SqlValue::Text(node_id.to_string())],
508 label: Some("edge-purge-incident".to_string()),
509 }
510}
511
512pub const EDGE_SYMMETRIC_CONFLICT_PROBE_SQL: &str = "SELECT id FROM graph_edges \
533 WHERE namespace = ?1 AND source_id = ?2 AND target_id = ?3 \
534 AND relation = ?4 AND id != ?5";
535
536pub const EDGE_SYMMETRIC_DELETE_NONCANONICAL_SQL: &str =
537 "DELETE FROM graph_edges WHERE namespace = ?1 AND id = ?2";
538
539pub const EDGE_SYMMETRIC_DELETE_NONCANONICAL_GUARDED_SQL: &str =
552 "DELETE FROM graph_edges WHERE namespace = ?1 AND id = ?2 \
553 AND updated_at = ?3 AND deleted_at IS ?4";
554
555pub const EDGE_SYMMETRIC_UPDATE_INPLACE_SQL: &str = "UPDATE graph_edges SET \
562 source_id = ?1, target_id = ?2, relation = ?3, \
563 weight = ?4, updated_at = ?5, metadata = ?6 \
564 WHERE namespace = ?7 AND id = ?8 \
565 AND updated_at = ?9 AND deleted_at IS ?10 \
566 AND ?5 > updated_at";
567
568pub fn edge_symmetric_conflict_probe_statement(
571 namespace: &str,
572 canon_src: Uuid,
573 canon_tgt: Uuid,
574 relation: EdgeRelation,
575 exclude_id: Uuid,
576) -> SqlStatement {
577 SqlStatement {
578 sql: EDGE_SYMMETRIC_CONFLICT_PROBE_SQL.to_string(),
579 params: vec![
580 SqlValue::Text(namespace.to_string()),
581 SqlValue::Text(canon_src.to_string()),
582 SqlValue::Text(canon_tgt.to_string()),
583 SqlValue::Text(relation.to_string()),
584 SqlValue::Text(exclude_id.to_string()),
585 ],
586 label: Some("edge-symmetric-conflict-probe".to_string()),
587 }
588}
589
590pub fn edge_symmetric_delete_noncanonical_statement(namespace: &str, id: Uuid) -> SqlStatement {
593 SqlStatement {
594 sql: EDGE_SYMMETRIC_DELETE_NONCANONICAL_SQL.to_string(),
595 params: vec![
596 SqlValue::Text(namespace.to_string()),
597 SqlValue::Text(id.to_string()),
598 ],
599 label: Some("edge-symmetric-delete-noncanonical".to_string()),
600 }
601}
602
603#[allow(clippy::too_many_arguments)]
607pub fn edge_symmetric_update_inplace_statement(
608 namespace: &str,
609 id: Uuid,
610 canon_src: Uuid,
611 canon_tgt: Uuid,
612 relation: EdgeRelation,
613 weight: f64,
614 updated_at_micros: i64,
615 metadata: Option<&str>,
616 expected_updated_at_micros: i64,
617 expected_deleted_at_micros: Option<i64>,
618) -> SqlStatement {
619 SqlStatement {
620 sql: EDGE_SYMMETRIC_UPDATE_INPLACE_SQL.to_string(),
621 params: vec![
622 SqlValue::Text(canon_src.to_string()),
623 SqlValue::Text(canon_tgt.to_string()),
624 SqlValue::Text(relation.to_string()),
625 SqlValue::Float(weight),
626 SqlValue::Integer(updated_at_micros),
627 match metadata {
628 Some(m) => SqlValue::Text(m.to_string()),
629 None => SqlValue::Null,
630 },
631 SqlValue::Text(namespace.to_string()),
632 SqlValue::Text(id.to_string()),
633 SqlValue::Integer(expected_updated_at_micros),
634 match expected_deleted_at_micros {
635 Some(value) => SqlValue::Integer(value),
636 None => SqlValue::Null,
637 },
638 ],
639 label: Some("edge-symmetric-update-inplace".to_string()),
640 }
641}
642
643#[allow(clippy::too_many_arguments)]
737pub fn edge_symmetric_delete_if_conflict_statement(
738 namespace: &str,
739 id: Uuid,
740 canon_src: Uuid,
741 canon_tgt: Uuid,
742 relation: EdgeRelation,
743 expected_updated_at_micros: i64,
744 expected_deleted_at_micros: Option<i64>,
745) -> SqlStatement {
746 SqlStatement {
747 sql: "DELETE FROM graph_edges \
748 WHERE namespace = ?1 AND id = ?2 \
749 AND updated_at = ?6 AND deleted_at IS ?7 \
750 AND EXISTS ( \
751 SELECT 1 FROM graph_edges \
752 WHERE namespace = ?1 AND source_id = ?3 AND target_id = ?4 \
753 AND relation = ?5 AND id != ?2 \
754 )"
755 .to_string(),
756 params: vec![
757 SqlValue::Text(namespace.to_string()),
758 SqlValue::Text(id.to_string()),
759 SqlValue::Text(canon_src.to_string()),
760 SqlValue::Text(canon_tgt.to_string()),
761 SqlValue::Text(relation.to_string()),
762 SqlValue::Integer(expected_updated_at_micros),
763 match expected_deleted_at_micros {
764 Some(value) => SqlValue::Integer(value),
765 None => SqlValue::Null,
766 },
767 ],
768 label: Some("edge-symmetric-delete-if-conflict".to_string()),
769 }
770}
771
772#[allow(clippy::too_many_arguments)]
781pub fn edge_symmetric_absorb_or_update_inplace_statement(
782 namespace: &str,
783 id: Uuid,
784 canon_src: Uuid,
785 canon_tgt: Uuid,
786 relation: EdgeRelation,
787 weight: f64,
788 updated_at_micros: i64,
789 metadata: Option<&str>,
790 target_backend: Option<&str>,
791 expected_updated_at_micros: i64,
792 expected_deleted_at_micros: Option<i64>,
793) -> SqlStatement {
794 SqlStatement {
795 sql: "UPDATE graph_edges SET \
796 source_id = CASE WHEN id = ?2 THEN ?3 ELSE source_id END, \
797 target_id = CASE WHEN id = ?2 THEN ?4 ELSE target_id END, \
798 relation = CASE WHEN id = ?2 THEN ?5 ELSE relation END, \
799 weight = CASE WHEN id = ?2 THEN ?6 ELSE weight END, \
800 updated_at = CASE WHEN id = ?2 THEN ?7 ELSE updated_at END, \
801 deleted_at = CASE WHEN id = ?2 THEN NULL ELSE deleted_at END, \
802 metadata = CASE WHEN id = ?2 THEN ?8 ELSE metadata END, \
803 target_backend = CASE WHEN id = ?2 THEN ?9 ELSE target_backend END \
804 WHERE namespace = ?1 \
805 AND ( \
806 (id = ?2 AND changes() = 0 AND updated_at = ?10 AND deleted_at IS ?11 \
807 AND ?7 > updated_at) \
808 OR (source_id = ?3 AND target_id = ?4 AND relation = ?5 \
809 AND id != ?2 AND changes() = 1) \
810 )"
811 .to_string(),
812 params: vec![
813 SqlValue::Text(namespace.to_string()),
814 SqlValue::Text(id.to_string()),
815 SqlValue::Text(canon_src.to_string()),
816 SqlValue::Text(canon_tgt.to_string()),
817 SqlValue::Text(relation.to_string()),
818 SqlValue::Float(weight),
819 SqlValue::Integer(updated_at_micros),
820 match metadata {
821 Some(m) => SqlValue::Text(m.to_string()),
822 None => SqlValue::Null,
823 },
824 match target_backend {
825 Some(b) => SqlValue::Text(b.to_string()),
826 None => SqlValue::Null,
827 },
828 SqlValue::Integer(expected_updated_at_micros),
829 match expected_deleted_at_micros {
830 Some(value) => SqlValue::Integer(value),
831 None => SqlValue::Null,
832 },
833 ],
834 label: Some("edge-symmetric-absorb-or-update-inplace".to_string()),
835 }
836}
837
838pub struct SqlGraphStore {
840 pool: Arc<ConnectionPool>,
841 is_file_backed: bool,
842 namespace: String,
846 writer_task: Option<WriterTaskHandle>,
847}
848
849impl SqlGraphStore {
850 pub fn new_scoped(
858 pool: Arc<ConnectionPool>,
859 is_file_backed: bool,
860 namespace: impl Into<String>,
861 ) -> Self {
862 let writer_task = pool.writer_task_handle().ok().flatten();
868
869 Self {
870 pool,
871 is_file_backed,
872 namespace: namespace.into(),
873 writer_task,
874 }
875 }
876
877 fn open_standalone_writer(&self) -> Result<rusqlite::Connection, StorageError> {
878 self.pool
879 .open_standalone_writer()
880 .map_err(|e| map_sqlite_err(e, "open_graph_writer"))
881 }
882
883 fn current_writer_task(
884 &self,
885 operation: &'static str,
886 ) -> Result<Option<WriterTaskHandle>, StorageError> {
887 self.pool
888 .writer_task_for_write(self.writer_task.as_ref(), operation)
889 }
890
891 async fn with_writer<F, R>(&self, op: &'static str, f: F) -> Result<R, StorageError>
898 where
899 F: FnOnce(&rusqlite::Connection) -> Result<R, rusqlite::Error> + Send + 'static,
900 R: Send + 'static,
901 {
902 if let Some(writer_task) = self.current_writer_task(op)? {
903 return writer_task
904 .send_bounded(move |conn| f(conn).map_err(|e| map_err(e, op)))
905 .await;
906 }
907
908 self.pool
909 .record_direct_route(crate::timeout_sink::Site::DirectRouteGraphGeneralWrite);
910 if self.is_file_backed {
911 let conn = self.open_standalone_writer()?;
912 let db = crate::timeout_sink::db_label(&self.pool);
913 tokio::task::spawn_blocking(move || {
914 f(&conn).map_err(|e| {
915 crate::timeout_sink::maybe_emit_busy(
916 &db,
917 crate::timeout_sink::Site::StandaloneGraph,
918 &e,
919 );
920 map_err(e, op)
921 })
922 })
923 .await
924 .map_err(|e| StorageError::driver(StorageCapability::Graph, op, e))?
925 } else {
926 let pool = Arc::clone(&self.pool);
927 tokio::task::spawn_blocking(move || {
928 let guard = pool.try_writer().map_err(|e| map_sqlite_err(e, op))?;
929 f(guard.conn()).map_err(|e| map_err(e, op))
930 })
931 .await
932 .map_err(|e| StorageError::driver(StorageCapability::Graph, op, e))?
933 }
934 }
935
936 async fn observed_edge_write(
937 &self,
938 operation: &'static str,
939 request: EdgeUpsertRequest,
940 guard_endpoints: bool,
941 ) -> Result<GuardedEdgeUpsertOutcome, StorageError> {
942 if let Some(writer_task) = self.current_writer_task(operation)? {
943 return writer_task
944 .send_bounded(move |conn| {
945 observed_edge_upsert(conn, &request, guard_endpoints)
946 .map_err(|error| map_err(error, operation))
947 })
948 .await;
949 }
950
951 let origin = self.pool.origin();
952 self.with_writer(operation, move |conn| {
953 conn.execute_batch("BEGIN IMMEDIATE")?;
954 let _tx_handle =
955 khive_storage::tx_registry::register_scoped(Some(operation.to_string()), origin);
956 let outcome = match observed_edge_upsert(conn, &request, guard_endpoints) {
957 Ok(outcome) => outcome,
958 Err(error) => {
959 let _ = conn.execute_batch("ROLLBACK");
960 return Err(error);
961 }
962 };
963 if let Err(error) = conn.execute_batch("COMMIT") {
964 let _ = conn.execute_batch("ROLLBACK");
965 return Err(error);
966 }
967 Ok(outcome)
968 })
969 .await
970 }
971
972 async fn observed_edge_batch_write(
973 &self,
974 operation: &'static str,
975 requests: Vec<EdgeUpsertRequest>,
976 guard_endpoints: bool,
977 ) -> Result<GuardedEdgeBatchUpsertOutcome, StorageError> {
978 if let Some(writer_task) = self.current_writer_task(operation)? {
979 return writer_task
980 .send_bounded(move |conn| {
981 observed_edge_batch_upsert(conn, &requests, guard_endpoints)
982 .map_err(|error| map_err(error, operation))
983 })
984 .await;
985 }
986
987 let origin = self.pool.origin();
988 self.with_writer(operation, move |conn| {
989 conn.execute_batch("BEGIN IMMEDIATE")?;
990 let _tx_handle =
991 khive_storage::tx_registry::register_scoped(Some(operation.to_string()), origin);
992 let outcome = match observed_edge_batch_upsert(conn, &requests, guard_endpoints) {
993 Ok(outcome) => outcome,
994 Err(error) => {
995 let _ = conn.execute_batch("ROLLBACK");
996 return Err(error);
997 }
998 };
999 if let Err(error) = conn.execute_batch("COMMIT") {
1000 let _ = conn.execute_batch("ROLLBACK");
1001 return Err(error);
1002 }
1003 Ok(outcome)
1004 })
1005 .await
1006 }
1007
1008 async fn with_reader<F, R>(&self, op: &'static str, f: F) -> Result<R, StorageError>
1009 where
1010 F: FnOnce(&rusqlite::Connection) -> Result<R, rusqlite::Error> + Send + 'static,
1011 R: Send + 'static,
1012 {
1013 super::run_pooled_store_read(
1014 Arc::clone(&self.pool),
1015 StorageCapability::Graph,
1016 op,
1017 move |conn| f(conn).map_err(|error| map_err(error, op)),
1018 )
1019 .await
1020 }
1021}
1022
1023fn report_graph_usage(queries: &std::sync::atomic::AtomicU64, rows: &std::sync::atomic::AtomicU64) {
1039 khive_storage::usage::count(
1040 khive_storage::usage::UsageUnit::DbRoundTrips,
1041 queries.load(std::sync::atomic::Ordering::Relaxed),
1042 );
1043 khive_storage::usage::count(
1044 khive_storage::usage::UsageUnit::GraphHops,
1045 rows.load(std::sync::atomic::Ordering::Relaxed),
1046 );
1047}
1048
1049fn read_edge(row: &rusqlite::Row<'_>) -> Result<Edge, rusqlite::Error> {
1050 let namespace: String = row.get(0)?;
1051 let id_str: String = row.get(1)?;
1052 let source_str: String = row.get(2)?;
1053 let target_str: String = row.get(3)?;
1054 let relation_str: String = row.get(4)?;
1055 let weight: f64 = row.get(5)?;
1056 let created_micros: i64 = row.get(6)?;
1057 let updated_micros: i64 = row.get(7)?;
1058 let deleted_micros: Option<i64> = row.get(8)?;
1059 let metadata_str: Option<String> = row.get(9)?;
1060 let target_backend: Option<String> = row.get(10)?;
1061
1062 let id = parse_uuid(&id_str)?;
1063 let source_id = parse_uuid(&source_str)?;
1064 let target_id = parse_uuid(&target_str)?;
1065 let created_at = micros_to_datetime(created_micros);
1066 let relation = relation_str.parse::<EdgeRelation>().map_err(|e| {
1067 rusqlite::Error::FromSqlConversionFailure(4, rusqlite::types::Type::Text, Box::new(e))
1068 })?;
1069 let metadata = match metadata_str {
1070 Some(s) => {
1071 let v = serde_json::from_str(&s).map_err(|e| {
1072 rusqlite::Error::FromSqlConversionFailure(
1073 9,
1074 rusqlite::types::Type::Text,
1075 Box::new(e),
1076 )
1077 })?;
1078 Some(v)
1079 }
1080 None => None,
1081 };
1082
1083 Ok(Edge {
1084 id: id.into(),
1085 namespace,
1086 source_id,
1087 target_id,
1088 relation,
1089 weight,
1090 created_at,
1091 updated_at: micros_to_datetime(updated_micros),
1092 deleted_at: deleted_micros.map(micros_to_datetime),
1093 metadata,
1094 target_backend,
1095 })
1096}
1097
1098fn parse_uuid(s: &str) -> Result<Uuid, rusqlite::Error> {
1099 Uuid::parse_str(s).map_err(|e| {
1100 rusqlite::Error::FromSqlConversionFailure(0, rusqlite::types::Type::Text, Box::new(e))
1101 })
1102}
1103
1104fn neighbor_extra_clause(
1110 query: &NeighborQuery,
1111 start_param_idx: usize,
1112 after: Option<&NeighborCursor>,
1113 neighbor_kinds: Option<&[String]>,
1114) -> (String, String, Vec<Box<dyn rusqlite::types::ToSql>>) {
1115 let mut conditions: Vec<String> = Vec::new();
1116 let mut extra_params: Vec<Box<dyn rusqlite::types::ToSql>> = Vec::new();
1117 let mut param_idx = start_param_idx;
1118
1119 if let Some(ref rels) = query.relations {
1120 if !rels.is_empty() {
1121 let placeholders: Vec<String> = rels
1122 .iter()
1123 .map(|r| {
1124 extra_params.push(Box::new(r.to_string()));
1125 let p = format!("?{}", param_idx);
1126 param_idx += 1;
1127 p
1128 })
1129 .collect();
1130 conditions.push(format!("relation IN ({})", placeholders.join(",")));
1131 }
1132 }
1133
1134 if let Some(min_w) = query.min_weight {
1135 extra_params.push(Box::new(min_w));
1136 conditions.push(format!("weight >= ?{}", param_idx));
1137 param_idx += 1;
1138 }
1139
1140 if let Some(cursor) = after {
1141 extra_params.push(Box::new(cursor.weight));
1142 let weight_idx = param_idx;
1143 param_idx += 1;
1144 extra_params.push(Box::new(cursor.node_id.to_string()));
1145 let node_idx = param_idx;
1146 param_idx += 1;
1147 extra_params.push(Box::new(cursor.edge_id.to_string()));
1148 let edge_idx = param_idx;
1149 param_idx += 1;
1150 conditions.push(format!(
1151 "(weight < ?{weight_idx} OR (weight = ?{weight_idx} AND node_id > ?{node_idx}) OR (weight = ?{weight_idx} AND node_id = ?{node_idx} AND edge_id > ?{edge_idx}))"
1152 ));
1153 }
1154
1155 if let Some(kinds) = neighbor_kinds.filter(|kinds| !kinds.is_empty()) {
1156 let placeholders: Vec<String> = kinds
1157 .iter()
1158 .map(|kind| {
1159 extra_params.push(Box::new(kind.clone()));
1160 let p = format!("?{param_idx}");
1161 param_idx += 1;
1162 p
1163 })
1164 .collect();
1165 let entity_placeholders = placeholders.join(",");
1166 let note_placeholders: Vec<String> = kinds
1167 .iter()
1168 .map(|kind| {
1169 extra_params.push(Box::new(kind.clone()));
1170 let p = format!("?{param_idx}");
1171 param_idx += 1;
1172 p
1173 })
1174 .collect();
1175 conditions.push(format!(
1176 "(EXISTS (SELECT 1 FROM entities AS neighbor_entities WHERE neighbor_entities.id = node_id AND neighbor_entities.namespace = ?1 AND neighbor_entities.deleted_at IS NULL AND neighbor_entities.kind IN ({entity_placeholders})) OR EXISTS (SELECT 1 FROM notes AS neighbor_notes WHERE neighbor_notes.id = node_id AND neighbor_notes.namespace = ?1 AND neighbor_notes.deleted_at IS NULL AND neighbor_notes.kind IN ({})))",
1177 note_placeholders.join(",")
1178 ));
1179 }
1180
1181 let where_extra = if conditions.is_empty() {
1182 String::new()
1183 } else {
1184 format!(" WHERE {}", conditions.join(" AND "))
1185 };
1186
1187 let limit_clause = if let Some(lim) = query.limit {
1188 extra_params.push(Box::new(lim as i64));
1189 format!(" LIMIT ?{}", param_idx)
1190 } else {
1191 String::new()
1192 };
1193
1194 (where_extra, limit_clause, extra_params)
1195}
1196
1197#[cfg(test)]
1205static NEIGHBOR_SELECT_COUNT: std::sync::atomic::AtomicUsize =
1206 std::sync::atomic::AtomicUsize::new(0);
1207
1208#[cfg(test)]
1209fn count_neighbor_select() {
1210 NEIGHBOR_SELECT_COUNT.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1211}
1212
1213#[cfg(not(test))]
1214fn count_neighbor_select() {}
1215
1216#[cfg(test)]
1217pub(crate) fn reset_neighbor_select_count() {
1218 NEIGHBOR_SELECT_COUNT.store(0, std::sync::atomic::Ordering::Relaxed);
1219}
1220
1221#[cfg(test)]
1222pub(crate) fn neighbor_select_count() -> usize {
1223 NEIGHBOR_SELECT_COUNT.load(std::sync::atomic::Ordering::Relaxed)
1224}
1225
1226fn micros_to_datetime(micros: i64) -> DateTime<Utc> {
1227 Utc.timestamp_micros(micros)
1228 .single()
1229 .unwrap_or_else(Utc::now)
1230}
1231
1232fn edge_order_clause(sort: &[SortOrder<EdgeSortField>]) -> String {
1239 if sort.is_empty() {
1240 return " ORDER BY created_at DESC, id DESC".to_string();
1241 }
1242 let mut parts: Vec<String> = sort
1243 .iter()
1244 .map(|s| {
1245 let dir = match s.direction {
1246 SortDirection::Asc => "ASC",
1247 SortDirection::Desc => "DESC",
1248 };
1249 format!("{} {}", edge_sort_col(&s.field), dir)
1250 })
1251 .collect();
1252 let dir = match sort.last().map(|s| &s.direction) {
1253 Some(SortDirection::Asc) => "ASC",
1254 _ => "DESC",
1255 };
1256 parts.push(format!("id {dir}"));
1257 format!(" ORDER BY {}", parts.join(", "))
1258}
1259
1260fn endpoint_base_case(column: &str) -> String {
1278 format!(
1279 "CASE WHEN EXISTS (SELECT 1 FROM entities be WHERE be.id = graph_edges.{column}) \
1280 THEN 'entity' \
1281 WHEN EXISTS (SELECT 1 FROM notes bn WHERE bn.id = graph_edges.{column}) \
1282 THEN 'note' ELSE 'none' END"
1283 )
1284}
1285
1286fn fold_endpoint_base_row(counts: &mut EdgeEndpointBaseCounts, source: &str, target: &str, n: u64) {
1290 let slot = match (source, target) {
1291 ("entity", "entity") => &mut counts.entity_entity,
1292 ("entity", "note") => &mut counts.entity_note,
1293 ("note", "entity") => &mut counts.note_entity,
1294 ("note", "note") => &mut counts.note_note,
1295 _ => &mut counts.unresolved,
1296 };
1297 *slot = slot.saturating_add(n);
1298}
1299
1300const LIVE_ENDPOINTS_CONDITION: &str = "NOT EXISTS (SELECT 1 FROM entities le \
1313 WHERE le.id = graph_edges.source_id AND le.deleted_at IS NOT NULL) \
1314 AND NOT EXISTS (SELECT 1 FROM entities le \
1315 WHERE le.id = graph_edges.target_id AND le.deleted_at IS NOT NULL) \
1316 AND NOT EXISTS (SELECT 1 FROM notes ln \
1317 WHERE ln.id = graph_edges.source_id AND ln.deleted_at IS NOT NULL) \
1318 AND NOT EXISTS (SELECT 1 FROM notes ln \
1319 WHERE ln.id = graph_edges.target_id AND ln.deleted_at IS NOT NULL)";
1320
1321fn with_live_endpoints(where_clause: &str) -> String {
1324 format!("{where_clause} AND {LIVE_ENDPOINTS_CONDITION}")
1325}
1326
1327fn build_edge_filter_sql(
1328 namespace: &str,
1329 filter: &EdgeFilter,
1330) -> (String, Vec<Box<dyn rusqlite::types::ToSql>>) {
1331 build_edge_filter_sql_for_namespaces(&[namespace.to_string()], filter)
1332}
1333
1334fn build_edge_filter_sql_for_namespaces(
1335 namespaces: &[String],
1336 filter: &EdgeFilter,
1337) -> (String, Vec<Box<dyn rusqlite::types::ToSql>>) {
1338 let params: Vec<Box<dyn rusqlite::types::ToSql>> = namespaces
1339 .iter()
1340 .map(|namespace| -> Box<dyn rusqlite::types::ToSql> { Box::new(namespace.clone()) })
1341 .collect();
1342 let namespace_condition = match namespaces.len() {
1343 0 => "0".to_string(),
1344 1 => "namespace = ?1".to_string(),
1345 _ => {
1346 let placeholders: Vec<String> =
1347 (1..=namespaces.len()).map(|i| format!("?{i}")).collect();
1348 format!("namespace IN ({})", placeholders.join(", "))
1349 }
1350 };
1351 build_edge_filter_conditions(namespace_condition, params, filter)
1352}
1353
1354fn build_edge_filter_sql_for_namespaces_json(
1369 namespaces_json: &str,
1370 filter: &EdgeFilter,
1371) -> (String, Vec<Box<dyn rusqlite::types::ToSql>>) {
1372 let params: Vec<Box<dyn rusqlite::types::ToSql>> = vec![Box::new(namespaces_json.to_string())];
1373 let namespace_condition = "namespace IN (SELECT value FROM json_each(?1))".to_string();
1374 build_edge_filter_conditions(namespace_condition, params, filter)
1375}
1376
1377fn build_edge_filter_conditions(
1378 namespace_condition: String,
1379 mut params: Vec<Box<dyn rusqlite::types::ToSql>>,
1380 filter: &EdgeFilter,
1381) -> (String, Vec<Box<dyn rusqlite::types::ToSql>>) {
1382 let mut conditions = vec![namespace_condition, "deleted_at IS NULL".to_string()];
1383
1384 if !filter.ids.is_empty() {
1385 let placeholders: Vec<String> = filter
1386 .ids
1387 .iter()
1388 .map(|id| {
1389 params.push(Box::new(id.to_string()));
1390 format!("?{}", params.len())
1391 })
1392 .collect();
1393 conditions.push(format!("id IN ({})", placeholders.join(",")));
1394 }
1395
1396 if !filter.source_ids.is_empty() {
1397 let placeholders: Vec<String> = filter
1398 .source_ids
1399 .iter()
1400 .map(|id| {
1401 params.push(Box::new(id.to_string()));
1402 format!("?{}", params.len())
1403 })
1404 .collect();
1405 conditions.push(format!("source_id IN ({})", placeholders.join(",")));
1406 }
1407
1408 if !filter.target_ids.is_empty() {
1409 let placeholders: Vec<String> = filter
1410 .target_ids
1411 .iter()
1412 .map(|id| {
1413 params.push(Box::new(id.to_string()));
1414 format!("?{}", params.len())
1415 })
1416 .collect();
1417 conditions.push(format!("target_id IN ({})", placeholders.join(",")));
1418 }
1419
1420 if !filter.relations.is_empty() {
1421 let placeholders: Vec<String> = filter
1422 .relations
1423 .iter()
1424 .map(|r| {
1425 params.push(Box::new(r.to_string()));
1426 format!("?{}", params.len())
1427 })
1428 .collect();
1429 conditions.push(format!("relation IN ({})", placeholders.join(",")));
1430 }
1431
1432 if let Some(min_w) = filter.min_weight {
1433 params.push(Box::new(min_w));
1434 conditions.push(format!("weight >= ?{}", params.len()));
1435 }
1436
1437 if let Some(max_w) = filter.max_weight {
1438 params.push(Box::new(max_w));
1439 conditions.push(format!("weight <= ?{}", params.len()));
1440 }
1441
1442 if let Some(ref time_range) = filter.created_at {
1443 if let Some(start) = time_range.start {
1444 params.push(Box::new(start.timestamp_micros()));
1445 conditions.push(format!("created_at >= ?{}", params.len()));
1446 }
1447 if let Some(end) = time_range.end {
1448 params.push(Box::new(end.timestamp_micros()));
1449 conditions.push(format!("created_at < ?{}", params.len()));
1450 }
1451 }
1452
1453 let clause = format!(" WHERE {}", conditions.join(" AND "));
1454 (clause, params)
1455}
1456
1457fn edge_sort_col(field: &EdgeSortField) -> &'static str {
1458 match field {
1459 EdgeSortField::CreatedAt => "created_at",
1460 EdgeSortField::Weight => "weight",
1461 EdgeSortField::Relation => "relation",
1462 }
1463}
1464
1465fn canonical_edge_endpoints(
1474 relation: EdgeRelation,
1475 source_id: Uuid,
1476 target_id: Uuid,
1477) -> (Uuid, Uuid) {
1478 if relation.is_symmetric() && target_id < source_id {
1479 (target_id, source_id)
1480 } else {
1481 (source_id, target_id)
1482 }
1483}
1484
1485fn edge_endpoints_exist(
1493 conn: &rusqlite::Connection,
1494 source_id: Uuid,
1495 target_id: Uuid,
1496) -> Result<MissingEndpoints, rusqlite::Error> {
1497 let src_exists = endpoint_exists_clause("?1");
1498 let tgt_exists = endpoint_exists_clause("?2");
1499 let sql = format!("SELECT ({src_exists}), ({tgt_exists})");
1500 conn.query_row(
1501 &sql,
1502 rusqlite::params![source_id.to_string(), target_id.to_string()],
1503 |row| {
1504 let src_exists: bool = row.get(0)?;
1505 let tgt_exists: bool = row.get(1)?;
1506 Ok(MissingEndpoints {
1507 source: !src_exists,
1508 target: !tgt_exists,
1509 })
1510 },
1511 )
1512}
1513
1514fn edge_by_natural_key_including_deleted(
1515 conn: &rusqlite::Connection,
1516 edge: &Edge,
1517) -> Result<Option<Edge>, rusqlite::Error> {
1518 let (source_id, target_id) =
1519 canonical_edge_endpoints(edge.relation, edge.source_id, edge.target_id);
1520 conn.query_row(
1521 "SELECT namespace, id, source_id, target_id, relation, weight, \
1522 created_at, updated_at, deleted_at, metadata, target_backend \
1523 FROM graph_edges \
1524 WHERE namespace = ?1 AND source_id = ?2 AND target_id = ?3 AND relation = ?4",
1525 rusqlite::params![
1526 &edge.namespace,
1527 source_id.to_string(),
1528 target_id.to_string(),
1529 edge.relation.as_str(),
1530 ],
1531 read_edge,
1532 )
1533 .optional()
1534}
1535
1536fn observed_edge_upsert(
1539 conn: &rusqlite::Connection,
1540 request: &EdgeUpsertRequest,
1541 guard_endpoints: bool,
1542) -> Result<GuardedEdgeUpsertOutcome, rusqlite::Error> {
1543 let (source_id, target_id) = canonical_edge_endpoints(
1544 request.edge.relation,
1545 request.edge.source_id,
1546 request.edge.target_id,
1547 );
1548 if guard_endpoints {
1549 #[cfg(test)]
1552 tests::insert_probe_seam::hook((source_id, target_id));
1553 let missing = edge_endpoints_exist(conn, source_id, target_id)?;
1554 if missing.any() {
1555 return Ok(GuardedEdgeUpsertOutcome::Refused(
1556 EdgeUpsertRefusal::MissingEndpoints(missing),
1557 ));
1558 }
1559 }
1560
1561 let previous = edge_by_natural_key_including_deleted(conn, &request.edge)?;
1562 if let Some(edge) = previous.as_ref() {
1563 if edge.deleted_at.is_some() && !request.resurrect {
1564 return Ok(GuardedEdgeUpsertOutcome::Refused(
1565 EdgeUpsertRefusal::ResurrectionRequired { edge: edge.clone() },
1566 ));
1567 }
1568 }
1569
1570 let statement = edge_upsert_statement_with_resurrection(&request.edge, request.resurrect);
1571 let mut stmt = conn.prepare(&statement.sql)?;
1572 bind_params(&mut stmt, &statement.params)?;
1573 let affected = stmt.raw_execute()?;
1574 if affected == 0 {
1575 let edge = edge_by_natural_key_including_deleted(conn, &request.edge)?
1576 .ok_or(rusqlite::Error::QueryReturnedNoRows)?;
1577 return Ok(GuardedEdgeUpsertOutcome::Refused(
1578 EdgeUpsertRefusal::ResurrectionRequired { edge },
1579 ));
1580 }
1581
1582 let edge = edge_by_natural_key_including_deleted(conn, &request.edge)?
1583 .ok_or(rusqlite::Error::QueryReturnedNoRows)?;
1584 let disposition = match previous.as_ref().and_then(|edge| edge.deleted_at) {
1585 None if previous.is_none() => EdgeUpsertDisposition::Created,
1586 None => EdgeUpsertDisposition::Updated,
1587 Some(_) => EdgeUpsertDisposition::Resurrected,
1588 };
1589 Ok(GuardedEdgeUpsertOutcome::Written(EdgeUpsertResult {
1590 edge,
1591 disposition,
1592 previous,
1593 }))
1594}
1595
1596fn observed_edge_batch_upsert(
1599 conn: &rusqlite::Connection,
1600 requests: &[EdgeUpsertRequest],
1601 guard_endpoints: bool,
1602) -> Result<GuardedEdgeBatchUpsertOutcome, rusqlite::Error> {
1603 for (entry_index, request) in requests.iter().enumerate() {
1607 let (source_id, target_id) = canonical_edge_endpoints(
1608 request.edge.relation,
1609 request.edge.source_id,
1610 request.edge.target_id,
1611 );
1612 if guard_endpoints {
1613 let missing = edge_endpoints_exist(conn, source_id, target_id)?;
1614 if missing.any() {
1615 return Ok(GuardedEdgeBatchUpsertOutcome {
1616 rows: Vec::new(),
1617 refusal: Some(GuardedEdgeBatchRefusal {
1618 entry_index,
1619 reason: EdgeUpsertRefusal::MissingEndpoints(missing),
1620 }),
1621 });
1622 }
1623 }
1624 if let Some(edge) = edge_by_natural_key_including_deleted(conn, &request.edge)? {
1625 if edge.deleted_at.is_some() && !request.resurrect {
1626 return Ok(GuardedEdgeBatchUpsertOutcome {
1627 rows: Vec::new(),
1628 refusal: Some(GuardedEdgeBatchRefusal {
1629 entry_index,
1630 reason: EdgeUpsertRefusal::ResurrectionRequired { edge },
1631 }),
1632 });
1633 }
1634 }
1635 }
1636
1637 let mut rows = Vec::with_capacity(requests.len());
1638 for request in requests {
1639 match observed_edge_upsert(conn, request, false)? {
1640 GuardedEdgeUpsertOutcome::Written(row) => rows.push(row),
1641 GuardedEdgeUpsertOutcome::Refused(_) => {
1642 return Err(rusqlite::Error::ExecuteReturnedResults)
1643 }
1644 }
1645 }
1646 Ok(GuardedEdgeBatchUpsertOutcome {
1647 rows,
1648 refusal: None,
1649 })
1650}
1651
1652fn traversal_neighbor_sql(
1657 direction: Direction,
1658 relation_count: usize,
1659 has_min_weight: bool,
1660) -> String {
1661 let (node_column, endpoint_column, index) = match direction {
1662 Direction::Out => ("target_id", "source_id", "idx_graph_edges_ns_src_rel"),
1663 Direction::In => ("source_id", "target_id", "idx_graph_edges_ns_tgt_rel"),
1664 Direction::Both => unreachable!("Direction::Both is split into indexed Out/In seeks"),
1665 };
1666 let mut sql = format!(
1667 "SELECT {node_column}, id, weight \
1668 FROM graph_edges INDEXED BY {index} \
1669 WHERE namespace = ?1 AND {endpoint_column} = ?2 AND deleted_at IS NULL"
1670 );
1671 if relation_count > 0 {
1672 let placeholders = (0..relation_count)
1673 .map(|offset| format!("?{}", 4 + offset))
1674 .collect::<Vec<_>>()
1675 .join(",");
1676 sql.push_str(&format!(" AND relation IN ({placeholders})"));
1677 }
1678 if has_min_weight {
1679 sql.push_str(&format!(" AND weight >= ?{}", 4 + relation_count));
1680 }
1681 sql.push_str(" LIMIT ?3");
1682 sql
1683}
1684
1685fn traversal_timeout_error(budget: &TraversalExecutionBudget) -> StorageError {
1686 StorageError::Timeout {
1687 operation: format!(
1688 "traverse ({}ms execution budget)",
1689 budget.max_duration().as_millis()
1690 )
1691 .into(),
1692 }
1693}
1694
1695fn traversal_work_error(budget: &TraversalExecutionBudget) -> StorageError {
1696 StorageError::InvalidInput {
1697 capability: StorageCapability::Graph,
1698 operation: "traverse".into(),
1699 message: format!(
1700 "traversal work budget exceeded after {} adjacency rows; \
1701 narrow roots, depth, relations, or result limit",
1702 budget.work_limit()
1703 ),
1704 }
1705}
1706
1707#[derive(Clone, Copy)]
1708struct TraversalFrontierNode {
1709 node_id: Uuid,
1710 depth: usize,
1711 total_weight: f64,
1712}
1713
1714#[allow(clippy::too_many_arguments)]
1715fn run_bounded_traversal(
1716 conn: &rusqlite::Connection,
1717 roots: Vec<Uuid>,
1718 opts: TraversalOptions,
1719 include_roots: bool,
1720 namespace: String,
1721 origin: khive_storage::tx_registry::TxOrigin,
1722 budget: TraversalExecutionBudget,
1723 counted_rows: &std::sync::atomic::AtomicU64,
1724 counted_queries: &std::sync::atomic::AtomicU64,
1725) -> Result<Vec<GraphPath>, StorageError> {
1726 let progress_timed_out = Arc::new(std::sync::atomic::AtomicBool::new(false));
1727 let callback_timed_out = Arc::clone(&progress_timed_out);
1728 let callback_budget = budget.clone();
1729 #[cfg(test)]
1730 let progress_seam_root = roots.first().copied();
1731 conn.progress_handler(
1732 1_000,
1733 Some(move || {
1734 if crate::read_cancellation::current_read_should_interrupt() {
1735 return true;
1736 }
1737 #[cfg(test)]
1738 if tests::traverse_progress_seam::hook(progress_seam_root) {
1739 callback_timed_out.store(true, std::sync::atomic::Ordering::Relaxed);
1740 return true;
1741 }
1742 let expired = callback_budget.is_expired();
1743 if expired {
1744 callback_timed_out.store(true, std::sync::atomic::Ordering::Relaxed);
1745 }
1746 expired
1747 }),
1748 )
1749 .map_err(|e| map_err(e, "traverse_progress_handler"))?;
1750
1751 let result = (|| {
1752 let result_limit = opts.effective_limit() as usize;
1753 let relation_count = opts.relations.as_ref().map_or(0, Vec::len);
1754 let directions = match opts.direction {
1755 Direction::Out => vec![Direction::Out],
1756 Direction::In => vec![Direction::In],
1757 Direction::Both => vec![Direction::Out, Direction::In],
1758 };
1759 let statements = directions
1760 .into_iter()
1761 .map(|direction| {
1762 traversal_neighbor_sql(direction, relation_count, opts.min_weight.is_some())
1763 })
1764 .collect::<Vec<_>>();
1765 let map_sql_error = |error| {
1766 if progress_timed_out.load(std::sync::atomic::Ordering::Relaxed) {
1767 traversal_timeout_error(&budget)
1768 } else {
1769 map_err(error, "traverse")
1770 }
1771 };
1772
1773 let mut all_paths = Vec::with_capacity(roots.len());
1774 for root_id in roots {
1775 let mut seen = HashSet::new();
1776 seen.insert(root_id);
1777 let mut frontier = VecDeque::from([TraversalFrontierNode {
1778 node_id: root_id,
1779 depth: 0,
1780 total_weight: 0.0,
1781 }]);
1782 let mut nodes = Vec::with_capacity(result_limit + usize::from(include_roots));
1783 if include_roots {
1784 nodes.push(PathNode {
1785 node_id: root_id,
1786 via_edge: None,
1787 depth: 0,
1788 name: None,
1789 kind: None,
1790 properties: None,
1791 weight: 0.0,
1792 });
1793 }
1794 let mut non_root_count = 0usize;
1795
1796 'root_walk: while non_root_count < result_limit {
1797 let Some(current) = frontier.pop_front() else {
1798 break;
1799 };
1800 if current.depth >= opts.max_depth {
1801 continue;
1802 }
1803 if budget.is_expired() {
1804 return Err(traversal_timeout_error(&budget));
1805 }
1806
1807 for sql in &statements {
1808 let row_cap = budget.remaining_work().saturating_add(1);
1809 let mut params: Vec<Box<dyn rusqlite::types::ToSql>> = vec![
1810 Box::new(namespace.clone()),
1811 Box::new(current.node_id.to_string()),
1812 Box::new(row_cap as i64),
1813 ];
1814 if let Some(relations) = &opts.relations {
1815 params.extend(relations.iter().map(|relation| {
1816 Box::new(relation.to_string()) as Box<dyn rusqlite::types::ToSql>
1817 }));
1818 }
1819 if let Some(min_weight) = opts.min_weight {
1820 params.push(Box::new(min_weight));
1821 }
1822 let param_refs = params
1823 .iter()
1824 .map(|param| param.as_ref())
1825 .collect::<Vec<&dyn rusqlite::types::ToSql>>();
1826
1827 let _snapshot = khive_storage::tx_registry::register_scoped(
1828 Some("graph_traverse_read".to_string()),
1829 origin.clone(),
1830 );
1831 let mut stmt = conn.prepare(sql).map_err(&map_sql_error)?;
1832 counted_queries.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1833 let mut rows = stmt.query(param_refs.as_slice()).map_err(&map_sql_error)?;
1834 while let Some(row) = rows.next().map_err(&map_sql_error)? {
1835 #[cfg(test)]
1838 tests::traverse_snapshot_seam::hook(current.node_id);
1839 counted_rows.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1840 if budget.is_expired() {
1841 return Err(traversal_timeout_error(&budget));
1842 }
1843 if !budget.try_consume_row() {
1844 return Err(traversal_work_error(&budget));
1845 }
1846 let node_str: String = row.get(0).map_err(&map_sql_error)?;
1847 let edge_str: String = row.get(1).map_err(&map_sql_error)?;
1848 let edge_weight: f64 = row.get(2).map_err(&map_sql_error)?;
1849 let node_id = parse_uuid(&node_str).map_err(&map_sql_error)?;
1850 if !seen.insert(node_id) {
1851 continue;
1852 }
1853 let via_edge = parse_uuid(&edge_str).map_err(&map_sql_error)?;
1854 let depth = current.depth + 1;
1855 let total_weight = current.total_weight + edge_weight;
1856 nodes.push(PathNode {
1857 node_id,
1858 via_edge: Some(via_edge),
1859 depth,
1860 name: None,
1861 kind: None,
1862 properties: None,
1863 weight: total_weight,
1864 });
1865 non_root_count += 1;
1866 if depth < opts.max_depth {
1867 frontier.push_back(TraversalFrontierNode {
1868 node_id,
1869 depth,
1870 total_weight,
1871 });
1872 }
1873 if non_root_count == result_limit {
1874 break 'root_walk;
1875 }
1876 }
1877 }
1878 }
1879
1880 if !nodes.is_empty() {
1881 let total_weight = nodes.iter().map(|node| node.weight).fold(0.0_f64, f64::max);
1882 all_paths.push(GraphPath {
1883 root_id,
1884 nodes,
1885 total_weight,
1886 });
1887 }
1888 }
1889 Ok(all_paths)
1890 })();
1891
1892 conn.progress_handler(0, None::<fn() -> bool>)
1893 .map_err(|e| map_err(e, "traverse_progress_handler_clear"))?;
1894 result
1895}
1896
1897impl SqlGraphStore {
1898 async fn query_neighbors_page(
1899 &self,
1900 operation: &'static str,
1901 node_id: Uuid,
1902 query: NeighborQuery,
1903 after: Option<NeighborCursor>,
1904 neighbor_kinds: Option<Vec<String>>,
1905 ) -> Result<Vec<NeighborHit>, StorageError> {
1906 count_neighbor_select();
1907
1908 let namespace = self.namespace.clone();
1909 let node_str = node_id.to_string();
1910 let counted_queries = Arc::new(std::sync::atomic::AtomicU64::new(0));
1911 let counted_rows = Arc::new(std::sync::atomic::AtomicU64::new(0));
1912 let closure_queries = Arc::clone(&counted_queries);
1913 let closure_rows = Arc::clone(&counted_rows);
1914 let result = self
1915 .with_reader(operation, move |conn| {
1916 let base_out = "SELECT target_id AS node_id, id AS edge_id, relation, weight \
1917 FROM graph_edges \
1918 WHERE namespace = ?1 AND source_id = ?2 AND deleted_at IS NULL";
1919 let base_in = "SELECT source_id AS node_id, id AS edge_id, relation, weight \
1920 FROM graph_edges \
1921 WHERE namespace = ?1 AND target_id = ?2 AND deleted_at IS NULL";
1922 let sql = match query.direction {
1923 Direction::Out => base_out.to_string(),
1924 Direction::In => base_in.to_string(),
1925 Direction::Both => format!("{} UNION ALL {}", base_out, base_in),
1926 };
1927 let (where_extra, limit_clause, extra_params) =
1928 neighbor_extra_clause(&query, 3, after.as_ref(), neighbor_kinds.as_deref());
1929 let full_sql = format!(
1930 "SELECT node_id, edge_id, relation, weight FROM ({}){} \
1931 ORDER BY weight DESC, node_id ASC, edge_id ASC{}",
1932 sql, where_extra, limit_clause
1933 );
1934 let mut stmt = conn.prepare(&full_sql)?;
1935 let mut all_params: Vec<Box<dyn rusqlite::types::ToSql>> = Vec::new();
1936 all_params.push(Box::new(namespace.clone()));
1937 all_params.push(Box::new(node_str.clone()));
1938 all_params.extend(extra_params);
1939 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
1940 all_params.iter().map(|p| p.as_ref()).collect();
1941
1942 closure_queries.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1943 let rows = stmt.query_map(param_refs.as_slice(), |row| {
1944 let nid_str: String = row.get(0)?;
1945 let eid_str: String = row.get(1)?;
1946 let relation_str: String = row.get(2)?;
1947 let weight: f64 = row.get(3)?;
1948 Ok((nid_str, eid_str, relation_str, weight))
1949 })?;
1950 let mut hits = Vec::new();
1951 for row in rows {
1952 let (nid_str, eid_str, relation_str, weight) = row?;
1953 closure_rows.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1954 let relation = relation_str.parse::<EdgeRelation>().map_err(|e| {
1955 rusqlite::Error::FromSqlConversionFailure(
1956 2,
1957 rusqlite::types::Type::Text,
1958 Box::new(e),
1959 )
1960 })?;
1961 hits.push(NeighborHit {
1962 node_id: parse_uuid(&nid_str)?,
1963 edge_id: parse_uuid(&eid_str)?,
1964 relation,
1965 weight,
1966 name: None,
1967 kind: None,
1968 entity_type: None,
1969 });
1970 }
1971 Ok(hits)
1972 })
1973 .await;
1974
1975 report_graph_usage(&counted_queries, &counted_rows);
1976 result
1977 }
1978}
1979
1980#[async_trait]
1981impl GraphStore for SqlGraphStore {
1982 async fn latest_annotating_note(
1983 &self,
1984 node_id: Uuid,
1985 kind: &str,
1986 tag: &str,
1987 ) -> Result<Option<(Uuid, i64)>, StorageError> {
1988 let namespace = self.namespace.clone();
1989 let node_id = node_id.to_string();
1990 let kind = kind.to_owned();
1991 let tag = tag.to_owned();
1992 self.with_reader("latest_annotating_note", move |conn| {
1993 conn.query_row(
1994 LATEST_ANNOTATING_NOTE_SQL,
1995 rusqlite::params![namespace, node_id, kind, tag],
1996 |row| {
1997 let id: String = row.get(0)?;
1998 Ok((parse_uuid(&id)?, row.get(1)?))
1999 },
2000 )
2001 .optional()
2002 })
2003 .await
2004 }
2005
2006 async fn upsert_edge(&self, edge: Edge) -> Result<(), StorageError> {
2007 self.upsert_edge_observed(EdgeUpsertRequest {
2008 edge,
2009 resurrect: false,
2010 })
2011 .await
2012 .map(|_| ())
2013 }
2014
2015 async fn upsert_edge_observed(
2016 &self,
2017 request: EdgeUpsertRequest,
2018 ) -> Result<EdgeUpsertResult, StorageError> {
2019 match self
2020 .observed_edge_write("upsert_edge_observed", request, false)
2021 .await?
2022 {
2023 GuardedEdgeUpsertOutcome::Written(result) => Ok(result),
2024 GuardedEdgeUpsertOutcome::Refused(EdgeUpsertRefusal::ResurrectionRequired { edge }) => {
2025 Err(resurrection_required_error("upsert_edge_observed", &edge))
2026 }
2027 GuardedEdgeUpsertOutcome::Refused(EdgeUpsertRefusal::MissingEndpoints(_)) => {
2028 Err(StorageError::Conflict {
2029 capability: StorageCapability::Graph,
2030 operation: "upsert_edge_observed".into(),
2031 message: "unguarded edge upsert reported a missing-endpoint refusal".into(),
2032 })
2033 }
2034 }
2035 }
2036
2037 async fn insert_edge_if_absent(&self, edge: Edge) -> Result<bool, StorageError> {
2038 let statement = edge_insert_if_absent_statement(&edge);
2039 self.with_writer("insert_edge_if_absent", move |conn| {
2040 let mut stmt = conn.prepare(&statement.sql)?;
2041 bind_params(&mut stmt, &statement.params)?;
2042 Ok(stmt.raw_execute()? > 0)
2043 })
2044 .await
2045 }
2046
2047 async fn replace_edge_if_unchanged(
2048 &self,
2049 edge: Edge,
2050 expected_updated_at: DateTime<Utc>,
2051 expected_deleted_at: Option<DateTime<Utc>>,
2052 ) -> Result<bool, StorageError> {
2053 let statement =
2054 edge_replace_if_unchanged_statement(&edge, expected_updated_at, expected_deleted_at);
2055 self.with_writer("replace_edge_if_unchanged", move |conn| {
2056 let mut stmt = conn.prepare(&statement.sql)?;
2057 bind_params(&mut stmt, &statement.params)?;
2058 Ok(stmt.raw_execute()? > 0)
2059 })
2060 .await
2061 }
2062
2063 async fn upsert_edges(&self, edges: Vec<Edge>) -> Result<BatchWriteSummary, StorageError> {
2064 let attempted = edges.len() as u64;
2065 let requests = edges
2066 .into_iter()
2067 .map(|edge| EdgeUpsertRequest {
2068 edge,
2069 resurrect: false,
2070 })
2071 .collect();
2072 let outcome = self
2073 .observed_edge_batch_write("upsert_edges", requests, false)
2074 .await?;
2075 if let Some(refusal) = outcome.refusal {
2076 return match refusal.reason {
2077 EdgeUpsertRefusal::ResurrectionRequired { edge } => {
2078 Err(resurrection_required_error("upsert_edges", &edge))
2079 }
2080 EdgeUpsertRefusal::MissingEndpoints(_) => Err(StorageError::Conflict {
2081 capability: StorageCapability::Graph,
2082 operation: "upsert_edges".into(),
2083 message: "unguarded edge batch reported a missing-endpoint refusal".into(),
2084 }),
2085 };
2086 }
2087 Ok(BatchWriteSummary {
2088 attempted,
2089 affected: outcome.rows.len() as u64,
2090 ..BatchWriteSummary::default()
2091 })
2092 }
2093
2094 async fn upsert_edge_guarded(&self, edge: Edge) -> Result<GuardedWriteOutcome, StorageError> {
2095 match self
2096 .upsert_edge_guarded_observed(EdgeUpsertRequest {
2097 edge,
2098 resurrect: false,
2099 })
2100 .await?
2101 {
2102 GuardedEdgeUpsertOutcome::Written(_) => Ok(GuardedWriteOutcome::Written),
2103 GuardedEdgeUpsertOutcome::Refused(EdgeUpsertRefusal::MissingEndpoints(missing)) => {
2104 Ok(GuardedWriteOutcome::Refused(missing))
2105 }
2106 GuardedEdgeUpsertOutcome::Refused(EdgeUpsertRefusal::ResurrectionRequired { edge }) => {
2107 Err(resurrection_required_error("upsert_edge_guarded", &edge))
2108 }
2109 }
2110 }
2111
2112 async fn upsert_edge_guarded_observed(
2113 &self,
2114 request: EdgeUpsertRequest,
2115 ) -> Result<GuardedEdgeUpsertOutcome, StorageError> {
2116 self.observed_edge_write("upsert_edge_guarded_observed", request, true)
2117 .await
2118 }
2119
2120 async fn upsert_edges_guarded(
2121 &self,
2122 edges: Vec<Edge>,
2123 ) -> Result<GuardedBatchOutcome, StorageError> {
2124 let attempted = edges.len() as u64;
2125 let requests = edges
2126 .iter()
2127 .cloned()
2128 .map(|edge| EdgeUpsertRequest {
2129 edge,
2130 resurrect: false,
2131 })
2132 .collect();
2133 let outcome = self.upsert_edges_guarded_observed(requests).await?;
2134 match outcome.refusal {
2135 None => Ok(GuardedBatchOutcome {
2136 summary: BatchWriteSummary {
2137 attempted,
2138 affected: outcome.rows.len() as u64,
2139 ..BatchWriteSummary::default()
2140 },
2141 refused: None,
2142 }),
2143 Some(refusal) => match refusal.reason {
2144 EdgeUpsertRefusal::MissingEndpoints(missing) => {
2145 let index = refusal.entry_index;
2146 let edge = &edges[index];
2147 let (source_id, target_id) =
2148 canonical_edge_endpoints(edge.relation, edge.source_id, edge.target_id);
2149 let message = format!(
2150 "batch entry {index}: edge endpoint no longer exists at write time: source \
2151 {source_id} or target {target_id}"
2152 );
2153 let mut summary = BatchWriteSummary {
2154 attempted,
2155 ..BatchWriteSummary::default()
2156 };
2157 summary.first_error = message.clone();
2160 let refusal = GuardedBatchRefusal {
2161 entry_index: index,
2162 missing,
2163 };
2164 for (failed_index, failed_edge) in edges.iter().enumerate() {
2165 refusal.record_failure(&mut summary, failed_index, failed_edge, &message);
2166 }
2167 Ok(GuardedBatchOutcome {
2168 summary,
2169 refused: Some(refusal),
2170 })
2171 }
2172 EdgeUpsertRefusal::ResurrectionRequired { edge } => {
2173 Err(resurrection_required_error("upsert_edges_guarded", &edge))
2174 }
2175 },
2176 }
2177 }
2178
2179 async fn upsert_edges_guarded_observed(
2180 &self,
2181 requests: Vec<EdgeUpsertRequest>,
2182 ) -> Result<GuardedEdgeBatchUpsertOutcome, StorageError> {
2183 self.observed_edge_batch_write("upsert_edges_guarded_observed", requests, true)
2184 .await
2185 }
2186
2187 async fn get_edge(&self, id: LinkId) -> Result<Option<Edge>, StorageError> {
2188 let id_str = Uuid::from(id).to_string();
2189
2190 self.with_reader("get_edge", move |conn| {
2191 let mut stmt = conn.prepare(
2192 "SELECT namespace, id, source_id, target_id, relation, weight, \
2193 created_at, updated_at, deleted_at, metadata, target_backend \
2194 FROM graph_edges WHERE id = ?1 AND deleted_at IS NULL",
2195 )?;
2196 let mut rows = stmt.query(rusqlite::params![id_str])?;
2197 match rows.next()? {
2198 Some(row) => Ok(Some(read_edge(row)?)),
2199 None => Ok(None),
2200 }
2201 })
2202 .await
2203 }
2204
2205 async fn get_edge_including_deleted(&self, id: LinkId) -> Result<Option<Edge>, StorageError> {
2206 let id_str = Uuid::from(id).to_string();
2207
2208 self.with_reader("get_edge_including_deleted", move |conn| {
2209 let mut stmt = conn.prepare(
2210 "SELECT namespace, id, source_id, target_id, relation, weight, \
2211 created_at, updated_at, deleted_at, metadata, target_backend \
2212 FROM graph_edges WHERE id = ?1",
2213 )?;
2214 let mut rows = stmt.query(rusqlite::params![id_str])?;
2215 match rows.next()? {
2216 Some(row) => Ok(Some(read_edge(row)?)),
2217 None => Ok(None),
2218 }
2219 })
2220 .await
2221 }
2222
2223 async fn edge_sequence(&self, id: Uuid) -> Result<Option<i64>, StorageError> {
2224 let id = id.to_string();
2225 self.with_reader("edge_sequence", move |conn| {
2226 conn.query_row(
2227 "SELECT seq FROM graph_edges_seq WHERE edge_id = ?1",
2228 rusqlite::params![id],
2229 |row| row.get(0),
2230 )
2231 .optional()
2232 })
2233 .await
2234 }
2235
2236 async fn edge_sequences(&self, ids: &[Uuid]) -> Result<Vec<(Uuid, i64)>, StorageError> {
2237 if ids.is_empty() {
2238 return Ok(Vec::new());
2239 }
2240 let ids = ids.to_vec();
2241 self.with_reader("edge_sequences", move |conn| {
2242 const CHUNK: usize = 900;
2243 let mut resolved = Vec::with_capacity(ids.len());
2244 for chunk in ids.chunks(CHUNK) {
2245 let placeholders = (1..=chunk.len())
2246 .map(|index| format!("?{index}"))
2247 .collect::<Vec<_>>()
2248 .join(", ");
2249 let sql = format!(
2250 "SELECT edge_id, seq FROM graph_edges_seq WHERE edge_id IN ({placeholders})"
2251 );
2252 let strings = chunk.iter().map(Uuid::to_string).collect::<Vec<_>>();
2253 let params = strings
2254 .iter()
2255 .map(|id| id as &dyn rusqlite::types::ToSql)
2256 .collect::<Vec<_>>();
2257 let mut stmt = conn.prepare(&sql)?;
2258 let rows = stmt.query_map(params.as_slice(), |row| {
2259 let id: String = row.get(0)?;
2260 Ok((parse_uuid(&id)?, row.get::<_, i64>(1)?))
2261 })?;
2262 resolved.extend(rows.collect::<Result<Vec<_>, _>>()?);
2263 }
2264 Ok(resolved)
2265 })
2266 .await
2267 }
2268
2269 async fn get_edge_by_natural_key_including_deleted(
2270 &self,
2271 namespace: &str,
2272 source_id: Uuid,
2273 target_id: Uuid,
2274 relation: EdgeRelation,
2275 ) -> Result<Option<Edge>, StorageError> {
2276 let namespace = namespace.to_string();
2277 let source_str = source_id.to_string();
2278 let target_str = target_id.to_string();
2279 let relation_str = relation.to_string();
2280
2281 self.with_reader("get_edge_by_natural_key_including_deleted", move |conn| {
2282 let mut stmt = conn.prepare(
2283 "SELECT namespace, id, source_id, target_id, relation, weight, \
2284 created_at, updated_at, deleted_at, metadata, target_backend \
2285 FROM graph_edges \
2286 WHERE namespace = ?1 AND source_id = ?2 AND target_id = ?3 AND relation = ?4",
2287 )?;
2288 let mut rows = stmt.query(rusqlite::params![
2289 namespace,
2290 source_str,
2291 target_str,
2292 relation_str
2293 ])?;
2294 match rows.next()? {
2295 Some(row) => Ok(Some(read_edge(row)?)),
2296 None => Ok(None),
2297 }
2298 })
2299 .await
2300 }
2301
2302 async fn get_edges(&self, ids: &[LinkId]) -> Result<Vec<Edge>, StorageError> {
2303 if ids.is_empty() {
2304 return Ok(Vec::new());
2305 }
2306 const CHUNK: usize = 900;
2308 let id_strs: Vec<String> = ids.iter().map(|id| Uuid::from(*id).to_string()).collect();
2309
2310 let mut result: Vec<Edge> = Vec::with_capacity(ids.len());
2311 for chunk in id_strs.chunks(CHUNK) {
2312 let chunk_owned: Vec<String> = chunk.to_vec();
2313 let edges = self
2314 .with_reader("get_edges", move |conn| {
2315 let placeholders: Vec<String> =
2316 (1..=chunk_owned.len()).map(|i| format!("?{}", i)).collect();
2317 let sql = format!(
2318 "SELECT namespace, id, source_id, target_id, relation, weight, \
2319 created_at, updated_at, deleted_at, metadata, target_backend \
2320 FROM graph_edges WHERE id IN ({}) AND deleted_at IS NULL",
2321 placeholders.join(",")
2322 );
2323 let mut stmt = conn.prepare(&sql)?;
2324 let params: Vec<&dyn rusqlite::types::ToSql> = chunk_owned
2325 .iter()
2326 .map(|s| s as &dyn rusqlite::types::ToSql)
2327 .collect();
2328 let rows = stmt.query_map(params.as_slice(), read_edge)?;
2329 let mut edges = Vec::new();
2330 for row in rows {
2331 edges.push(row?);
2332 }
2333 Ok(edges)
2334 })
2335 .await?;
2336 result.extend(edges);
2337 }
2338 Ok(result)
2339 }
2340
2341 async fn batch_neighbors(
2342 &self,
2343 sources: &[Uuid],
2344 query: NeighborQuery,
2345 ) -> Result<Vec<(Uuid, NeighborHit)>, StorageError> {
2346 use khive_storage::types::Direction;
2347
2348 if sources.is_empty() {
2349 return Ok(Vec::new());
2350 }
2351 let mut seen_sources = HashSet::with_capacity(sources.len());
2352 let unique_sources: Vec<Uuid> = sources
2353 .iter()
2354 .copied()
2355 .filter(|source| seen_sources.insert(*source))
2356 .collect();
2357 const CHUNK_SIZE: usize = 880;
2358
2359 let namespace = self.namespace.clone();
2360 let mut result: Vec<(Uuid, NeighborHit)> = Vec::new();
2361 let counted_queries = Arc::new(std::sync::atomic::AtomicU64::new(0));
2364 let counted_rows = Arc::new(std::sync::atomic::AtomicU64::new(0));
2365
2366 for chunk in unique_sources.chunks(CHUNK_SIZE) {
2367 let chunk_owned: Vec<Uuid> = chunk.to_vec();
2368 let query_clone = query.clone();
2369 let ns = namespace.clone();
2370 let closure_queries = Arc::clone(&counted_queries);
2371 let closure_rows = Arc::clone(&counted_rows);
2372
2373 let chunk_result = self
2374 .with_reader("batch_neighbors", move |conn| {
2375 let src_strs: Vec<String> = chunk_owned.iter().map(|u| u.to_string()).collect();
2376
2377 let sources_json = serde_json::to_string(&src_strs).map_err(|error| {
2378 rusqlite::Error::ToSqlConversionFailure(Box::new(error))
2379 })?;
2380
2381 let build_inner_sql =
2382 |direction_out: bool,
2383 q: &NeighborQuery|
2384 -> (String, Vec<String>, Option<f64>) {
2385 let (filter_col, node_col) = if direction_out {
2386 ("source_id", "target_id")
2387 } else {
2388 ("target_id", "source_id")
2389 };
2390
2391 let mut rel_params: Vec<String> = Vec::new();
2392 let mut conditions: Vec<String> = Vec::new();
2393 let mut param_idx = 3;
2394
2395 if let Some(ref rels) = q.relations {
2396 if !rels.is_empty() {
2397 let ps: Vec<String> = rels
2398 .iter()
2399 .map(|r| {
2400 rel_params.push(r.to_string());
2401 let p = format!("?{param_idx}");
2402 param_idx += 1;
2403 p
2404 })
2405 .collect();
2406 conditions
2407 .push(format!("edges.relation IN ({})", ps.join(",")));
2408 }
2409 }
2410
2411 let min_weight_val = if let Some(min_w) = q.min_weight {
2414 conditions.push(format!("edges.weight >= ?{param_idx}"));
2415 Some(min_w)
2416 } else {
2417 None
2418 };
2419
2420 let where_extra = if conditions.is_empty() {
2421 String::new()
2422 } else {
2423 format!(" AND {}", conditions.join(" AND "))
2424 };
2425
2426 let sql = format!(
2427 "SELECT requested.origin_id, edges.{node_col} AS node_id, \
2428 edges.id AS edge_id, edges.relation, edges.weight \
2429 FROM requested CROSS JOIN graph_edges AS edges \
2430 ON edges.{filter_col} = requested.origin_id \
2431 WHERE edges.namespace = ?1 \
2432 AND edges.deleted_at IS NULL{where_extra}",
2433 );
2434 (sql, rel_params, min_weight_val)
2435 };
2436
2437 let mut all_params: Vec<Box<dyn rusqlite::types::ToSql>> = Vec::new();
2438 all_params.push(Box::new(ns.to_string()));
2439 all_params.push(Box::new(sources_json));
2440
2441 let (combined_inner, rel_params, min_weight_val) = match query_clone.direction {
2442 Direction::Out => build_inner_sql(true, &query_clone),
2443 Direction::In => build_inner_sql(false, &query_clone),
2444 Direction::Both => {
2445 let (out_sql, rel_params, min_weight_val) =
2446 build_inner_sql(true, &query_clone);
2447 let (in_sql, _, _) = build_inner_sql(false, &query_clone);
2448 (
2449 format!("{out_sql} UNION ALL {in_sql}"),
2450 rel_params,
2451 min_weight_val,
2452 )
2453 }
2454 };
2455
2456 for relation in rel_params {
2457 all_params.push(Box::new(relation));
2458 }
2459 if let Some(min_weight) = min_weight_val {
2460 all_params.push(Box::new(min_weight));
2461 }
2462 let limit_param_idx = all_params.len() + 1;
2463
2464 let full_sql = if let Some(lim) = query_clone.limit {
2472 all_params.push(Box::new(lim as i64));
2473 format!(
2474 "WITH requested(origin_id) AS (\
2475 SELECT value FROM json_each(?2)\
2476 ) SELECT origin_id, node_id, edge_id, relation, weight \
2477 FROM (SELECT *, ROW_NUMBER() OVER (PARTITION BY origin_id \
2478 ORDER BY weight DESC, node_id ASC) AS rn \
2479 FROM ({combined_inner})) WHERE rn <= ?{limit_param_idx}",
2480 )
2481 } else {
2482 format!(
2483 "WITH requested(origin_id) AS (\
2484 SELECT value FROM json_each(?2)\
2485 ) SELECT origin_id, node_id, edge_id, relation, weight \
2486 FROM ({combined_inner})",
2487 )
2488 };
2489
2490 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
2491 all_params.iter().map(|p| p.as_ref()).collect();
2492
2493 let mut stmt = conn.prepare(&full_sql)?;
2494 closure_queries.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2495 let rows = stmt.query_map(param_refs.as_slice(), |row| {
2496 let origin_str: String = row.get(0)?;
2497 let nid_str: String = row.get(1)?;
2498 let eid_str: String = row.get(2)?;
2499 let relation_str: String = row.get(3)?;
2500 let weight: f64 = row.get(4)?;
2501 Ok((origin_str, nid_str, eid_str, relation_str, weight))
2502 })?;
2503
2504 let mut pairs = Vec::new();
2505 for row in rows {
2506 let (origin_str, nid_str, eid_str, relation_str, weight) = row?;
2507 closure_rows.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2508 let origin = parse_uuid(&origin_str)?;
2509 let node_id = parse_uuid(&nid_str)?;
2510 let edge_id = parse_uuid(&eid_str)?;
2511 let relation = relation_str.parse::<EdgeRelation>().map_err(|e| {
2512 rusqlite::Error::FromSqlConversionFailure(
2513 3,
2514 rusqlite::types::Type::Text,
2515 Box::new(e),
2516 )
2517 })?;
2518 pairs.push((
2519 origin,
2520 NeighborHit {
2521 node_id,
2522 edge_id,
2523 relation,
2524 weight,
2525 name: None,
2526 kind: None,
2527 entity_type: None,
2528 },
2529 ));
2530 }
2531 Ok(pairs)
2532 })
2533 .await;
2534 let pairs = match chunk_result {
2535 Ok(pairs) => pairs,
2536 Err(e) => {
2537 report_graph_usage(&counted_queries, &counted_rows);
2538 return Err(e);
2539 }
2540 };
2541 result.extend(pairs);
2542 }
2543 report_graph_usage(&counted_queries, &counted_rows);
2544
2545 let requested: HashSet<Uuid> = unique_sources.iter().copied().collect();
2546 let mut grouped: HashMap<Uuid, Vec<NeighborHit>> =
2547 HashMap::with_capacity(unique_sources.len());
2548 for (origin, hit) in result {
2549 if !requested.contains(&origin) {
2550 return Err(StorageError::Internal(format!(
2551 "batch_neighbors returned unrequested origin {origin}"
2552 )));
2553 }
2554 grouped.entry(origin).or_default().push(hit);
2555 }
2556
2557 for hits in grouped.values_mut() {
2558 hits.sort_by(|a, b| {
2559 b.weight
2560 .partial_cmp(&a.weight)
2561 .unwrap_or(std::cmp::Ordering::Equal)
2562 .then(a.node_id.cmp(&b.node_id))
2563 .then(a.edge_id.cmp(&b.edge_id))
2564 });
2565 }
2566
2567 let mut ordered = Vec::new();
2568 for &source in sources {
2569 if let Some(hits) = grouped.get(&source) {
2570 ordered.extend(hits.iter().cloned().map(|hit| (source, hit)));
2571 }
2572 }
2573 Ok(ordered)
2574 }
2575
2576 async fn delete_edge(&self, id: LinkId, mode: DeleteMode) -> Result<bool, StorageError> {
2577 let id = Uuid::from(id);
2578 let statement = match mode {
2579 DeleteMode::Soft => {
2580 edge_soft_delete_statement(id, chrono::Utc::now().timestamp_micros())
2581 }
2582 DeleteMode::Hard => edge_hard_delete_statement(id),
2583 };
2584 self.with_writer("delete_edge", move |conn| {
2585 let mut stmt = conn.prepare(&statement.sql)?;
2586 bind_params(&mut stmt, &statement.params)?;
2587 Ok(stmt.raw_execute()? > 0)
2588 })
2589 .await
2590 }
2591
2592 async fn query_edges(
2593 &self,
2594 filter: EdgeFilter,
2595 sort: Vec<SortOrder<EdgeSortField>>,
2596 page: PageRequest,
2597 ) -> Result<Page<Edge>, StorageError> {
2598 let namespace = self.namespace.clone();
2599 let limit_i64 = i64::from(page.limit);
2600 let offset_i64 = i64::try_from(page.offset).map_err(|_| StorageError::InvalidInput {
2601 capability: StorageCapability::Graph,
2602 operation: "query_edges".into(),
2603 message: format!(
2604 "PageRequest: offset must be <= i64::MAX, got {}",
2605 page.offset
2606 ),
2607 })?;
2608 self.with_reader("query_edges", move |conn| {
2609 let (where_clause, mut all_params) = build_edge_filter_sql(&namespace, &filter);
2610 let order_clause = edge_order_clause(&sort);
2611 all_params.push(Box::new(limit_i64));
2612 all_params.push(Box::new(offset_i64));
2613
2614 let limit_idx = all_params.len() - 1;
2615 let offset_idx = all_params.len();
2616
2617 let data_sql = format!(
2618 "SELECT namespace, id, source_id, target_id, relation, weight, \
2619 created_at, updated_at, deleted_at, metadata, target_backend \
2620 FROM graph_edges{}{} LIMIT ?{} OFFSET ?{}",
2621 where_clause, order_clause, limit_idx, offset_idx,
2622 );
2623
2624 let mut stmt = conn.prepare(&data_sql)?;
2625 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
2626 all_params.iter().map(|p| p.as_ref()).collect();
2627 let rows = stmt.query_map(param_refs.as_slice(), read_edge)?;
2628
2629 let mut items = Vec::new();
2630 for row in rows {
2631 items.push(row?);
2632 }
2633
2634 Ok(Page { items, total: None })
2635 })
2636 .await
2637 }
2638
2639 async fn count_edges(&self, filter: EdgeFilter) -> Result<u64, StorageError> {
2640 let namespace = self.namespace.clone();
2641 self.with_reader("count_edges", move |conn| {
2642 let (where_clause, params) = build_edge_filter_sql(&namespace, &filter);
2643 let sql = format!(
2644 "SELECT COUNT(*) FROM graph_edges{}",
2645 with_live_endpoints(&where_clause)
2646 );
2647 let mut stmt = conn.prepare(&sql)?;
2648 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
2649 params.iter().map(|p| p.as_ref()).collect();
2650 let count: i64 = stmt.query_row(param_refs.as_slice(), |row| row.get(0))?;
2651 Ok(count as u64)
2652 })
2653 .await
2654 }
2655
2656 async fn count_edges_in_namespaces(
2657 &self,
2658 namespaces: &[String],
2659 filter: EdgeFilter,
2660 ) -> Result<u64, StorageError> {
2661 let namespaces: Vec<String> = namespaces
2662 .iter()
2663 .cloned()
2664 .collect::<HashSet<_>>()
2665 .into_iter()
2666 .collect();
2667 self.with_reader("count_edges_in_namespaces", move |conn| {
2668 let mut total = 0;
2669 for chunk in namespaces.chunks(NAMESPACE_COUNT_CHUNK_SIZE) {
2670 let (where_clause, params) = build_edge_filter_sql_for_namespaces(chunk, &filter);
2671 let sql = format!(
2672 "SELECT COUNT(*) FROM graph_edges{}",
2673 with_live_endpoints(&where_clause)
2674 );
2675 let mut stmt = conn.prepare(&sql)?;
2676 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
2677 params.iter().map(|p| p.as_ref()).collect();
2678 let count: i64 = stmt.query_row(param_refs.as_slice(), |row| row.get(0))?;
2679 total += count as u64;
2680 }
2681 Ok(total)
2682 })
2683 .await
2684 }
2685
2686 async fn query_edges_in_namespaces(
2687 &self,
2688 namespaces: &[String],
2689 filter: EdgeFilter,
2690 sort: Vec<SortOrder<EdgeSortField>>,
2691 page: PageRequest,
2692 ) -> Result<Page<Edge>, StorageError> {
2693 let namespaces: Vec<String> = {
2708 let mut seen = HashSet::new();
2709 namespaces
2710 .iter()
2711 .filter(|ns| seen.insert((*ns).clone()))
2712 .cloned()
2713 .collect()
2714 };
2715 let limit_i64 = i64::from(page.limit);
2716 let offset_i64 = i64::try_from(page.offset).map_err(|_| StorageError::InvalidInput {
2717 capability: StorageCapability::Graph,
2718 operation: "query_edges_in_namespaces".into(),
2719 message: format!(
2720 "PageRequest: offset must be <= i64::MAX, got {}",
2721 page.offset
2722 ),
2723 })?;
2724 self.with_reader("query_edges_in_namespaces", move |conn| {
2725 let namespaces_json = serde_json::to_string(&namespaces)
2726 .map_err(|error| rusqlite::Error::ToSqlConversionFailure(Box::new(error)))?;
2727
2728 let (where_clause, mut all_params) =
2729 build_edge_filter_sql_for_namespaces_json(&namespaces_json, &filter);
2730 let order_clause = edge_order_clause(&sort);
2731 all_params.push(Box::new(limit_i64));
2732 all_params.push(Box::new(offset_i64));
2733
2734 let limit_idx = all_params.len() - 1;
2735 let offset_idx = all_params.len();
2736
2737 let data_sql = format!(
2738 "SELECT namespace, id, source_id, target_id, relation, weight, \
2739 created_at, updated_at, deleted_at, metadata, target_backend \
2740 FROM graph_edges{}{} LIMIT ?{} OFFSET ?{}",
2741 where_clause, order_clause, limit_idx, offset_idx,
2742 );
2743
2744 let mut stmt = conn.prepare(&data_sql)?;
2745 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
2746 all_params.iter().map(|p| p.as_ref()).collect();
2747 let rows = stmt.query_map(param_refs.as_slice(), read_edge)?;
2748
2749 let mut items = Vec::new();
2750 for row in rows {
2751 items.push(row?);
2752 }
2753
2754 Ok(Page { items, total: None })
2755 })
2756 .await
2757 }
2758
2759 async fn count_edges_by_relation(&self) -> Result<Vec<(EdgeRelation, u64)>, StorageError> {
2760 let namespace = self.namespace.clone();
2761 self.with_reader("count_edges_by_relation", move |conn| {
2762 let sql = format!(
2763 "SELECT relation, COUNT(*) FROM graph_edges \
2764 WHERE namespace = ?1 AND deleted_at IS NULL AND {LIVE_ENDPOINTS_CONDITION} \
2765 GROUP BY relation"
2766 );
2767 let mut stmt = conn.prepare(&sql)?;
2768 let rows = stmt.query_map([&namespace], |row| {
2769 let relation_str: String = row.get(0)?;
2770 let count: i64 = row.get(1)?;
2771 Ok((relation_str, count))
2772 })?;
2773 let mut out = Vec::new();
2774 for row in rows {
2775 let (relation_str, count) = row?;
2776 let relation = relation_str.parse::<EdgeRelation>().map_err(|e| {
2777 rusqlite::Error::FromSqlConversionFailure(
2778 0,
2779 rusqlite::types::Type::Text,
2780 Box::new(e),
2781 )
2782 })?;
2783 out.push((relation, count as u64));
2784 }
2785 Ok(out)
2786 })
2787 .await
2788 }
2789
2790 async fn count_edges_by_relation_in_namespaces(
2791 &self,
2792 namespaces: &[String],
2793 ) -> Result<Vec<(EdgeRelation, u64)>, StorageError> {
2794 let namespaces: Vec<String> = namespaces
2795 .iter()
2796 .cloned()
2797 .collect::<HashSet<_>>()
2798 .into_iter()
2799 .collect();
2800 self.with_reader("count_edges_by_relation_in_namespaces", move |conn| {
2801 let mut totals = HashMap::new();
2802 for chunk in namespaces.chunks(NAMESPACE_COUNT_CHUNK_SIZE) {
2803 let (where_clause, params) =
2804 build_edge_filter_sql_for_namespaces(chunk, &EdgeFilter::default());
2805 let sql = format!(
2806 "SELECT relation, COUNT(*) FROM graph_edges{} GROUP BY relation",
2807 with_live_endpoints(&where_clause)
2808 );
2809 let mut stmt = conn.prepare(&sql)?;
2810 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
2811 params.iter().map(|p| p.as_ref()).collect();
2812 let rows = stmt.query_map(param_refs.as_slice(), |row| {
2813 let relation_str: String = row.get(0)?;
2814 let count: i64 = row.get(1)?;
2815 Ok((relation_str, count))
2816 })?;
2817 for row in rows {
2818 let (relation_str, count) = row?;
2819 let relation = relation_str.parse::<EdgeRelation>().map_err(|e| {
2820 rusqlite::Error::FromSqlConversionFailure(
2821 0,
2822 rusqlite::types::Type::Text,
2823 Box::new(e),
2824 )
2825 })?;
2826 *totals.entry(relation).or_insert(0) += count as u64;
2827 }
2828 }
2829 Ok(totals.into_iter().collect())
2830 })
2831 .await
2832 }
2833
2834 async fn count_edges_by_endpoint_base(&self) -> Result<EdgeEndpointBaseCounts, StorageError> {
2835 let namespace = self.namespace.clone();
2836 self.with_reader("count_edges_by_endpoint_base", move |conn| {
2837 let source_case = endpoint_base_case("source_id");
2838 let target_case = endpoint_base_case("target_id");
2839 let sql = format!(
2840 "SELECT {source_case} AS source_base, {target_case} AS target_base, COUNT(*) \
2841 FROM graph_edges \
2842 WHERE namespace = ?1 AND deleted_at IS NULL AND {LIVE_ENDPOINTS_CONDITION} \
2843 GROUP BY source_base, target_base"
2844 );
2845 let mut stmt = conn.prepare(&sql)?;
2846 let rows = stmt.query_map([&namespace], |row| {
2847 let source: String = row.get(0)?;
2848 let target: String = row.get(1)?;
2849 let count: i64 = row.get(2)?;
2850 Ok((source, target, count))
2851 })?;
2852 let mut counts = EdgeEndpointBaseCounts::default();
2853 for row in rows {
2854 let (source, target, count) = row?;
2855 fold_endpoint_base_row(&mut counts, &source, &target, count as u64);
2856 }
2857 Ok(counts)
2858 })
2859 .await
2860 }
2861
2862 async fn count_edges_by_endpoint_base_in_namespaces(
2863 &self,
2864 namespaces: &[String],
2865 ) -> Result<EdgeEndpointBaseCounts, StorageError> {
2866 let namespaces: Vec<String> = namespaces
2867 .iter()
2868 .cloned()
2869 .collect::<HashSet<_>>()
2870 .into_iter()
2871 .collect();
2872 self.with_reader("count_edges_by_endpoint_base_in_namespaces", move |conn| {
2873 let source_case = endpoint_base_case("source_id");
2874 let target_case = endpoint_base_case("target_id");
2875 let mut counts = EdgeEndpointBaseCounts::default();
2876 for chunk in namespaces.chunks(NAMESPACE_COUNT_CHUNK_SIZE) {
2877 let (where_clause, params) =
2878 build_edge_filter_sql_for_namespaces(chunk, &EdgeFilter::default());
2879 let sql = format!(
2880 "SELECT {source_case} AS source_base, {target_case} AS target_base, COUNT(*) \
2881 FROM graph_edges{} GROUP BY source_base, target_base",
2882 with_live_endpoints(&where_clause)
2883 );
2884 let mut stmt = conn.prepare(&sql)?;
2885 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
2886 params.iter().map(|p| p.as_ref()).collect();
2887 let rows = stmt.query_map(param_refs.as_slice(), |row| {
2888 let source: String = row.get(0)?;
2889 let target: String = row.get(1)?;
2890 let count: i64 = row.get(2)?;
2891 Ok((source, target, count))
2892 })?;
2893 for row in rows {
2894 let (source, target, count) = row?;
2895 fold_endpoint_base_row(&mut counts, &source, &target, count as u64);
2896 }
2897 }
2898 Ok(counts)
2899 })
2900 .await
2901 }
2902
2903 async fn query_edges_after(
2904 &self,
2905 filter: EdgeFilter,
2906 after: Option<Uuid>,
2907 limit: u32,
2908 ) -> Result<EdgeSeekPage, StorageError> {
2909 let namespace = self.namespace.clone();
2910 let limit_usize = limit as usize;
2911 let probe_limit_i64 = i64::from(limit) + 1;
2912 self.with_reader("query_edges_after", move |conn| {
2913 let (mut where_clause, mut params) = build_edge_filter_sql(&namespace, &filter);
2914 if let Some(cursor) = after {
2915 params.push(Box::new(cursor.to_string()));
2916 where_clause.push_str(&format!(" AND id > ?{}", params.len()));
2917 }
2918 params.push(Box::new(probe_limit_i64));
2919 let limit_idx = params.len();
2920
2921 let data_sql = format!(
2927 "SELECT namespace, id, source_id, target_id, relation, weight, \
2928 created_at, updated_at, deleted_at, metadata, target_backend \
2929 FROM graph_edges{} ORDER BY id ASC LIMIT ?{}",
2930 where_clause, limit_idx,
2931 );
2932
2933 let mut stmt = conn.prepare(&data_sql)?;
2934 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
2935 params.iter().map(|p| p.as_ref()).collect();
2936 let rows = stmt.query_map(param_refs.as_slice(), read_edge)?;
2937
2938 let mut items = Vec::new();
2939 for row in rows {
2940 items.push(row?);
2941 }
2942 let has_more = items.len() > limit_usize;
2943 if has_more {
2944 items.truncate(limit_usize);
2945 }
2946 let next_after = if has_more {
2947 items.last().map(|e| Uuid::from(e.id))
2948 } else {
2949 None
2950 };
2951
2952 Ok(EdgeSeekPage { items, next_after })
2953 })
2954 .await
2955 }
2956
2957 async fn query_edges_sequence_after(
2958 &self,
2959 filter: EdgeFilter,
2960 after: Option<SeekCursor>,
2961 limit: u32,
2962 ) -> Result<SeekPage<Edge>, StorageError> {
2963 if limit == 0 {
2964 return Ok(SeekPage::default());
2965 }
2966 let namespace = self.namespace.clone();
2967 let limit_usize = limit as usize;
2968 let probe_limit_i64 = i64::from(limit) + 1;
2969 self.with_reader("query_edges_sequence_after", move |conn| {
2970 let (mut where_clause, mut params) = build_edge_filter_sql(&namespace, &filter);
2971 if let Some(cursor) = after {
2972 params.push(Box::new(cursor.sequence));
2973 where_clause.push_str(&format!(" AND graph_edges_seq.seq > ?{}", params.len()));
2974 }
2975 params.push(Box::new(probe_limit_i64));
2976 let limit_idx = params.len();
2977 let sql = format!(
2980 "SELECT namespace, id, source_id, target_id, relation, weight, \
2981 created_at, updated_at, deleted_at, metadata, target_backend, \
2982 graph_edges_seq.seq \
2983 FROM graph_edges_seq CROSS JOIN graph_edges \
2984 ON graph_edges.id = graph_edges_seq.edge_id{where_clause} \
2985 ORDER BY graph_edges_seq.seq ASC LIMIT ?{limit_idx}"
2986 );
2987 let mut stmt = conn.prepare(&sql)?;
2988 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
2989 params.iter().map(|param| param.as_ref()).collect();
2990 let rows = stmt.query_map(param_refs.as_slice(), |row| {
2991 Ok((read_edge(row)?, row.get::<_, i64>(11)?))
2992 })?;
2993 let mut entries = rows.collect::<Result<Vec<_>, _>>()?;
2994 let has_more = entries.len() > limit_usize;
2995 if has_more {
2996 entries.truncate(limit_usize);
2997 }
2998 let next_after = if has_more {
2999 entries.last().map(|(edge, sequence)| SeekCursor {
3000 sequence: *sequence,
3001 id: Uuid::from(edge.id),
3002 })
3003 } else {
3004 None
3005 };
3006 let items = entries.into_iter().map(|(edge, _)| edge).collect();
3007 Ok(SeekPage { items, next_after })
3008 })
3009 .await
3010 }
3011
3012 async fn neighbors(
3013 &self,
3014 node_id: Uuid,
3015 query: NeighborQuery,
3016 ) -> Result<Vec<NeighborHit>, StorageError> {
3017 self.query_neighbors_page("neighbors", node_id, query, None, None)
3018 .await
3019 }
3020
3021 async fn neighbors_page(
3022 &self,
3023 node_id: Uuid,
3024 query: NeighborQuery,
3025 after: Option<NeighborCursor>,
3026 neighbor_kinds: Option<Vec<String>>,
3027 ) -> Result<Vec<NeighborHit>, StorageError> {
3028 self.query_neighbors_page("neighbors_page", node_id, query, after, neighbor_kinds)
3029 .await
3030 }
3031
3032 async fn neighbors_both_directions(
3037 &self,
3038 node_id: Uuid,
3039 query: NeighborQuery,
3040 ) -> Result<Vec<DirectedNeighborHit>, StorageError> {
3041 count_neighbor_select();
3042
3043 let namespace = self.namespace.clone();
3044 let node_str = node_id.to_string();
3045
3046 let counted_queries = Arc::new(std::sync::atomic::AtomicU64::new(0));
3047 let counted_rows = Arc::new(std::sync::atomic::AtomicU64::new(0));
3048 let closure_queries = Arc::clone(&counted_queries);
3049 let closure_rows = Arc::clone(&counted_rows);
3050 let result = self
3051 .with_reader("neighbors_both_directions", move |conn| {
3052 let base_out = "SELECT target_id AS node_id, id AS edge_id, relation, weight, \
3053 'out' AS dir \
3054 FROM graph_edges \
3055 WHERE namespace = ?1 AND source_id = ?2 AND deleted_at IS NULL";
3056 let base_in = "SELECT source_id AS node_id, id AS edge_id, relation, weight, \
3057 'in' AS dir \
3058 FROM graph_edges \
3059 WHERE namespace = ?1 AND target_id = ?2 AND deleted_at IS NULL";
3060 let sql = format!("{} UNION ALL {}", base_out, base_in);
3061
3062 let (where_extra, limit_clause, extra_params) =
3063 neighbor_extra_clause(&query, 3, None, None);
3064
3065 let full_sql = format!(
3074 "SELECT node_id, edge_id, relation, weight, dir FROM ({}){} \
3075 ORDER BY weight DESC, node_id ASC, \
3076 CASE dir WHEN 'out' THEN 0 ELSE 1 END ASC, edge_id ASC{}",
3077 sql, where_extra, limit_clause
3078 );
3079
3080 let mut stmt = conn.prepare(&full_sql)?;
3081
3082 let mut all_params: Vec<Box<dyn rusqlite::types::ToSql>> = Vec::new();
3083 all_params.push(Box::new(namespace.clone()));
3084 all_params.push(Box::new(node_str.clone()));
3085 all_params.extend(extra_params);
3086
3087 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
3088 all_params.iter().map(|p| p.as_ref()).collect();
3089
3090 closure_queries.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3091 let rows = stmt.query_map(param_refs.as_slice(), |row| {
3092 let nid_str: String = row.get(0)?;
3093 let eid_str: String = row.get(1)?;
3094 let relation_str: String = row.get(2)?;
3095 let weight: f64 = row.get(3)?;
3096 let dir_str: String = row.get(4)?;
3097 Ok((nid_str, eid_str, relation_str, weight, dir_str))
3098 })?;
3099
3100 let mut hits = Vec::new();
3101 for row in rows {
3102 let (nid_str, eid_str, relation_str, weight, dir_str) = row?;
3103 closure_rows.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3104 let relation = relation_str.parse::<EdgeRelation>().map_err(|e| {
3105 rusqlite::Error::FromSqlConversionFailure(
3106 2,
3107 rusqlite::types::Type::Text,
3108 Box::new(e),
3109 )
3110 })?;
3111 let direction = if dir_str == "out" {
3112 Direction::Out
3113 } else {
3114 Direction::In
3115 };
3116 hits.push(DirectedNeighborHit {
3117 hit: NeighborHit {
3118 node_id: parse_uuid(&nid_str)?,
3119 edge_id: parse_uuid(&eid_str)?,
3120 relation,
3121 weight,
3122 name: None,
3123 kind: None,
3124 entity_type: None,
3125 },
3126 direction,
3127 });
3128 }
3129
3130 Ok(hits)
3131 })
3132 .await;
3133
3134 report_graph_usage(&counted_queries, &counted_rows);
3135 result
3136 }
3137
3138 async fn traverse(&self, request: TraversalRequest) -> Result<Vec<GraphPath>, StorageError> {
3139 request
3140 .validate()
3141 .map_err(|message| StorageError::InvalidInput {
3142 capability: StorageCapability::Graph,
3143 operation: "traverse".into(),
3144 message,
3145 })?;
3146 if request.roots.is_empty() {
3147 return Ok(Vec::new());
3148 }
3149
3150 let mut distinct_roots = HashSet::with_capacity(request.roots.len());
3151 let roots = request
3152 .roots
3153 .iter()
3154 .copied()
3155 .filter(|root| distinct_roots.insert(*root))
3156 .collect::<Vec<_>>();
3157 let opts = request.options;
3158 let include_roots = request.include_roots;
3159 let namespace = self.namespace.clone();
3160 let origin = self.pool.origin();
3161 let budget = request.execution_budget;
3162 let counted_rows = Arc::new(std::sync::atomic::AtomicU64::new(0));
3168 let counted_queries = Arc::new(std::sync::atomic::AtomicU64::new(0));
3169 let closure_rows = Arc::clone(&counted_rows);
3170 let closure_queries = Arc::clone(&counted_queries);
3171 let result = self
3172 .with_reader("traverse", move |conn| {
3173 Ok(run_bounded_traversal(
3174 conn,
3175 roots,
3176 opts,
3177 include_roots,
3178 namespace,
3179 origin,
3180 budget,
3181 closure_rows.as_ref(),
3182 closure_queries.as_ref(),
3183 ))
3184 })
3185 .await
3186 .and_then(|inner| inner);
3187
3188 khive_storage::usage::count(
3197 khive_storage::usage::UsageUnit::DbRoundTrips,
3198 counted_queries.load(std::sync::atomic::Ordering::Relaxed),
3199 );
3200 khive_storage::usage::count(
3201 khive_storage::usage::UsageUnit::GraphHops,
3202 counted_rows.load(std::sync::atomic::Ordering::Relaxed),
3203 );
3204
3205 result
3206 }
3207
3208 async fn purge_incident_edges(&self, node_id: Uuid) -> Result<u64, StorageError> {
3209 let statement = purge_incident_edges_statement(node_id);
3213 self.with_writer("purge_incident_edges", move |conn| {
3214 let mut stmt = conn.prepare(&statement.sql)?;
3215 bind_params(&mut stmt, &statement.params)?;
3216 Ok(stmt.raw_execute()? as u64)
3217 })
3218 .await
3219 }
3220}
3221
3222const GRAPH_DDL: &str = include_str!("../../sql/graph-ddl.sql");
3227
3228pub(crate) fn ensure_graph_schema(conn: &rusqlite::Connection) -> Result<(), rusqlite::Error> {
3229 conn.execute_batch(GRAPH_DDL)
3230}
3231
3232#[cfg(test)]
3233#[path = "graph_annotation_tests.rs"]
3234mod annotation_tests;
3235
3236#[cfg(test)]
3237#[path = "graph_tests.rs"]
3238mod tests;