1#[path = "graph/write_transaction.rs"]
4mod write_transaction;
5use write_transaction::run_graph_mutation_transaction;
6
7#[path = "graph/symmetric_update.rs"]
8mod symmetric_update;
9pub use symmetric_update::{
10 edge_symmetric_conflict_probe_statement, edge_symmetric_delete_noncanonical_statement,
11 edge_symmetric_update_inplace_statement, EDGE_SYMMETRIC_CONFLICT_PROBE_SQL,
12 EDGE_SYMMETRIC_DELETE_NONCANONICAL_GUARDED_SQL, EDGE_SYMMETRIC_DELETE_NONCANONICAL_SQL,
13 EDGE_SYMMETRIC_UPDATE_INPLACE_SQL,
14};
15
16use std::collections::{HashMap, HashSet, VecDeque};
17use std::sync::Arc;
18
19use async_trait::async_trait;
20use chrono::{DateTime, TimeZone, Utc};
21use rusqlite::OptionalExtension;
22use uuid::Uuid;
23
24use khive_storage::error::StorageError;
25use khive_storage::graph::{
26 CommitAnnotationGuard, CommitAnnotationInsertOutcome, SymmetricEdgeUpdateOutcome,
27 SymmetricEdgeUpdateRequest,
28};
29use khive_storage::types::{
30 BatchWriteSummary, DeleteMode, DirectedNeighborHit, Direction, Edge, EdgeEndpointBaseCounts,
31 EdgeFilter, EdgeSeekPage, EdgeSortField, EdgeUpsertDisposition, EdgeUpsertRefusal,
32 EdgeUpsertRequest, EdgeUpsertResult, GraphPath, GuardedBatchOutcome, GuardedBatchRefusal,
33 GuardedEdgeBatchRefusal, GuardedEdgeBatchUpsertOutcome, GuardedEdgeUpsertOutcome,
34 GuardedWriteOutcome, MissingEndpoints, NeighborCursor, NeighborHit, NeighborQuery, Page,
35 PageRequest, PathNode, SeekCursor, SeekPage, SortDirection, SortOrder, SqlStatement, SqlValue,
36 TraversalExecutionBudget, TraversalOptions, TraversalRequest,
37};
38use khive_storage::LinkId;
39use khive_storage::StorageCapability;
40use khive_storage::{Event, GraphStore, StorageResult};
41use khive_types::EdgeRelation;
42
43use crate::error::SqliteError;
44use crate::pool::ConnectionPool;
45use crate::sql_bridge::bind_params;
46use crate::writer_task::WriterTaskHandle;
47
48fn map_err(e: rusqlite::Error, op: &'static str) -> StorageError {
50 StorageError::driver(StorageCapability::Graph, op, e)
51}
52
53fn map_sqlite_err(e: SqliteError, op: &'static str) -> StorageError {
54 e.into_storage_error(StorageCapability::Graph, op)
55}
56
57fn resurrection_required_error(operation: &'static str, edge: &Edge) -> StorageError {
58 StorageError::Conflict {
59 capability: StorageCapability::Graph,
60 operation: operation.into(),
61 message: format!(
62 "edge {} is soft-deleted; explicit resurrection is required",
63 edge.id
64 ),
65 }
66}
67
68const NAMESPACE_COUNT_CHUNK_SIZE: usize = 500;
69
70const LATEST_ANNOTATING_NOTE_SQL: &str = r#"WITH incident AS MATERIALIZED (
79 SELECT source_id, deleted_at
80 FROM graph_edges INDEXED BY idx_graph_edges_ns_tgt_rel
81 WHERE namespace = ?1 AND target_id = ?2 AND relation = 'annotates'
82 LIMIT 3
83)
84SELECT result.id, result.created_at
85FROM notes AS result
86WHERE result.id = CASE WHEN (SELECT count(*) FROM incident) <= 2 THEN (
87 SELECT n.id
88 FROM incident AS e CROSS JOIN notes AS n
89 WHERE n.id = e.source_id AND e.deleted_at IS NULL
90 AND n.deleted_at IS NULL AND n.kind = ?3
91 AND EXISTS (SELECT 1 FROM json_each(CASE
92 WHEN json_type(n.properties, '$.tags') = 'array'
93 THEN json_extract(n.properties, '$.tags') ELSE '[]' END) AS tag
94 WHERE tag.type = 'text' AND tag.value = ?4 COLLATE BINARY)
95 ORDER BY n.created_at DESC, n.id ASC LIMIT 1
96) ELSE (
97 SELECT n.id
98 FROM notes AS n INDEXED BY idx_notes_created
99 WHERE n.deleted_at IS NULL AND n.kind = ?3
100 AND EXISTS (SELECT 1 FROM json_each(CASE
101 WHEN json_type(n.properties, '$.tags') = 'array'
102 THEN json_extract(n.properties, '$.tags') ELSE '[]' END) AS tag
103 WHERE tag.type = 'text' AND tag.value = ?4 COLLATE BINARY)
104 AND EXISTS (SELECT 1 FROM graph_edges AS e INDEXED BY idx_graph_edges_unique_triple
105 WHERE e.namespace = ?1 AND e.source_id = n.id AND e.target_id = ?2
106 AND e.relation = 'annotates' AND e.deleted_at IS NULL)
107 ORDER BY n.created_at DESC, n.id ASC LIMIT 1
108) END"#;
109
110const LATEST_ANNOTATING_NOTE_WITH_PROPERTY_SQL: &str = r#"WITH incident AS MATERIALIZED (
114 SELECT source_id, deleted_at
115 FROM graph_edges INDEXED BY idx_graph_edges_ns_tgt_rel
116 WHERE namespace = ?1 AND target_id = ?2 AND relation = 'annotates'
117 LIMIT 3
118)
119SELECT result.id, result.created_at
120FROM notes AS result
121WHERE result.id = CASE WHEN (SELECT count(*) FROM incident) <= 2 THEN (
122 SELECT n.id
123 FROM incident AS e CROSS JOIN notes AS n
124 WHERE n.id = e.source_id AND e.deleted_at IS NULL
125 AND n.deleted_at IS NULL AND n.kind = ?3
126 AND EXISTS (SELECT 1 FROM json_each(CASE
127 WHEN json_type(n.properties, '$.tags') = 'array'
128 THEN json_extract(n.properties, '$.tags') ELSE '[]' END) AS tag
129 WHERE tag.type = 'text' AND tag.value = ?4 COLLATE BINARY)
130 AND EXISTS (SELECT 1 FROM json_each(n.properties) AS property
131 WHERE property.key = ?5 COLLATE BINARY AND property.type = 'text'
132 AND property.value = ?6 COLLATE BINARY)
133 ORDER BY n.created_at DESC, n.id ASC LIMIT 1
134) ELSE (
135 SELECT n.id
136 FROM notes AS n INDEXED BY idx_notes_created
137 WHERE n.deleted_at IS NULL AND n.kind = ?3
138 AND EXISTS (SELECT 1 FROM json_each(CASE
139 WHEN json_type(n.properties, '$.tags') = 'array'
140 THEN json_extract(n.properties, '$.tags') ELSE '[]' END) AS tag
141 WHERE tag.type = 'text' AND tag.value = ?4 COLLATE BINARY)
142 AND EXISTS (SELECT 1 FROM json_each(n.properties) AS property
143 WHERE property.key = ?5 COLLATE BINARY AND property.type = 'text'
144 AND property.value = ?6 COLLATE BINARY)
145 AND EXISTS (SELECT 1 FROM graph_edges AS e INDEXED BY idx_graph_edges_unique_triple
146 WHERE e.namespace = ?1 AND e.source_id = n.id AND e.target_id = ?2
147 AND e.relation = 'annotates' AND e.deleted_at IS NULL)
148 ORDER BY n.created_at DESC, n.id ASC LIMIT 1
149) END"#;
150
151fn edge_conflict_clause(resurrect: bool) -> String {
163 let deleted_at = if resurrect {
164 "NULL"
165 } else {
166 "graph_edges.deleted_at"
167 };
168 let predicate = if resurrect {
169 ""
170 } else {
171 " WHERE graph_edges.deleted_at IS NULL"
172 };
173 format!(
174 "weight = excluded.weight, \
175 updated_at = excluded.updated_at, \
176 deleted_at = {deleted_at}, \
177 metadata = excluded.metadata, \
178 target_backend = excluded.target_backend{predicate}"
179 )
180}
181
182fn endpoint_exists_clause(id_param: &str) -> String {
193 format!(
194 "EXISTS (SELECT 1 FROM entities WHERE id = {id_param} AND deleted_at IS NULL) \
195 OR EXISTS (SELECT 1 FROM notes WHERE id = {id_param} AND deleted_at IS NULL) \
196 OR EXISTS (SELECT 1 FROM events WHERE id = {id_param}) \
197 OR EXISTS (SELECT 1 FROM graph_edges WHERE id = {id_param} AND deleted_at IS NULL)"
198 )
199}
200
201pub fn edge_snapshot_assertion_statement(edge: &Edge, require_endpoints: bool) -> SqlStatement {
205 let mut sql = "SELECT 1 FROM graph_edges WHERE id=?1 AND namespace=?2 \
206 AND source_id=?3 AND target_id=?4 AND relation=?5 \
207 AND updated_at=?6 AND deleted_at IS NULL"
208 .to_string();
209 if require_endpoints {
210 sql.push_str(&format!(
211 " AND ({}) AND ({})",
212 endpoint_exists_clause("?3"),
213 endpoint_exists_clause("?4")
214 ));
215 }
216 SqlStatement {
217 sql,
218 params: vec![
219 SqlValue::Text(Uuid::from(edge.id).to_string()),
220 SqlValue::Text(edge.namespace.clone()),
221 SqlValue::Text(edge.source_id.to_string()),
222 SqlValue::Text(edge.target_id.to_string()),
223 SqlValue::Text(edge.relation.to_string()),
224 SqlValue::Integer(edge.updated_at.timestamp_micros()),
225 ],
226 label: Some("edge-snapshot-assertion".into()),
227 }
228}
229
230pub fn edge_upsert_statement(edge: &Edge) -> SqlStatement {
234 edge_upsert_statement_with_resurrection(edge, false)
235}
236
237pub fn edge_upsert_statement_with_resurrection(edge: &Edge, resurrect: bool) -> SqlStatement {
239 let (source_id, target_id) =
240 canonical_edge_endpoints(edge.relation, edge.source_id, edge.target_id);
241 let metadata_str = edge
242 .metadata
243 .as_ref()
244 .map(|v| serde_json::to_string(v).unwrap_or_default());
245 let conflict_clause = edge_conflict_clause(resurrect);
246 SqlStatement {
247 sql: format!(
248 "INSERT INTO graph_edges \
249 (namespace, id, source_id, target_id, relation, weight, \
250 created_at, updated_at, deleted_at, metadata, target_backend) \
251 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11) \
252 ON CONFLICT(namespace, id) DO UPDATE SET \
253 source_id = excluded.source_id, \
254 target_id = excluded.target_id, \
255 relation = excluded.relation, \
256 {conflict_clause} \
257 ON CONFLICT(namespace, source_id, target_id, relation) DO UPDATE SET \
258 {conflict_clause}"
259 ),
260 params: vec![
261 SqlValue::Text(edge.namespace.clone()),
262 SqlValue::Text(Uuid::from(edge.id).to_string()),
263 SqlValue::Text(source_id.to_string()),
264 SqlValue::Text(target_id.to_string()),
265 SqlValue::Text(edge.relation.to_string()),
266 SqlValue::Float(edge.weight),
267 SqlValue::Integer(edge.created_at.timestamp_micros()),
268 SqlValue::Integer(edge.updated_at.timestamp_micros()),
269 match edge.deleted_at {
270 Some(t) => SqlValue::Integer(t.timestamp_micros()),
271 None => SqlValue::Null,
272 },
273 match metadata_str {
274 Some(m) => SqlValue::Text(m),
275 None => SqlValue::Null,
276 },
277 match &edge.target_backend {
278 Some(b) => SqlValue::Text(b.clone()),
279 None => SqlValue::Null,
280 },
281 ],
282 label: Some("edge-upsert".to_string()),
283 }
284}
285
286pub fn edge_insert_only_guarded_by_endpoints_statement(edge: &Edge) -> SqlStatement {
291 let mut statement = edge_upsert_statement(edge);
292 let src_exists = endpoint_exists_clause("?3");
293 let tgt_exists = endpoint_exists_clause("?4");
294 statement.sql = format!(
295 "INSERT INTO graph_edges \
296 (namespace, id, source_id, target_id, relation, weight, \
297 created_at, updated_at, deleted_at, metadata, target_backend) \
298 SELECT ?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11 \
299 WHERE ({src_exists}) AND ({tgt_exists})"
300 );
301 statement.label = Some("edge-insert-only-where-endpoints-exist".to_string());
302 statement
303}
304
305pub fn edge_insert_if_absent_statement(edge: &Edge) -> SqlStatement {
309 let mut statement = edge_upsert_statement(edge);
310 statement.sql = "INSERT INTO graph_edges \
311 (namespace, id, source_id, target_id, relation, weight, \
312 created_at, updated_at, deleted_at, metadata, target_backend) \
313 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11) \
314 ON CONFLICT DO NOTHING"
315 .to_string();
316 statement.label = Some("edge-insert-if-absent".to_string());
317 statement
318}
319
320pub fn edge_replace_if_unchanged_statement(
331 edge: &Edge,
332 expected_updated_at: DateTime<Utc>,
333 expected_deleted_at: Option<DateTime<Utc>>,
334) -> SqlStatement {
335 let (source_id, target_id) =
336 canonical_edge_endpoints(edge.relation, edge.source_id, edge.target_id);
337 let metadata_str = edge
338 .metadata
339 .as_ref()
340 .map(|v| serde_json::to_string(v).unwrap_or_default());
341 SqlStatement {
342 sql: "UPDATE graph_edges SET \
343 namespace = ?1, source_id = ?2, target_id = ?3, relation = ?4, weight = ?5, \
344 updated_at = ?6, deleted_at = ?7, metadata = ?8, target_backend = ?9 \
345 WHERE id = ?10 AND updated_at = ?11 AND deleted_at IS ?12 \
346 AND ?6 > updated_at"
347 .to_string(),
348 params: vec![
349 SqlValue::Text(edge.namespace.clone()),
350 SqlValue::Text(source_id.to_string()),
351 SqlValue::Text(target_id.to_string()),
352 SqlValue::Text(edge.relation.to_string()),
353 SqlValue::Float(edge.weight),
354 SqlValue::Integer(edge.updated_at.timestamp_micros()),
355 match edge.deleted_at {
356 Some(t) => SqlValue::Integer(t.timestamp_micros()),
357 None => SqlValue::Null,
358 },
359 match metadata_str {
360 Some(m) => SqlValue::Text(m),
361 None => SqlValue::Null,
362 },
363 match &edge.target_backend {
364 Some(b) => SqlValue::Text(b.clone()),
365 None => SqlValue::Null,
366 },
367 SqlValue::Text(Uuid::from(edge.id).to_string()),
368 SqlValue::Integer(expected_updated_at.timestamp_micros()),
369 match expected_deleted_at {
370 Some(value) => SqlValue::Integer(value.timestamp_micros()),
371 None => SqlValue::Null,
372 },
373 ],
374 label: Some("edge-replace-if-unchanged".to_string()),
375 }
376}
377
378#[allow(clippy::too_many_arguments)]
398pub fn edge_insert_guarded_by_endpoints_statement(
399 namespace: &str,
400 edge_id: Uuid,
401 source_id: Uuid,
402 target_id: Uuid,
403 relation: EdgeRelation,
404 weight: f64,
405 now: i64,
406 metadata: Option<&str>,
407) -> SqlStatement {
408 edge_insert_guarded_by_endpoints_with_resurrection_statement(
409 namespace, source_id, target_id, edge_id, relation, weight, now, metadata, false,
410 )
411}
412
413#[allow(clippy::too_many_arguments)]
415pub fn edge_insert_guarded_by_endpoints_with_resurrection_statement(
416 namespace: &str,
417 source_id: Uuid,
418 target_id: Uuid,
419 edge_id: Uuid,
420 relation: EdgeRelation,
421 weight: f64,
422 now: i64,
423 metadata: Option<&str>,
424 resurrect: bool,
425) -> SqlStatement {
426 let src_exists = endpoint_exists_clause("?3");
427 let tgt_exists = endpoint_exists_clause("?4");
428 let conflict_clause = edge_conflict_clause(resurrect);
429 SqlStatement {
430 sql: format!(
431 "INSERT INTO graph_edges \
432 (namespace, id, source_id, target_id, relation, weight, \
433 created_at, updated_at, metadata) \
434 SELECT ?1, ?2, ?3, ?4, ?5, ?6, ?7, ?7, ?8 \
435 WHERE ({src_exists}) AND ({tgt_exists}) \
436 ON CONFLICT(namespace, source_id, target_id, relation) DO UPDATE SET \
437 {conflict_clause}"
438 ),
439 params: vec![
440 SqlValue::Text(namespace.to_string()),
441 SqlValue::Text(edge_id.to_string()),
442 SqlValue::Text(source_id.to_string()),
443 SqlValue::Text(target_id.to_string()),
444 SqlValue::Text(relation.as_str().to_string()),
445 SqlValue::Float(weight),
446 SqlValue::Integer(now),
447 match metadata {
448 Some(m) => SqlValue::Text(m.to_string()),
449 None => SqlValue::Null,
450 },
451 ],
452 label: Some("atomic-link-insert-edge-where-exists".to_string()),
453 }
454}
455
456#[allow(clippy::too_many_arguments)]
460pub fn edge_insert_new_guarded_by_endpoints_statement(
461 namespace: &str,
462 edge_id: Uuid,
463 source_id: Uuid,
464 target_id: Uuid,
465 relation: EdgeRelation,
466 weight: f64,
467 now: i64,
468 metadata: Option<&str>,
469) -> SqlStatement {
470 let src_exists = endpoint_exists_clause("?3");
471 let tgt_exists = endpoint_exists_clause("?4");
472 SqlStatement {
473 sql: format!(
474 "INSERT INTO graph_edges \
475 (namespace, id, source_id, target_id, relation, weight, \
476 created_at, updated_at, metadata) \
477 SELECT ?1, ?2, ?3, ?4, ?5, ?6, ?7, ?7, ?8 \
478 WHERE ({src_exists}) AND ({tgt_exists}) \
479 ON CONFLICT DO NOTHING"
480 ),
481 params: vec![
482 SqlValue::Text(namespace.to_string()),
483 SqlValue::Text(edge_id.to_string()),
484 SqlValue::Text(source_id.to_string()),
485 SqlValue::Text(target_id.to_string()),
486 SqlValue::Text(relation.to_string()),
487 SqlValue::Float(weight),
488 SqlValue::Integer(now),
489 metadata
490 .map(|value| SqlValue::Text(value.to_string()))
491 .unwrap_or(SqlValue::Null),
492 ],
493 label: Some("edge-link-create-if-absent-and-endpoints-exist".to_string()),
494 }
495}
496
497pub fn edge_link_replace_if_unchanged_and_endpoints_exist_statement(
503 previous: &Edge,
504 weight: f64,
505 now: i64,
506 metadata: Option<&str>,
507) -> SqlStatement {
508 let src_exists = endpoint_exists_clause("?6");
509 let tgt_exists = endpoint_exists_clause("?7");
510 SqlStatement {
511 sql: format!(
512 "UPDATE graph_edges SET \
513 weight = ?1, updated_at = ?2, deleted_at = NULL, \
514 metadata = ?3, target_backend = NULL \
515 WHERE namespace = ?4 AND id = ?5 \
516 AND source_id = ?6 AND target_id = ?7 AND relation = ?8 \
517 AND updated_at = ?9 AND deleted_at IS ?10 AND ?2 > updated_at \
518 AND ({src_exists}) AND ({tgt_exists})"
519 ),
520 params: vec![
521 SqlValue::Float(weight),
522 SqlValue::Integer(now),
523 metadata
524 .map(|value| SqlValue::Text(value.to_string()))
525 .unwrap_or(SqlValue::Null),
526 SqlValue::Text(previous.namespace.clone()),
527 SqlValue::Text(Uuid::from(previous.id).to_string()),
528 SqlValue::Text(previous.source_id.to_string()),
529 SqlValue::Text(previous.target_id.to_string()),
530 SqlValue::Text(previous.relation.to_string()),
531 SqlValue::Integer(previous.updated_at.timestamp_micros()),
532 previous
533 .deleted_at
534 .map(|value| SqlValue::Integer(value.timestamp_micros()))
535 .unwrap_or(SqlValue::Null),
536 ],
537 label: Some("edge-link-replace-if-unchanged-and-endpoints-exist".to_string()),
538 }
539}
540
541pub fn edge_soft_delete_statement(id: Uuid, now: i64) -> SqlStatement {
543 SqlStatement {
544 sql: "UPDATE graph_edges SET deleted_at = ?2, updated_at = ?2 \
545 WHERE id = ?1 AND deleted_at IS NULL"
546 .to_string(),
547 params: vec![SqlValue::Text(id.to_string()), SqlValue::Integer(now)],
548 label: Some("edge-delete-soft".to_string()),
549 }
550}
551
552pub fn edge_hard_delete_statement(id: Uuid) -> SqlStatement {
554 SqlStatement {
555 sql: "DELETE FROM graph_edges WHERE id = ?1".to_string(),
556 params: vec![SqlValue::Text(id.to_string())],
557 label: Some("edge-delete-hard".to_string()),
558 }
559}
560
561pub(crate) const EDGE_NAMESPACES_CTE: &str = "WITH RECURSIVE edge_ns(namespace) AS ( \
572 SELECT (SELECT min(namespace) FROM graph_edges) \
573 UNION ALL \
574 SELECT (SELECT min(namespace) FROM graph_edges WHERE namespace > edge_ns.namespace) \
575 FROM edge_ns WHERE edge_ns.namespace IS NOT NULL)";
576
577pub fn purge_incident_edges_statement(node_id: Uuid) -> SqlStatement {
583 SqlStatement {
584 sql: format!(
585 "DELETE FROM graph_edges WHERE id IN ({EDGE_NAMESPACES_CTE} \
586 SELECT e.id FROM edge_ns CROSS JOIN graph_edges AS e \
587 ON e.namespace = edge_ns.namespace AND e.source_id = ?1 \
588 UNION ALL \
589 SELECT e.id FROM edge_ns CROSS JOIN graph_edges AS e \
590 ON e.namespace = edge_ns.namespace AND e.target_id = ?1)"
591 ),
592 params: vec![SqlValue::Text(node_id.to_string())],
593 label: Some("edge-purge-incident".to_string()),
594 }
595}
596
597#[allow(clippy::too_many_arguments)]
691pub fn edge_symmetric_delete_if_conflict_statement(
692 namespace: &str,
693 id: Uuid,
694 canon_src: Uuid,
695 canon_tgt: Uuid,
696 relation: EdgeRelation,
697 expected_updated_at_micros: i64,
698 expected_deleted_at_micros: Option<i64>,
699) -> SqlStatement {
700 SqlStatement {
701 sql: "DELETE FROM graph_edges \
702 WHERE namespace = ?1 AND id = ?2 \
703 AND updated_at = ?6 AND deleted_at IS ?7 \
704 AND EXISTS ( \
705 SELECT 1 FROM graph_edges \
706 WHERE namespace = ?1 AND source_id = ?3 AND target_id = ?4 \
707 AND relation = ?5 AND id != ?2 \
708 )"
709 .to_string(),
710 params: vec![
711 SqlValue::Text(namespace.to_string()),
712 SqlValue::Text(id.to_string()),
713 SqlValue::Text(canon_src.to_string()),
714 SqlValue::Text(canon_tgt.to_string()),
715 SqlValue::Text(relation.to_string()),
716 SqlValue::Integer(expected_updated_at_micros),
717 match expected_deleted_at_micros {
718 Some(value) => SqlValue::Integer(value),
719 None => SqlValue::Null,
720 },
721 ],
722 label: Some("edge-symmetric-delete-if-conflict".to_string()),
723 }
724}
725
726#[allow(clippy::too_many_arguments)]
735pub fn edge_symmetric_absorb_or_update_inplace_statement(
736 namespace: &str,
737 id: Uuid,
738 canon_src: Uuid,
739 canon_tgt: Uuid,
740 relation: EdgeRelation,
741 weight: f64,
742 updated_at_micros: i64,
743 metadata: Option<&str>,
744 target_backend: Option<&str>,
745 expected_updated_at_micros: i64,
746 expected_deleted_at_micros: Option<i64>,
747) -> SqlStatement {
748 SqlStatement {
749 sql: "UPDATE graph_edges SET \
750 source_id = CASE WHEN id = ?2 THEN ?3 ELSE source_id END, \
751 target_id = CASE WHEN id = ?2 THEN ?4 ELSE target_id END, \
752 relation = CASE WHEN id = ?2 THEN ?5 ELSE relation END, \
753 weight = CASE WHEN id = ?2 THEN ?6 ELSE weight END, \
754 updated_at = CASE WHEN id = ?2 THEN ?7 ELSE updated_at END, \
755 deleted_at = CASE WHEN id = ?2 THEN NULL ELSE deleted_at END, \
756 metadata = CASE WHEN id = ?2 THEN ?8 ELSE metadata END, \
757 target_backend = CASE WHEN id = ?2 THEN ?9 ELSE target_backend END \
758 WHERE namespace = ?1 \
759 AND ( \
760 (id = ?2 AND changes() = 0 AND updated_at = ?10 AND deleted_at IS ?11 \
761 AND ?7 > updated_at) \
762 OR (source_id = ?3 AND target_id = ?4 AND relation = ?5 \
763 AND id != ?2 AND changes() = 1) \
764 )"
765 .to_string(),
766 params: vec![
767 SqlValue::Text(namespace.to_string()),
768 SqlValue::Text(id.to_string()),
769 SqlValue::Text(canon_src.to_string()),
770 SqlValue::Text(canon_tgt.to_string()),
771 SqlValue::Text(relation.to_string()),
772 SqlValue::Float(weight),
773 SqlValue::Integer(updated_at_micros),
774 match metadata {
775 Some(m) => SqlValue::Text(m.to_string()),
776 None => SqlValue::Null,
777 },
778 match target_backend {
779 Some(b) => SqlValue::Text(b.to_string()),
780 None => SqlValue::Null,
781 },
782 SqlValue::Integer(expected_updated_at_micros),
783 match expected_deleted_at_micros {
784 Some(value) => SqlValue::Integer(value),
785 None => SqlValue::Null,
786 },
787 ],
788 label: Some("edge-symmetric-absorb-or-update-inplace".to_string()),
789 }
790}
791
792#[doc(hidden)]
796pub struct GraphDocumentGuard {
797 pub namespace: String,
798 pub id: Uuid,
799 pub expected_blob_ref: String,
800}
801
802#[doc(hidden)]
806pub struct GraphEdgeSnapshotGuard {
807 pub namespace: String,
808 pub source_id: Uuid,
809 pub target_id: Uuid,
810 pub relation: EdgeRelation,
811 pub expected: Option<Edge>,
812}
813
814#[doc(hidden)]
818#[derive(Default)]
819pub struct GraphMutationPreconditions {
820 pub document: Option<GraphDocumentGuard>,
821 pub edges: Vec<GraphEdgeSnapshotGuard>,
822}
823
824#[doc(hidden)]
828pub enum GraphMutationRequest {
829 Single {
830 request: EdgeUpsertRequest,
831 guard_endpoints: bool,
832 },
833 Batch {
834 requests: Vec<EdgeUpsertRequest>,
835 guard_endpoints: bool,
836 },
837 CommitAnnotation {
838 edge: Edge,
839 guard: CommitAnnotationGuard,
840 },
841}
842
843#[doc(hidden)]
847pub enum GraphMutationOutcome {
848 Single(GuardedEdgeUpsertOutcome),
849 Batch(GuardedEdgeBatchUpsertOutcome),
850 CommitAnnotation(CommitAnnotationInsertOutcome),
851}
852
853#[doc(hidden)]
857pub struct GraphMutationEventOutcome {
858 pub mutation: GraphMutationOutcome,
859 pub retired: Vec<Edge>,
860}
861
862const GRAPH_MUTATION_EVENTS_OP: &str = "compose_graph_mutation_events";
863
864fn count_graph_mutation_events(outcome: &GraphMutationEventOutcome) {
865 let written = match &outcome.mutation {
866 GraphMutationOutcome::Single(GuardedEdgeUpsertOutcome::Written(_)) => 1,
867 GraphMutationOutcome::Batch(outcome) if outcome.refusal.is_none() => outcome.rows.len(),
868 GraphMutationOutcome::CommitAnnotation(CommitAnnotationInsertOutcome::Created(_)) => 1,
869 _ => 0,
870 };
871 khive_storage::usage::count(
872 khive_storage::usage::UsageUnit::EventRows,
873 (written + outcome.retired.len()) as u64,
874 );
875}
876
877#[doc(hidden)]
881pub async fn compose_graph_mutation_events<F>(
882 backend: &crate::StorageBackend,
883 mutation: GraphMutationRequest,
884 preconditions: GraphMutationPreconditions,
885 retirements: Vec<Edge>,
886 make_events: F,
887) -> StorageResult<GraphMutationEventOutcome>
888where
889 F: FnOnce(&GraphMutationEventOutcome) -> StorageResult<Vec<Event>> + Send + 'static,
890{
891 if backend.is_read_only() {
892 return Err(StorageError::Pool {
893 operation: GRAPH_MUTATION_EVENTS_OP.into(),
894 message: "backend is read-only".into(),
895 });
896 }
897 backend
900 .graph()
901 .map_err(|error| map_sqlite_err(error, GRAPH_MUTATION_EVENTS_OP))?;
902 backend
903 .events()
904 .map_err(|error| map_sqlite_err(error, GRAPH_MUTATION_EVENTS_OP))?;
905 let pool = backend.pool_arc();
906 if let Some(writer_task) = pool.writer_task_for_write(None, GRAPH_MUTATION_EVENTS_OP)? {
907 return writer_task
908 .send_bounded_at(
909 move |conn| {
910 graph_mutation_events_enlisted(
911 conn,
912 mutation,
913 preconditions,
914 retirements,
915 make_events,
916 )
917 },
918 GRAPH_MUTATION_EVENTS_OP,
919 )
920 .await
921 .inspect(count_graph_mutation_events)
922 .inspect_err(|error| khive_storage::usage::account_event_write(Err(error)));
923 }
924 pool.record_direct_route(crate::timeout_sink::Site::DirectRouteGraphGeneralWrite);
925 let is_file_backed = backend.is_file_backed();
926 tokio::task::spawn_blocking(move || {
927 if is_file_backed {
928 let admission = pool.write_admission();
929 let _lease = admission
930 .acquire()
931 .map_err(|error| map_sqlite_err(error, GRAPH_MUTATION_EVENTS_OP))?;
932 let conn = pool
933 .open_standalone_writer_for_admitted_operation()
934 .map_err(|error| map_sqlite_err(error, GRAPH_MUTATION_EVENTS_OP))?;
935 run_graph_mutation_transaction(&pool, &conn, false, move |conn| {
936 graph_mutation_events_enlisted(
937 conn,
938 mutation,
939 preconditions,
940 retirements,
941 make_events,
942 )
943 })
944 } else {
945 let guard = pool
946 .writer_for_admitted_operation()
947 .map_err(|error| map_sqlite_err(error, GRAPH_MUTATION_EVENTS_OP))?;
948 run_graph_mutation_transaction(&pool, guard.conn(), true, move |conn| {
949 graph_mutation_events_enlisted(
950 conn,
951 mutation,
952 preconditions,
953 retirements,
954 make_events,
955 )
956 })
957 }
958 })
959 .await
960 .map_err(|error| {
961 StorageError::driver(StorageCapability::Graph, GRAPH_MUTATION_EVENTS_OP, error)
962 })?
963 .inspect(count_graph_mutation_events)
964 .inspect_err(|error| khive_storage::usage::account_event_write(Err(error)))
965}
966
967fn graph_mutation_conflict(message: &'static str) -> StorageError {
968 StorageError::Conflict {
969 capability: StorageCapability::Graph,
970 operation: GRAPH_MUTATION_EVENTS_OP.into(),
971 message: message.into(),
972 }
973}
974
975fn edge_snapshot_matches(actual: &Edge, expected: &Edge) -> bool {
976 actual.id == expected.id
977 && actual.namespace == expected.namespace
978 && actual.source_id == expected.source_id
979 && actual.target_id == expected.target_id
980 && actual.relation == expected.relation
981 && actual.weight.to_bits() == expected.weight.to_bits()
982 && actual.metadata == expected.metadata
983 && actual.target_backend == expected.target_backend
984 && actual.created_at == expected.created_at
985 && actual.updated_at == expected.updated_at
986 && actual.deleted_at == expected.deleted_at
987}
988
989fn check_graph_mutation_preconditions(
990 conn: &rusqlite::Connection,
991 preconditions: &GraphMutationPreconditions,
992 retirements: &[Edge],
993) -> StorageResult<()> {
994 if let Some(document) = &preconditions.document {
995 let matches: bool = conn
996 .query_row(
997 "SELECT EXISTS(SELECT 1 FROM entities WHERE id=?1 AND namespace=?2 \
998 AND deleted_at IS NULL AND json_type(properties, '$.blob_ref')='text' \
999 AND json_extract(properties, '$.blob_ref')=?3 COLLATE BINARY)",
1000 rusqlite::params![
1001 document.id.to_string(),
1002 &document.namespace,
1003 &document.expected_blob_ref,
1004 ],
1005 |row| row.get(0),
1006 )
1007 .map_err(|error| map_err(error, GRAPH_MUTATION_EVENTS_OP))?;
1008 if !matches {
1009 return Err(graph_mutation_conflict(
1010 "source document body reference changed",
1011 ));
1012 }
1013 }
1014 for guard in &preconditions.edges {
1015 let actual = edge_by_natural_key_parts_including_deleted(
1016 conn,
1017 &guard.namespace,
1018 guard.source_id,
1019 guard.target_id,
1020 guard.relation,
1021 )
1022 .map_err(|error| map_err(error, GRAPH_MUTATION_EVENTS_OP))?;
1023 let matches = match (actual.as_ref(), guard.expected.as_ref()) {
1024 (None, None) => true,
1025 (Some(actual), Some(expected)) => edge_snapshot_matches(actual, expected),
1026 _ => false,
1027 };
1028 if !matches {
1029 return Err(graph_mutation_conflict("edge ownership snapshot changed"));
1030 }
1031 }
1032 for expected in retirements {
1033 let actual = edge_by_natural_key_including_deleted(conn, expected)
1034 .map_err(|error| map_err(error, GRAPH_MUTATION_EVENTS_OP))?;
1035 if expected.deleted_at.is_some()
1036 || !actual
1037 .as_ref()
1038 .is_some_and(|actual| edge_snapshot_matches(actual, expected))
1039 {
1040 return Err(graph_mutation_conflict("edge retirement snapshot changed"));
1041 }
1042 }
1043 Ok(())
1044}
1045
1046fn graph_mutation_events_enlisted<F>(
1048 conn: &rusqlite::Connection,
1049 mutation: GraphMutationRequest,
1050 preconditions: GraphMutationPreconditions,
1051 retirements: Vec<Edge>,
1052 make_events: F,
1053) -> StorageResult<GraphMutationEventOutcome>
1054where
1055 F: FnOnce(&GraphMutationEventOutcome) -> StorageResult<Vec<Event>>,
1056{
1057 if !retirements.is_empty() && !matches!(&mutation, GraphMutationRequest::Batch { .. }) {
1058 return Err(StorageError::InvalidInput {
1059 capability: StorageCapability::Graph,
1060 operation: GRAPH_MUTATION_EVENTS_OP.into(),
1061 message: "retirements require a batch mutation".into(),
1062 });
1063 }
1064 check_graph_mutation_preconditions(conn, &preconditions, &retirements)?;
1065 let mutation = match mutation {
1066 GraphMutationRequest::Single {
1067 request,
1068 guard_endpoints,
1069 } => GraphMutationOutcome::Single(
1070 observed_edge_upsert(conn, &request, guard_endpoints)
1071 .map_err(|error| map_err(error, GRAPH_MUTATION_EVENTS_OP))?,
1072 ),
1073 GraphMutationRequest::Batch {
1074 requests,
1075 guard_endpoints,
1076 } => {
1077 let outcome = observed_edge_batch_upsert(conn, &requests, guard_endpoints)
1078 .map_err(|error| map_err(error, GRAPH_MUTATION_EVENTS_OP))?;
1079 if outcome.refusal.is_none() && outcome.rows.len() != requests.len() {
1080 return Err(graph_mutation_conflict(
1081 "edge result count differs from request count",
1082 ));
1083 }
1084 GraphMutationOutcome::Batch(outcome)
1085 }
1086 GraphMutationRequest::CommitAnnotation { edge, guard } => {
1087 if edge.relation != EdgeRelation::Annotates || edge.deleted_at.is_some() {
1088 return Err(StorageError::InvalidInput {
1089 capability: StorageCapability::Graph,
1090 operation: GRAPH_MUTATION_EVENTS_OP.into(),
1091 message: "expected a live annotates edge".into(),
1092 });
1093 }
1094 GraphMutationOutcome::CommitAnnotation(
1095 conditional_commit_annotation_insert(conn, edge, &guard)
1096 .map_err(|error| map_err(error, GRAPH_MUTATION_EVENTS_OP))?,
1097 )
1098 }
1099 };
1100 let written = match &mutation {
1101 GraphMutationOutcome::Single(GuardedEdgeUpsertOutcome::Written(_)) => 1,
1102 GraphMutationOutcome::Batch(outcome) if outcome.refusal.is_none() => outcome.rows.len(),
1103 GraphMutationOutcome::CommitAnnotation(CommitAnnotationInsertOutcome::Created(_)) => 1,
1104 _ => {
1105 return Ok(GraphMutationEventOutcome {
1106 mutation,
1107 retired: Vec::new(),
1108 })
1109 }
1110 };
1111 let mut retired = Vec::with_capacity(retirements.len());
1112 for edge in retirements {
1113 let statement =
1114 edge_soft_delete_statement(Uuid::from(edge.id), Utc::now().timestamp_micros());
1115 let mut stmt = conn
1116 .prepare(&statement.sql)
1117 .map_err(|error| map_err(error, GRAPH_MUTATION_EVENTS_OP))?;
1118 bind_params(&mut stmt, &statement.params)
1119 .map_err(|error| map_err(error, GRAPH_MUTATION_EVENTS_OP))?;
1120 if stmt
1121 .raw_execute()
1122 .map_err(|error| map_err(error, GRAPH_MUTATION_EVENTS_OP))?
1123 != 1
1124 {
1125 return Err(graph_mutation_conflict(
1126 "edge retirement changed during mutation",
1127 ));
1128 }
1129 retired.push(edge);
1130 }
1131 let outcome = GraphMutationEventOutcome { mutation, retired };
1132 let expected_events = written + outcome.retired.len();
1133 if expected_events == 0 {
1134 return Ok(outcome);
1135 }
1136 let events = make_events(&outcome)?;
1137 if events.len() != expected_events {
1138 return Err(graph_mutation_conflict(
1139 "event result count differs from mutation count",
1140 ));
1141 }
1142 for event in &events {
1143 super::event::append_event_in_transaction(conn, event).map_err(|error| {
1144 StorageError::driver(StorageCapability::Events, GRAPH_MUTATION_EVENTS_OP, error)
1145 })?;
1146 }
1147 Ok(outcome)
1148}
1149
1150pub struct SqlGraphStore {
1152 pool: Arc<ConnectionPool>,
1153 index_repair: Option<super::index_repair::IndexRepairContext>,
1154 is_file_backed: bool,
1155 namespace: String,
1159 writer_task: Option<WriterTaskHandle>,
1160}
1161
1162impl SqlGraphStore {
1163 pub fn new_scoped(
1171 pool: Arc<ConnectionPool>,
1172 is_file_backed: bool,
1173 namespace: impl Into<String>,
1174 ) -> Self {
1175 let writer_task = pool.writer_task_handle().ok().flatten();
1181
1182 Self {
1183 pool,
1184 index_repair: None,
1185 is_file_backed,
1186 namespace: namespace.into(),
1187 writer_task,
1188 }
1189 }
1190
1191 pub(crate) fn with_index_repair(
1192 mut self,
1193 repair: super::index_repair::IndexRepairContext,
1194 ) -> Self {
1195 self.index_repair = Some(repair);
1196 self
1197 }
1198
1199 async fn with_indexed_reader<F, R>(&self, op: &'static str, read: F) -> Result<R, StorageError>
1200 where
1201 F: FnMut(&rusqlite::Connection) -> Result<R, rusqlite::Error> + Send + 'static,
1202 R: Send + 'static,
1203 {
1204 super::index_repair::run_indexed_read(
1205 Arc::clone(&self.pool),
1206 self.index_repair.clone(),
1207 StorageCapability::Graph,
1208 op,
1209 read,
1210 )
1211 .await
1212 }
1213
1214 fn current_writer_task(
1215 &self,
1216 operation: &'static str,
1217 ) -> Result<Option<WriterTaskHandle>, StorageError> {
1218 self.pool
1219 .writer_task_for_write(self.writer_task.as_ref(), operation)
1220 }
1221
1222 async fn with_writer<F, R>(&self, op: &'static str, f: F) -> Result<R, StorageError>
1229 where
1230 F: FnOnce(&rusqlite::Connection) -> Result<R, rusqlite::Error> + Send + 'static,
1231 R: Send + 'static,
1232 {
1233 if let Some(writer_task) = self.current_writer_task(op)? {
1234 return writer_task
1235 .send_bounded_at(move |conn| f(conn).map_err(|e| map_err(e, op)), op)
1236 .await;
1237 }
1238
1239 self.pool
1240 .record_direct_route(crate::timeout_sink::Site::DirectRouteGraphGeneralWrite);
1241 let pool = Arc::clone(&self.pool);
1242 let db = crate::timeout_sink::db_label(&pool);
1243 let is_file_backed = self.is_file_backed;
1244 tokio::task::spawn_blocking(move || {
1245 let result =
1246 pool.execute_direct_transaction(StorageCapability::Graph, op, move |conn| {
1247 f(conn).map_err(|error| map_err(error, op))
1248 });
1249 if is_file_backed {
1250 if let Err(error) = &result {
1251 crate::timeout_sink::maybe_emit_busy_storage_error(
1252 &db,
1253 crate::timeout_sink::Site::StandaloneGraph,
1254 error,
1255 );
1256 }
1257 }
1258 result
1259 })
1260 .await
1261 .map_err(|e| StorageError::driver(StorageCapability::Graph, op, e))?
1262 }
1263
1264 async fn observed_edge_write(
1265 &self,
1266 operation: &'static str,
1267 request: EdgeUpsertRequest,
1268 guard_endpoints: bool,
1269 ) -> Result<GuardedEdgeUpsertOutcome, StorageError> {
1270 if let Some(writer_task) = self.current_writer_task(operation)? {
1271 return writer_task
1272 .send_bounded_at(
1273 move |conn| {
1274 observed_edge_upsert(conn, &request, guard_endpoints)
1275 .map_err(|error| map_err(error, operation))
1276 },
1277 operation,
1278 )
1279 .await;
1280 }
1281
1282 self.with_writer(operation, move |conn| {
1283 observed_edge_upsert(conn, &request, guard_endpoints)
1284 })
1285 .await
1286 }
1287
1288 async fn observed_edge_batch_write(
1289 &self,
1290 operation: &'static str,
1291 requests: Vec<EdgeUpsertRequest>,
1292 guard_endpoints: bool,
1293 ) -> Result<GuardedEdgeBatchUpsertOutcome, StorageError> {
1294 if let Some(writer_task) = self.current_writer_task(operation)? {
1295 return writer_task
1296 .send_bounded_at(
1297 move |conn| {
1298 observed_edge_batch_upsert(conn, &requests, guard_endpoints)
1299 .map_err(|error| map_err(error, operation))
1300 },
1301 operation,
1302 )
1303 .await;
1304 }
1305
1306 self.with_writer(operation, move |conn| {
1307 observed_edge_batch_upsert(conn, &requests, guard_endpoints)
1308 })
1309 .await
1310 }
1311
1312 async fn with_reader<F, R>(&self, op: &'static str, f: F) -> Result<R, StorageError>
1313 where
1314 F: FnOnce(&rusqlite::Connection) -> Result<R, rusqlite::Error> + Send + 'static,
1315 R: Send + 'static,
1316 {
1317 super::run_pooled_store_read(
1318 Arc::clone(&self.pool),
1319 StorageCapability::Graph,
1320 op,
1321 move |conn| f(conn).map_err(|error| map_err(error, op)),
1322 )
1323 .await
1324 }
1325}
1326
1327fn report_graph_usage(queries: &std::sync::atomic::AtomicU64, rows: &std::sync::atomic::AtomicU64) {
1343 khive_storage::usage::count(
1344 khive_storage::usage::UsageUnit::DbRoundTrips,
1345 queries.load(std::sync::atomic::Ordering::Relaxed),
1346 );
1347 khive_storage::usage::count(
1348 khive_storage::usage::UsageUnit::GraphHops,
1349 rows.load(std::sync::atomic::Ordering::Relaxed),
1350 );
1351}
1352
1353fn read_edge(row: &rusqlite::Row<'_>) -> Result<Edge, rusqlite::Error> {
1354 let namespace: String = row.get(0)?;
1355 let id_str: String = row.get(1)?;
1356 let source_str: String = row.get(2)?;
1357 let target_str: String = row.get(3)?;
1358 let relation_str: String = row.get(4)?;
1359 let weight: f64 = row.get(5)?;
1360 let created_micros: i64 = row.get(6)?;
1361 let updated_micros: i64 = row.get(7)?;
1362 let deleted_micros: Option<i64> = row.get(8)?;
1363 let metadata_str: Option<String> = row.get(9)?;
1364 let target_backend: Option<String> = row.get(10)?;
1365
1366 let id = parse_uuid(&id_str)?;
1367 let source_id = parse_uuid(&source_str)?;
1368 let target_id = parse_uuid(&target_str)?;
1369 let created_at = micros_to_datetime(created_micros);
1370 let relation = relation_str.parse::<EdgeRelation>().map_err(|e| {
1371 rusqlite::Error::FromSqlConversionFailure(4, rusqlite::types::Type::Text, Box::new(e))
1372 })?;
1373 let metadata = match metadata_str {
1374 Some(s) => {
1375 let v = serde_json::from_str(&s).map_err(|e| {
1376 rusqlite::Error::FromSqlConversionFailure(
1377 9,
1378 rusqlite::types::Type::Text,
1379 Box::new(e),
1380 )
1381 })?;
1382 Some(v)
1383 }
1384 None => None,
1385 };
1386
1387 Ok(Edge {
1388 id: id.into(),
1389 namespace,
1390 source_id,
1391 target_id,
1392 relation,
1393 weight,
1394 created_at,
1395 updated_at: micros_to_datetime(updated_micros),
1396 deleted_at: deleted_micros.map(micros_to_datetime),
1397 metadata,
1398 target_backend,
1399 })
1400}
1401
1402fn parse_uuid(s: &str) -> Result<Uuid, rusqlite::Error> {
1403 Uuid::parse_str(s).map_err(|e| {
1404 rusqlite::Error::FromSqlConversionFailure(0, rusqlite::types::Type::Text, Box::new(e))
1405 })
1406}
1407
1408fn neighbor_extra_clause(
1414 query: &NeighborQuery,
1415 start_param_idx: usize,
1416 after: Option<&NeighborCursor>,
1417 neighbor_kinds: Option<&[String]>,
1418) -> (String, String, Vec<Box<dyn rusqlite::types::ToSql>>) {
1419 let mut conditions: Vec<String> = Vec::new();
1420 let mut extra_params: Vec<Box<dyn rusqlite::types::ToSql>> = Vec::new();
1421 let mut param_idx = start_param_idx;
1422
1423 if let Some(ref rels) = query.relations {
1424 if !rels.is_empty() {
1425 let placeholders: Vec<String> = rels
1426 .iter()
1427 .map(|r| {
1428 extra_params.push(Box::new(r.to_string()));
1429 let p = format!("?{}", param_idx);
1430 param_idx += 1;
1431 p
1432 })
1433 .collect();
1434 conditions.push(format!("relation IN ({})", placeholders.join(",")));
1435 }
1436 }
1437
1438 if let Some(min_w) = query.min_weight {
1439 extra_params.push(Box::new(min_w));
1440 conditions.push(format!("weight >= ?{}", param_idx));
1441 param_idx += 1;
1442 }
1443
1444 if let Some(cursor) = after {
1445 extra_params.push(Box::new(cursor.weight));
1446 let weight_idx = param_idx;
1447 param_idx += 1;
1448 extra_params.push(Box::new(cursor.node_id.to_string()));
1449 let node_idx = param_idx;
1450 param_idx += 1;
1451 extra_params.push(Box::new(cursor.edge_id.to_string()));
1452 let edge_idx = param_idx;
1453 param_idx += 1;
1454 conditions.push(format!(
1455 "(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}))"
1456 ));
1457 }
1458
1459 if let Some(kinds) = neighbor_kinds.filter(|kinds| !kinds.is_empty()) {
1460 let placeholders: Vec<String> = kinds
1461 .iter()
1462 .map(|kind| {
1463 extra_params.push(Box::new(kind.clone()));
1464 let p = format!("?{param_idx}");
1465 param_idx += 1;
1466 p
1467 })
1468 .collect();
1469 let entity_placeholders = placeholders.join(",");
1470 let note_placeholders: Vec<String> = kinds
1471 .iter()
1472 .map(|kind| {
1473 extra_params.push(Box::new(kind.clone()));
1474 let p = format!("?{param_idx}");
1475 param_idx += 1;
1476 p
1477 })
1478 .collect();
1479 conditions.push(format!(
1480 "(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 ({})))",
1481 note_placeholders.join(",")
1482 ));
1483 }
1484
1485 let where_extra = if conditions.is_empty() {
1486 String::new()
1487 } else {
1488 format!(" WHERE {}", conditions.join(" AND "))
1489 };
1490
1491 let limit_clause = if let Some(lim) = query.limit {
1492 extra_params.push(Box::new(lim as i64));
1493 format!(" LIMIT ?{}", param_idx)
1494 } else {
1495 String::new()
1496 };
1497
1498 (where_extra, limit_clause, extra_params)
1499}
1500
1501#[cfg(test)]
1509static NEIGHBOR_SELECT_COUNT: std::sync::atomic::AtomicUsize =
1510 std::sync::atomic::AtomicUsize::new(0);
1511
1512#[cfg(test)]
1513fn count_neighbor_select() {
1514 NEIGHBOR_SELECT_COUNT.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1515}
1516
1517#[cfg(not(test))]
1518fn count_neighbor_select() {}
1519
1520#[cfg(test)]
1521pub(crate) fn reset_neighbor_select_count() {
1522 NEIGHBOR_SELECT_COUNT.store(0, std::sync::atomic::Ordering::Relaxed);
1523}
1524
1525#[cfg(test)]
1526pub(crate) fn neighbor_select_count() -> usize {
1527 NEIGHBOR_SELECT_COUNT.load(std::sync::atomic::Ordering::Relaxed)
1528}
1529
1530fn micros_to_datetime(micros: i64) -> DateTime<Utc> {
1531 Utc.timestamp_micros(micros)
1532 .single()
1533 .unwrap_or_else(Utc::now)
1534}
1535
1536fn edge_order_clause(sort: &[SortOrder<EdgeSortField>]) -> String {
1543 if sort.is_empty() {
1544 return " ORDER BY created_at DESC, id DESC".to_string();
1545 }
1546 let mut parts: Vec<String> = sort
1547 .iter()
1548 .map(|s| {
1549 let dir = match s.direction {
1550 SortDirection::Asc => "ASC",
1551 SortDirection::Desc => "DESC",
1552 };
1553 format!("{} {}", edge_sort_col(&s.field), dir)
1554 })
1555 .collect();
1556 let dir = match sort.last().map(|s| &s.direction) {
1557 Some(SortDirection::Asc) => "ASC",
1558 _ => "DESC",
1559 };
1560 parts.push(format!("id {dir}"));
1561 format!(" ORDER BY {}", parts.join(", "))
1562}
1563
1564fn endpoint_base_case(column: &str) -> String {
1582 format!(
1583 "CASE WHEN EXISTS (SELECT 1 FROM entities be WHERE be.id = graph_edges.{column}) \
1584 THEN 'entity' \
1585 WHEN EXISTS (SELECT 1 FROM notes bn WHERE bn.id = graph_edges.{column}) \
1586 THEN 'note' ELSE 'none' END"
1587 )
1588}
1589
1590fn fold_endpoint_base_row(counts: &mut EdgeEndpointBaseCounts, source: &str, target: &str, n: u64) {
1594 let slot = match (source, target) {
1595 ("entity", "entity") => &mut counts.entity_entity,
1596 ("entity", "note") => &mut counts.entity_note,
1597 ("note", "entity") => &mut counts.note_entity,
1598 ("note", "note") => &mut counts.note_note,
1599 _ => &mut counts.unresolved,
1600 };
1601 *slot = slot.saturating_add(n);
1602}
1603
1604const LIVE_ENDPOINTS_CONDITION: &str = "NOT EXISTS (SELECT 1 FROM entities le \
1617 WHERE le.id = graph_edges.source_id AND le.deleted_at IS NOT NULL) \
1618 AND NOT EXISTS (SELECT 1 FROM entities le \
1619 WHERE le.id = graph_edges.target_id AND le.deleted_at IS NOT NULL) \
1620 AND NOT EXISTS (SELECT 1 FROM notes ln \
1621 WHERE ln.id = graph_edges.source_id AND ln.deleted_at IS NOT NULL) \
1622 AND NOT EXISTS (SELECT 1 FROM notes ln \
1623 WHERE ln.id = graph_edges.target_id AND ln.deleted_at IS NOT NULL)";
1624
1625fn with_live_endpoints(where_clause: &str) -> String {
1628 format!("{where_clause} AND {LIVE_ENDPOINTS_CONDITION}")
1629}
1630
1631fn build_edge_filter_sql(
1632 namespace: &str,
1633 filter: &EdgeFilter,
1634) -> (String, Vec<Box<dyn rusqlite::types::ToSql>>) {
1635 build_edge_filter_sql_for_namespaces(&[namespace.to_string()], filter)
1636}
1637
1638fn build_edge_filter_sql_for_namespaces(
1639 namespaces: &[String],
1640 filter: &EdgeFilter,
1641) -> (String, Vec<Box<dyn rusqlite::types::ToSql>>) {
1642 let params: Vec<Box<dyn rusqlite::types::ToSql>> = namespaces
1643 .iter()
1644 .map(|namespace| -> Box<dyn rusqlite::types::ToSql> { Box::new(namespace.clone()) })
1645 .collect();
1646 let namespace_condition = match namespaces.len() {
1647 0 => "0".to_string(),
1648 1 => "namespace = ?1".to_string(),
1649 _ => {
1650 let placeholders: Vec<String> =
1651 (1..=namespaces.len()).map(|i| format!("?{i}")).collect();
1652 format!("namespace IN ({})", placeholders.join(", "))
1653 }
1654 };
1655 build_edge_filter_conditions(namespace_condition, params, filter)
1656}
1657
1658fn build_edge_filter_sql_for_namespaces_json(
1673 namespaces_json: &str,
1674 filter: &EdgeFilter,
1675) -> (String, Vec<Box<dyn rusqlite::types::ToSql>>) {
1676 let params: Vec<Box<dyn rusqlite::types::ToSql>> = vec![Box::new(namespaces_json.to_string())];
1677 let namespace_condition = "namespace IN (SELECT value FROM json_each(?1))".to_string();
1678 build_edge_filter_conditions(namespace_condition, params, filter)
1679}
1680
1681fn build_edge_filter_conditions(
1682 namespace_condition: String,
1683 mut params: Vec<Box<dyn rusqlite::types::ToSql>>,
1684 filter: &EdgeFilter,
1685) -> (String, Vec<Box<dyn rusqlite::types::ToSql>>) {
1686 let mut conditions = vec![namespace_condition, "deleted_at IS NULL".to_string()];
1687
1688 if !filter.ids.is_empty() {
1689 let placeholders: Vec<String> = filter
1690 .ids
1691 .iter()
1692 .map(|id| {
1693 params.push(Box::new(id.to_string()));
1694 format!("?{}", params.len())
1695 })
1696 .collect();
1697 conditions.push(format!("id IN ({})", placeholders.join(",")));
1698 }
1699
1700 if !filter.source_ids.is_empty() {
1701 let placeholders: Vec<String> = filter
1702 .source_ids
1703 .iter()
1704 .map(|id| {
1705 params.push(Box::new(id.to_string()));
1706 format!("?{}", params.len())
1707 })
1708 .collect();
1709 conditions.push(format!("source_id IN ({})", placeholders.join(",")));
1710 }
1711
1712 if !filter.target_ids.is_empty() {
1713 let placeholders: Vec<String> = filter
1714 .target_ids
1715 .iter()
1716 .map(|id| {
1717 params.push(Box::new(id.to_string()));
1718 format!("?{}", params.len())
1719 })
1720 .collect();
1721 conditions.push(format!("target_id IN ({})", placeholders.join(",")));
1722 }
1723
1724 if !filter.relations.is_empty() {
1725 let placeholders: Vec<String> = filter
1726 .relations
1727 .iter()
1728 .map(|r| {
1729 params.push(Box::new(r.to_string()));
1730 format!("?{}", params.len())
1731 })
1732 .collect();
1733 conditions.push(format!("relation IN ({})", placeholders.join(",")));
1734 }
1735
1736 if let Some(min_w) = filter.min_weight {
1737 params.push(Box::new(min_w));
1738 conditions.push(format!("weight >= ?{}", params.len()));
1739 }
1740
1741 if let Some(max_w) = filter.max_weight {
1742 params.push(Box::new(max_w));
1743 conditions.push(format!("weight <= ?{}", params.len()));
1744 }
1745
1746 if let Some(ref time_range) = filter.created_at {
1747 if let Some(start) = time_range.start {
1748 params.push(Box::new(start.timestamp_micros()));
1749 conditions.push(format!("created_at >= ?{}", params.len()));
1750 }
1751 if let Some(end) = time_range.end {
1752 params.push(Box::new(end.timestamp_micros()));
1753 conditions.push(format!("created_at < ?{}", params.len()));
1754 }
1755 }
1756
1757 let clause = format!(" WHERE {}", conditions.join(" AND "));
1758 (clause, params)
1759}
1760
1761fn edge_sort_col(field: &EdgeSortField) -> &'static str {
1762 match field {
1763 EdgeSortField::CreatedAt => "created_at",
1764 EdgeSortField::Weight => "weight",
1765 EdgeSortField::Relation => "relation",
1766 }
1767}
1768
1769fn canonical_edge_endpoints(
1778 relation: EdgeRelation,
1779 source_id: Uuid,
1780 target_id: Uuid,
1781) -> (Uuid, Uuid) {
1782 relation.canonical_endpoints(source_id, target_id)
1783}
1784
1785fn edge_endpoints_exist(
1793 conn: &rusqlite::Connection,
1794 source_id: Uuid,
1795 target_id: Uuid,
1796) -> Result<MissingEndpoints, rusqlite::Error> {
1797 let src_exists = endpoint_exists_clause("?1");
1798 let tgt_exists = endpoint_exists_clause("?2");
1799 let sql = format!("SELECT ({src_exists}), ({tgt_exists})");
1800 conn.query_row(
1801 &sql,
1802 rusqlite::params![source_id.to_string(), target_id.to_string()],
1803 |row| {
1804 let src_exists: bool = row.get(0)?;
1805 let tgt_exists: bool = row.get(1)?;
1806 Ok(MissingEndpoints {
1807 source: !src_exists,
1808 target: !tgt_exists,
1809 })
1810 },
1811 )
1812}
1813
1814fn edge_by_natural_key_including_deleted(
1815 conn: &rusqlite::Connection,
1816 edge: &Edge,
1817) -> Result<Option<Edge>, rusqlite::Error> {
1818 edge_by_natural_key_parts_including_deleted(
1819 conn,
1820 &edge.namespace,
1821 edge.source_id,
1822 edge.target_id,
1823 edge.relation,
1824 )
1825}
1826
1827fn edge_by_natural_key_parts_including_deleted(
1828 conn: &rusqlite::Connection,
1829 namespace: &str,
1830 source_id: Uuid,
1831 target_id: Uuid,
1832 relation: EdgeRelation,
1833) -> Result<Option<Edge>, rusqlite::Error> {
1834 let (source_id, target_id) = canonical_edge_endpoints(relation, source_id, target_id);
1835 conn.query_row(
1836 "SELECT namespace, id, source_id, target_id, relation, weight, \
1837 created_at, updated_at, deleted_at, metadata, target_backend \
1838 FROM graph_edges \
1839 WHERE namespace = ?1 AND source_id = ?2 AND target_id = ?3 AND relation = ?4",
1840 rusqlite::params![
1841 namespace,
1842 source_id.to_string(),
1843 target_id.to_string(),
1844 relation.as_str(),
1845 ],
1846 read_edge,
1847 )
1848 .optional()
1849}
1850
1851fn conditional_commit_annotation_insert(
1856 conn: &rusqlite::Connection,
1857 edge: Edge,
1858 guard: &CommitAnnotationGuard,
1859) -> Result<CommitAnnotationInsertOutcome, rusqlite::Error> {
1860 let cursor_matches: bool = conn.query_row(
1864 "SELECT (SELECT EXISTS(SELECT 1 FROM git_mirror_cursor \
1865 WHERE project_id=?1 AND kind='commits' \
1866 AND typeof(cursor_value)='text' \
1867 AND CAST(cursor_value AS BLOB)=?2 AND updated_at=?3)) \
1868 AND (SELECT EXISTS(SELECT 1 FROM git_mirror_cursor \
1869 WHERE project_id=?1 AND kind='commits_checkpoint' \
1870 AND typeof(cursor_value)='text' \
1871 AND CAST(cursor_value AS BLOB)=?4 AND updated_at=?5))",
1872 rusqlite::params![
1873 edge.target_id.to_string(),
1874 guard.commits.value,
1875 guard.commits.updated_at,
1876 guard.checkpoint.value,
1877 guard.checkpoint.updated_at,
1878 ],
1879 |row| row.get(0),
1880 )?;
1881 let metadata = edge
1882 .metadata
1883 .as_ref()
1884 .map(serde_json::to_string)
1885 .transpose()
1886 .map_err(|error| rusqlite::Error::ToSqlConversionFailure(Box::new(error)))?;
1887 let affected = conn.execute(
1888 include_str!("../../sql/commit-annotation-insert.sql"),
1889 rusqlite::params![
1890 edge.namespace,
1891 Uuid::from(edge.id).to_string(),
1892 edge.source_id.to_string(),
1893 edge.target_id.to_string(),
1894 edge.weight,
1895 edge.created_at.timestamp_micros(),
1896 edge.updated_at.timestamp_micros(),
1897 metadata,
1898 guard.expected_sha,
1899 guard.source_identity,
1900 cursor_matches,
1901 ],
1902 )?;
1903 if affected == 1 {
1904 return Ok(CommitAnnotationInsertOutcome::Created(edge));
1905 }
1906 let source_live: bool = conn.query_row(
1907 include_str!("../../sql/commit-annotation-source-live-select.sql"),
1908 rusqlite::params![
1909 edge.source_id.to_string(),
1910 edge.namespace,
1911 guard.expected_sha
1912 ],
1913 |row| row.get(0),
1914 )?;
1915 if !source_live {
1916 return Ok(CommitAnnotationInsertOutcome::SourceChanged);
1917 }
1918 let target_live: bool = conn.query_row(
1919 "SELECT EXISTS(SELECT 1 FROM entities WHERE id=?1 AND namespace=?2 \
1920 AND kind='project' AND deleted_at IS NULL \
1921 AND json_type(properties, '$.repo_slug')='text' \
1922 AND json_extract(properties, '$.repo_slug')=?3 COLLATE BINARY)",
1923 rusqlite::params![
1924 edge.target_id.to_string(),
1925 edge.namespace,
1926 guard.source_identity
1927 ],
1928 |row| row.get(0),
1929 )?;
1930 if !target_live {
1931 return Ok(CommitAnnotationInsertOutcome::TargetChanged);
1932 }
1933 if !cursor_matches {
1934 return Ok(CommitAnnotationInsertOutcome::CursorChanged);
1935 }
1936 match edge_by_natural_key_including_deleted(conn, &edge)? {
1937 Some(existing) if existing.deleted_at.is_some() => {
1938 Ok(CommitAnnotationInsertOutcome::Tombstoned)
1939 }
1940 Some(_) => Ok(CommitAnnotationInsertOutcome::ExistingLive),
1941 None => Err(rusqlite::Error::QueryReturnedNoRows),
1942 }
1943}
1944
1945fn observed_edge_upsert(
1948 conn: &rusqlite::Connection,
1949 request: &EdgeUpsertRequest,
1950 guard_endpoints: bool,
1951) -> Result<GuardedEdgeUpsertOutcome, rusqlite::Error> {
1952 let (source_id, target_id) = canonical_edge_endpoints(
1953 request.edge.relation,
1954 request.edge.source_id,
1955 request.edge.target_id,
1956 );
1957 if guard_endpoints {
1958 #[cfg(test)]
1961 tests::insert_probe_seam::hook((source_id, target_id));
1962 let missing = edge_endpoints_exist(conn, source_id, target_id)?;
1963 if missing.any() {
1964 return Ok(GuardedEdgeUpsertOutcome::Refused(
1965 EdgeUpsertRefusal::MissingEndpoints(missing),
1966 ));
1967 }
1968 }
1969
1970 let previous = edge_by_natural_key_including_deleted(conn, &request.edge)?;
1971 if let Some(edge) = previous.as_ref() {
1972 if edge.deleted_at.is_some() && !request.resurrect {
1973 return Ok(GuardedEdgeUpsertOutcome::Refused(
1974 EdgeUpsertRefusal::ResurrectionRequired { edge: edge.clone() },
1975 ));
1976 }
1977 }
1978
1979 let statement = edge_upsert_statement_with_resurrection(&request.edge, request.resurrect);
1980 let mut stmt = conn.prepare(&statement.sql)?;
1981 bind_params(&mut stmt, &statement.params)?;
1982 let affected = stmt.raw_execute()?;
1983 if affected == 0 {
1984 let edge = edge_by_natural_key_including_deleted(conn, &request.edge)?
1985 .ok_or(rusqlite::Error::QueryReturnedNoRows)?;
1986 return Ok(GuardedEdgeUpsertOutcome::Refused(
1987 EdgeUpsertRefusal::ResurrectionRequired { edge },
1988 ));
1989 }
1990
1991 let edge = edge_by_natural_key_including_deleted(conn, &request.edge)?
1992 .ok_or(rusqlite::Error::QueryReturnedNoRows)?;
1993 let disposition = match previous.as_ref().and_then(|edge| edge.deleted_at) {
1994 None if previous.is_none() => EdgeUpsertDisposition::Created,
1995 None => EdgeUpsertDisposition::Updated,
1996 Some(_) => EdgeUpsertDisposition::Resurrected,
1997 };
1998 Ok(GuardedEdgeUpsertOutcome::Written(EdgeUpsertResult {
1999 edge,
2000 disposition,
2001 previous,
2002 }))
2003}
2004
2005fn observed_edge_batch_upsert(
2008 conn: &rusqlite::Connection,
2009 requests: &[EdgeUpsertRequest],
2010 guard_endpoints: bool,
2011) -> Result<GuardedEdgeBatchUpsertOutcome, rusqlite::Error> {
2012 for (entry_index, request) in requests.iter().enumerate() {
2016 let (source_id, target_id) = canonical_edge_endpoints(
2017 request.edge.relation,
2018 request.edge.source_id,
2019 request.edge.target_id,
2020 );
2021 if guard_endpoints {
2022 let missing = edge_endpoints_exist(conn, source_id, target_id)?;
2023 if missing.any() {
2024 return Ok(GuardedEdgeBatchUpsertOutcome {
2025 rows: Vec::new(),
2026 refusal: Some(GuardedEdgeBatchRefusal {
2027 entry_index,
2028 reason: EdgeUpsertRefusal::MissingEndpoints(missing),
2029 }),
2030 });
2031 }
2032 }
2033 if let Some(edge) = edge_by_natural_key_including_deleted(conn, &request.edge)? {
2034 if edge.deleted_at.is_some() && !request.resurrect {
2035 return Ok(GuardedEdgeBatchUpsertOutcome {
2036 rows: Vec::new(),
2037 refusal: Some(GuardedEdgeBatchRefusal {
2038 entry_index,
2039 reason: EdgeUpsertRefusal::ResurrectionRequired { edge },
2040 }),
2041 });
2042 }
2043 }
2044 }
2045
2046 let mut rows = Vec::with_capacity(requests.len());
2047 for request in requests {
2048 match observed_edge_upsert(conn, request, false)? {
2049 GuardedEdgeUpsertOutcome::Written(row) => rows.push(row),
2050 GuardedEdgeUpsertOutcome::Refused(_) => {
2051 return Err(rusqlite::Error::ExecuteReturnedResults)
2052 }
2053 }
2054 }
2055 Ok(GuardedEdgeBatchUpsertOutcome {
2056 rows,
2057 refusal: None,
2058 })
2059}
2060
2061fn traversal_neighbor_sql(
2066 direction: Direction,
2067 relation_count: usize,
2068 has_min_weight: bool,
2069) -> String {
2070 let (node_column, endpoint_column, index) = match direction {
2071 Direction::Out => ("target_id", "source_id", "idx_graph_edges_ns_src_rel"),
2072 Direction::In => ("source_id", "target_id", "idx_graph_edges_ns_tgt_rel"),
2073 Direction::Both => unreachable!("Direction::Both is split into indexed Out/In seeks"),
2074 };
2075 let mut sql = format!(
2076 "SELECT {node_column}, id, weight \
2077 FROM graph_edges INDEXED BY {index} \
2078 WHERE namespace = ?1 AND {endpoint_column} = ?2 AND deleted_at IS NULL"
2079 );
2080 if relation_count > 0 {
2081 let placeholders = (0..relation_count)
2082 .map(|offset| format!("?{}", 4 + offset))
2083 .collect::<Vec<_>>()
2084 .join(",");
2085 sql.push_str(&format!(" AND relation IN ({placeholders})"));
2086 }
2087 if has_min_weight {
2088 sql.push_str(&format!(" AND weight >= ?{}", 4 + relation_count));
2089 }
2090 sql.push_str(" LIMIT ?3");
2091 sql
2092}
2093
2094fn traversal_timeout_error(budget: &TraversalExecutionBudget) -> StorageError {
2095 StorageError::Timeout {
2096 operation: format!(
2097 "traverse ({}ms execution budget)",
2098 budget.max_duration().as_millis()
2099 )
2100 .into(),
2101 }
2102}
2103
2104fn traversal_work_error(budget: &TraversalExecutionBudget) -> StorageError {
2105 StorageError::InvalidInput {
2106 capability: StorageCapability::Graph,
2107 operation: "traverse".into(),
2108 message: format!(
2109 "traversal work budget exceeded after {} adjacency rows; \
2110 narrow roots, depth, relations, or result limit",
2111 budget.work_limit()
2112 ),
2113 }
2114}
2115
2116#[derive(Clone, Copy)]
2117struct TraversalFrontierNode {
2118 node_id: Uuid,
2119 depth: usize,
2120 total_weight: f64,
2121}
2122
2123#[allow(clippy::too_many_arguments)]
2124fn run_bounded_traversal(
2125 conn: &rusqlite::Connection,
2126 roots: Vec<Uuid>,
2127 opts: TraversalOptions,
2128 include_roots: bool,
2129 namespace: String,
2130 origin: khive_storage::tx_registry::TxOrigin,
2131 budget: TraversalExecutionBudget,
2132 counted_rows: &std::sync::atomic::AtomicU64,
2133 counted_queries: &std::sync::atomic::AtomicU64,
2134) -> Result<Vec<GraphPath>, StorageError> {
2135 let progress_timed_out = Arc::new(std::sync::atomic::AtomicBool::new(false));
2136 let callback_timed_out = Arc::clone(&progress_timed_out);
2137 let callback_budget = budget.clone();
2138 #[cfg(test)]
2139 let progress_seam_root = roots.first().copied();
2140 conn.progress_handler(
2141 1_000,
2142 Some(move || {
2143 if crate::read_cancellation::current_read_should_interrupt() {
2144 return true;
2145 }
2146 #[cfg(test)]
2147 if tests::traverse_progress_seam::hook(progress_seam_root) {
2148 callback_timed_out.store(true, std::sync::atomic::Ordering::Relaxed);
2149 return true;
2150 }
2151 let expired = callback_budget.is_expired();
2152 if expired {
2153 callback_timed_out.store(true, std::sync::atomic::Ordering::Relaxed);
2154 }
2155 expired
2156 }),
2157 )
2158 .map_err(|e| map_err(e, "traverse_progress_handler"))?;
2159
2160 let result = (|| {
2161 let result_limit = opts.effective_limit() as usize;
2162 let relation_count = opts.relations.as_ref().map_or(0, Vec::len);
2163 let directions = match opts.direction {
2164 Direction::Out => vec![Direction::Out],
2165 Direction::In => vec![Direction::In],
2166 Direction::Both => vec![Direction::Out, Direction::In],
2167 };
2168 let statements = directions
2169 .into_iter()
2170 .map(|direction| {
2171 traversal_neighbor_sql(direction, relation_count, opts.min_weight.is_some())
2172 })
2173 .collect::<Vec<_>>();
2174 let map_sql_error = |error| {
2175 if progress_timed_out.load(std::sync::atomic::Ordering::Relaxed) {
2176 traversal_timeout_error(&budget)
2177 } else {
2178 map_err(error, "traverse")
2179 }
2180 };
2181
2182 let mut all_paths = Vec::with_capacity(roots.len());
2183 for root_id in roots {
2184 let mut seen = HashSet::new();
2185 seen.insert(root_id);
2186 let mut frontier = VecDeque::from([TraversalFrontierNode {
2187 node_id: root_id,
2188 depth: 0,
2189 total_weight: 0.0,
2190 }]);
2191 let mut nodes = Vec::with_capacity(result_limit + usize::from(include_roots));
2192 if include_roots {
2193 nodes.push(PathNode {
2194 node_id: root_id,
2195 via_edge: None,
2196 depth: 0,
2197 name: None,
2198 kind: None,
2199 properties: None,
2200 weight: 0.0,
2201 });
2202 }
2203 let mut non_root_count = 0usize;
2204
2205 'root_walk: while non_root_count < result_limit {
2206 let Some(current) = frontier.pop_front() else {
2207 break;
2208 };
2209 if current.depth >= opts.max_depth {
2210 continue;
2211 }
2212 if budget.is_expired() {
2213 return Err(traversal_timeout_error(&budget));
2214 }
2215
2216 for sql in &statements {
2217 let row_cap = budget.remaining_work().saturating_add(1);
2218 let mut params: Vec<Box<dyn rusqlite::types::ToSql>> = vec![
2219 Box::new(namespace.clone()),
2220 Box::new(current.node_id.to_string()),
2221 Box::new(row_cap as i64),
2222 ];
2223 if let Some(relations) = &opts.relations {
2224 params.extend(relations.iter().map(|relation| {
2225 Box::new(relation.to_string()) as Box<dyn rusqlite::types::ToSql>
2226 }));
2227 }
2228 if let Some(min_weight) = opts.min_weight {
2229 params.push(Box::new(min_weight));
2230 }
2231 let param_refs = params
2232 .iter()
2233 .map(|param| param.as_ref())
2234 .collect::<Vec<&dyn rusqlite::types::ToSql>>();
2235
2236 let _snapshot = khive_storage::tx_registry::register_scoped(
2237 Some("graph_traverse_read".to_string()),
2238 origin.clone(),
2239 );
2240 let mut stmt = conn.prepare(sql).map_err(&map_sql_error)?;
2241 counted_queries.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2242 let mut rows = stmt.query(param_refs.as_slice()).map_err(&map_sql_error)?;
2243 while let Some(row) = rows.next().map_err(&map_sql_error)? {
2244 #[cfg(test)]
2247 tests::traverse_snapshot_seam::hook(current.node_id);
2248 counted_rows.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2249 if budget.is_expired() {
2250 return Err(traversal_timeout_error(&budget));
2251 }
2252 if !budget.try_consume_row() {
2253 return Err(traversal_work_error(&budget));
2254 }
2255 let node_str: String = row.get(0).map_err(&map_sql_error)?;
2256 let edge_str: String = row.get(1).map_err(&map_sql_error)?;
2257 let edge_weight: f64 = row.get(2).map_err(&map_sql_error)?;
2258 let node_id = parse_uuid(&node_str).map_err(&map_sql_error)?;
2259 if !seen.insert(node_id) {
2260 continue;
2261 }
2262 let via_edge = parse_uuid(&edge_str).map_err(&map_sql_error)?;
2263 let depth = current.depth + 1;
2264 let total_weight = current.total_weight + edge_weight;
2265 nodes.push(PathNode {
2266 node_id,
2267 via_edge: Some(via_edge),
2268 depth,
2269 name: None,
2270 kind: None,
2271 properties: None,
2272 weight: total_weight,
2273 });
2274 non_root_count += 1;
2275 if depth < opts.max_depth {
2276 frontier.push_back(TraversalFrontierNode {
2277 node_id,
2278 depth,
2279 total_weight,
2280 });
2281 }
2282 if non_root_count == result_limit {
2283 break 'root_walk;
2284 }
2285 }
2286 }
2287 }
2288
2289 if !nodes.is_empty() {
2290 let total_weight = nodes.iter().map(|node| node.weight).fold(0.0_f64, f64::max);
2291 all_paths.push(GraphPath {
2292 root_id,
2293 nodes,
2294 total_weight,
2295 });
2296 }
2297 }
2298 Ok(all_paths)
2299 })();
2300
2301 conn.progress_handler(0, None::<fn() -> bool>)
2302 .map_err(|e| map_err(e, "traverse_progress_handler_clear"))?;
2303 result
2304}
2305
2306impl SqlGraphStore {
2307 async fn query_neighbors_page(
2308 &self,
2309 operation: &'static str,
2310 node_id: Uuid,
2311 query: NeighborQuery,
2312 after: Option<NeighborCursor>,
2313 neighbor_kinds: Option<Vec<String>>,
2314 ) -> Result<Vec<NeighborHit>, StorageError> {
2315 count_neighbor_select();
2316
2317 let namespace = self.namespace.clone();
2318 let node_str = node_id.to_string();
2319 let counted_queries = Arc::new(std::sync::atomic::AtomicU64::new(0));
2320 let counted_rows = Arc::new(std::sync::atomic::AtomicU64::new(0));
2321 let closure_queries = Arc::clone(&counted_queries);
2322 let closure_rows = Arc::clone(&counted_rows);
2323 let result = self
2324 .with_reader(operation, move |conn| {
2325 let base_out = "SELECT target_id AS node_id, id AS edge_id, relation, weight \
2326 FROM graph_edges \
2327 WHERE namespace = ?1 AND source_id = ?2 AND deleted_at IS NULL";
2328 let base_in = "SELECT source_id AS node_id, id AS edge_id, relation, weight \
2329 FROM graph_edges \
2330 WHERE namespace = ?1 AND target_id = ?2 AND deleted_at IS NULL";
2331 let sql = match query.direction {
2332 Direction::Out => base_out.to_string(),
2333 Direction::In => base_in.to_string(),
2334 Direction::Both => format!("{} UNION ALL {}", base_out, base_in),
2335 };
2336 let (where_extra, limit_clause, extra_params) =
2337 neighbor_extra_clause(&query, 3, after.as_ref(), neighbor_kinds.as_deref());
2338 let full_sql = format!(
2339 "SELECT node_id, edge_id, relation, weight FROM ({}){} \
2340 ORDER BY weight DESC, node_id ASC, edge_id ASC{}",
2341 sql, where_extra, limit_clause
2342 );
2343 let mut stmt = conn.prepare(&full_sql)?;
2344 let mut all_params: Vec<Box<dyn rusqlite::types::ToSql>> = Vec::new();
2345 all_params.push(Box::new(namespace.clone()));
2346 all_params.push(Box::new(node_str.clone()));
2347 all_params.extend(extra_params);
2348 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
2349 all_params.iter().map(|p| p.as_ref()).collect();
2350
2351 closure_queries.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2352 let rows = stmt.query_map(param_refs.as_slice(), |row| {
2353 let nid_str: String = row.get(0)?;
2354 let eid_str: String = row.get(1)?;
2355 let relation_str: String = row.get(2)?;
2356 let weight: f64 = row.get(3)?;
2357 Ok((nid_str, eid_str, relation_str, weight))
2358 })?;
2359 let mut hits = Vec::new();
2360 for row in rows {
2361 let (nid_str, eid_str, relation_str, weight) = row?;
2362 closure_rows.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2363 let relation = relation_str.parse::<EdgeRelation>().map_err(|e| {
2364 rusqlite::Error::FromSqlConversionFailure(
2365 2,
2366 rusqlite::types::Type::Text,
2367 Box::new(e),
2368 )
2369 })?;
2370 hits.push(NeighborHit {
2371 node_id: parse_uuid(&nid_str)?,
2372 edge_id: parse_uuid(&eid_str)?,
2373 relation,
2374 weight,
2375 name: None,
2376 kind: None,
2377 entity_type: None,
2378 });
2379 }
2380 Ok(hits)
2381 })
2382 .await;
2383
2384 report_graph_usage(&counted_queries, &counted_rows);
2385 result
2386 }
2387}
2388
2389#[async_trait]
2390impl GraphStore for SqlGraphStore {
2391 async fn update_symmetric_edge_if_unchanged(
2392 &self,
2393 request: SymmetricEdgeUpdateRequest,
2394 ) -> StorageResult<SymmetricEdgeUpdateOutcome> {
2395 self.update_symmetric_edge(request).await
2396 }
2397
2398 async fn latest_annotating_note(
2399 &self,
2400 node_id: Uuid,
2401 kind: &str,
2402 tag: &str,
2403 ) -> Result<Option<(Uuid, i64)>, StorageError> {
2404 let namespace = self.namespace.clone();
2405 let node_id = node_id.to_string();
2406 let kind = kind.to_owned();
2407 let tag = tag.to_owned();
2408 self.with_indexed_reader("latest_annotating_note", move |conn| {
2409 conn.query_row(
2410 LATEST_ANNOTATING_NOTE_SQL,
2411 rusqlite::params![namespace, node_id, kind, tag],
2412 |row| {
2413 let id: String = row.get(0)?;
2414 Ok((parse_uuid(&id)?, row.get(1)?))
2415 },
2416 )
2417 .optional()
2418 })
2419 .await
2420 }
2421
2422 async fn latest_annotating_note_with_property(
2423 &self,
2424 node_id: Uuid,
2425 kind: &str,
2426 tag: &str,
2427 property_key: &str,
2428 property_value: &str,
2429 ) -> Result<Option<(Uuid, i64)>, StorageError> {
2430 let namespace = self.namespace.clone();
2431 let node_id = node_id.to_string();
2432 let kind = kind.to_owned();
2433 let tag = tag.to_owned();
2434 let property_key = property_key.to_owned();
2435 let property_value = property_value.to_owned();
2436 self.with_indexed_reader("latest_annotating_note_with_property", move |conn| {
2437 conn.query_row(
2438 LATEST_ANNOTATING_NOTE_WITH_PROPERTY_SQL,
2439 rusqlite::params![namespace, node_id, kind, tag, property_key, property_value],
2440 |row| {
2441 let id: String = row.get(0)?;
2442 Ok((parse_uuid(&id)?, row.get(1)?))
2443 },
2444 )
2445 .optional()
2446 })
2447 .await
2448 }
2449
2450 async fn upsert_edge(&self, edge: Edge) -> Result<(), StorageError> {
2451 self.upsert_edge_observed(EdgeUpsertRequest {
2452 edge,
2453 resurrect: false,
2454 })
2455 .await
2456 .map(|_| ())
2457 }
2458
2459 async fn upsert_edge_observed(
2460 &self,
2461 request: EdgeUpsertRequest,
2462 ) -> Result<EdgeUpsertResult, StorageError> {
2463 match self
2464 .observed_edge_write("upsert_edge_observed", request, false)
2465 .await?
2466 {
2467 GuardedEdgeUpsertOutcome::Written(result) => Ok(result),
2468 GuardedEdgeUpsertOutcome::Refused(EdgeUpsertRefusal::ResurrectionRequired { edge }) => {
2469 Err(resurrection_required_error("upsert_edge_observed", &edge))
2470 }
2471 GuardedEdgeUpsertOutcome::Refused(EdgeUpsertRefusal::MissingEndpoints(_)) => {
2472 Err(StorageError::Conflict {
2473 capability: StorageCapability::Graph,
2474 operation: "upsert_edge_observed".into(),
2475 message: "unguarded edge upsert reported a missing-endpoint refusal".into(),
2476 })
2477 }
2478 }
2479 }
2480
2481 async fn insert_edge_if_absent(&self, edge: Edge) -> Result<bool, StorageError> {
2482 let statement = edge_insert_if_absent_statement(&edge);
2483 self.with_writer("insert_edge_if_absent", move |conn| {
2484 let mut stmt = conn.prepare(&statement.sql)?;
2485 bind_params(&mut stmt, &statement.params)?;
2486 Ok(stmt.raw_execute()? > 0)
2487 })
2488 .await
2489 }
2490
2491 async fn insert_commit_annotation_if_absent(
2492 &self,
2493 edge: Edge,
2494 guard: CommitAnnotationGuard,
2495 ) -> Result<CommitAnnotationInsertOutcome, StorageError> {
2496 const OP: &str = "insert_commit_annotation_if_absent";
2497 if edge.relation != EdgeRelation::Annotates || edge.deleted_at.is_some() {
2498 return Err(StorageError::InvalidInput {
2499 capability: StorageCapability::Graph,
2500 operation: OP.into(),
2501 message: "expected a live annotates edge".into(),
2502 });
2503 }
2504 if let Some(writer_task) = self.current_writer_task(OP)? {
2505 return writer_task
2506 .send_bounded_at(
2507 move |conn| {
2508 conditional_commit_annotation_insert(conn, edge, &guard)
2509 .map_err(|error| map_err(error, OP))
2510 },
2511 OP,
2512 )
2513 .await;
2514 }
2515 let origin = self.pool.origin();
2516 self.with_writer(OP, move |conn| {
2517 conn.execute_batch("BEGIN IMMEDIATE")?;
2518 let _tx_handle =
2519 khive_storage::tx_registry::register_scoped(Some(OP.to_string()), origin);
2520 let outcome = match conditional_commit_annotation_insert(conn, edge, &guard) {
2521 Ok(outcome) => outcome,
2522 Err(error) => {
2523 let _ = conn.execute_batch("ROLLBACK");
2524 return Err(error);
2525 }
2526 };
2527 if let Err(error) = conn.execute_batch("COMMIT") {
2528 let _ = conn.execute_batch("ROLLBACK");
2529 return Err(error);
2530 }
2531 Ok(outcome)
2532 })
2533 .await
2534 }
2535
2536 async fn replace_edge_if_unchanged(
2537 &self,
2538 edge: Edge,
2539 expected_updated_at: DateTime<Utc>,
2540 expected_deleted_at: Option<DateTime<Utc>>,
2541 ) -> Result<bool, StorageError> {
2542 let statement =
2543 edge_replace_if_unchanged_statement(&edge, expected_updated_at, expected_deleted_at);
2544 self.with_writer("replace_edge_if_unchanged", move |conn| {
2545 let mut stmt = conn.prepare(&statement.sql)?;
2546 bind_params(&mut stmt, &statement.params)?;
2547 Ok(stmt.raw_execute()? > 0)
2548 })
2549 .await
2550 }
2551
2552 async fn upsert_edges(&self, edges: Vec<Edge>) -> Result<BatchWriteSummary, StorageError> {
2553 let attempted = edges.len() as u64;
2554 let requests = edges
2555 .into_iter()
2556 .map(|edge| EdgeUpsertRequest {
2557 edge,
2558 resurrect: false,
2559 })
2560 .collect();
2561 let outcome = self
2562 .observed_edge_batch_write("upsert_edges", requests, false)
2563 .await?;
2564 if let Some(refusal) = outcome.refusal {
2565 return match refusal.reason {
2566 EdgeUpsertRefusal::ResurrectionRequired { edge } => {
2567 Err(resurrection_required_error("upsert_edges", &edge))
2568 }
2569 EdgeUpsertRefusal::MissingEndpoints(_) => Err(StorageError::Conflict {
2570 capability: StorageCapability::Graph,
2571 operation: "upsert_edges".into(),
2572 message: "unguarded edge batch reported a missing-endpoint refusal".into(),
2573 }),
2574 };
2575 }
2576 Ok(BatchWriteSummary {
2577 attempted,
2578 affected: outcome.rows.len() as u64,
2579 ..BatchWriteSummary::default()
2580 })
2581 }
2582
2583 async fn upsert_edge_guarded(&self, edge: Edge) -> Result<GuardedWriteOutcome, StorageError> {
2584 match self
2585 .upsert_edge_guarded_observed(EdgeUpsertRequest {
2586 edge,
2587 resurrect: false,
2588 })
2589 .await?
2590 {
2591 GuardedEdgeUpsertOutcome::Written(_) => Ok(GuardedWriteOutcome::Written),
2592 GuardedEdgeUpsertOutcome::Refused(EdgeUpsertRefusal::MissingEndpoints(missing)) => {
2593 Ok(GuardedWriteOutcome::Refused(missing))
2594 }
2595 GuardedEdgeUpsertOutcome::Refused(EdgeUpsertRefusal::ResurrectionRequired { edge }) => {
2596 Err(resurrection_required_error("upsert_edge_guarded", &edge))
2597 }
2598 }
2599 }
2600
2601 async fn upsert_edge_guarded_observed(
2602 &self,
2603 request: EdgeUpsertRequest,
2604 ) -> Result<GuardedEdgeUpsertOutcome, StorageError> {
2605 self.observed_edge_write("upsert_edge_guarded_observed", request, true)
2606 .await
2607 }
2608
2609 async fn upsert_edges_guarded(
2610 &self,
2611 edges: Vec<Edge>,
2612 ) -> Result<GuardedBatchOutcome, StorageError> {
2613 let attempted = edges.len() as u64;
2614 let requests = edges
2615 .iter()
2616 .cloned()
2617 .map(|edge| EdgeUpsertRequest {
2618 edge,
2619 resurrect: false,
2620 })
2621 .collect();
2622 let outcome = self.upsert_edges_guarded_observed(requests).await?;
2623 match outcome.refusal {
2624 None => Ok(GuardedBatchOutcome {
2625 summary: BatchWriteSummary {
2626 attempted,
2627 affected: outcome.rows.len() as u64,
2628 ..BatchWriteSummary::default()
2629 },
2630 refused: None,
2631 }),
2632 Some(refusal) => match refusal.reason {
2633 EdgeUpsertRefusal::MissingEndpoints(missing) => {
2634 let index = refusal.entry_index;
2635 let edge = &edges[index];
2636 let (source_id, target_id) =
2637 canonical_edge_endpoints(edge.relation, edge.source_id, edge.target_id);
2638 let message = format!(
2639 "batch entry {index}: edge endpoint no longer exists at write time: source \
2640 {source_id} or target {target_id}"
2641 );
2642 let mut summary = BatchWriteSummary {
2643 attempted,
2644 ..BatchWriteSummary::default()
2645 };
2646 summary.first_error = message.clone();
2649 let refusal = GuardedBatchRefusal {
2650 entry_index: index,
2651 missing,
2652 };
2653 for (failed_index, failed_edge) in edges.iter().enumerate() {
2654 refusal.record_failure(&mut summary, failed_index, failed_edge, &message);
2655 }
2656 Ok(GuardedBatchOutcome {
2657 summary,
2658 refused: Some(refusal),
2659 })
2660 }
2661 EdgeUpsertRefusal::ResurrectionRequired { edge } => {
2662 Err(resurrection_required_error("upsert_edges_guarded", &edge))
2663 }
2664 },
2665 }
2666 }
2667
2668 async fn upsert_edges_guarded_observed(
2669 &self,
2670 requests: Vec<EdgeUpsertRequest>,
2671 ) -> Result<GuardedEdgeBatchUpsertOutcome, StorageError> {
2672 self.observed_edge_batch_write("upsert_edges_guarded_observed", requests, true)
2673 .await
2674 }
2675
2676 async fn get_edge(&self, id: LinkId) -> Result<Option<Edge>, StorageError> {
2677 let id_str = Uuid::from(id).to_string();
2678
2679 self.with_reader("get_edge", move |conn| {
2680 let mut stmt = conn.prepare(
2681 "SELECT namespace, id, source_id, target_id, relation, weight, \
2682 created_at, updated_at, deleted_at, metadata, target_backend \
2683 FROM graph_edges WHERE id = ?1 AND deleted_at IS NULL",
2684 )?;
2685 let mut rows = stmt.query(rusqlite::params![id_str])?;
2686 match rows.next()? {
2687 Some(row) => Ok(Some(read_edge(row)?)),
2688 None => Ok(None),
2689 }
2690 })
2691 .await
2692 }
2693
2694 async fn get_edge_including_deleted(&self, id: LinkId) -> Result<Option<Edge>, StorageError> {
2695 let id_str = Uuid::from(id).to_string();
2696
2697 self.with_reader("get_edge_including_deleted", move |conn| {
2698 let mut stmt = conn.prepare(
2699 "SELECT namespace, id, source_id, target_id, relation, weight, \
2700 created_at, updated_at, deleted_at, metadata, target_backend \
2701 FROM graph_edges WHERE id = ?1",
2702 )?;
2703 let mut rows = stmt.query(rusqlite::params![id_str])?;
2704 match rows.next()? {
2705 Some(row) => Ok(Some(read_edge(row)?)),
2706 None => Ok(None),
2707 }
2708 })
2709 .await
2710 }
2711
2712 async fn edge_sequence(&self, id: Uuid) -> Result<Option<i64>, StorageError> {
2713 let id = id.to_string();
2714 self.with_reader("edge_sequence", move |conn| {
2715 conn.query_row(
2716 "SELECT seq FROM graph_edges_seq WHERE edge_id = ?1",
2717 rusqlite::params![id],
2718 |row| row.get(0),
2719 )
2720 .optional()
2721 })
2722 .await
2723 }
2724
2725 async fn edge_sequences(&self, ids: &[Uuid]) -> Result<Vec<(Uuid, i64)>, StorageError> {
2726 if ids.is_empty() {
2727 return Ok(Vec::new());
2728 }
2729 let ids = ids.to_vec();
2730 self.with_reader("edge_sequences", move |conn| {
2731 const CHUNK: usize = 900;
2732 let mut resolved = Vec::with_capacity(ids.len());
2733 for chunk in ids.chunks(CHUNK) {
2734 let placeholders = (1..=chunk.len())
2735 .map(|index| format!("?{index}"))
2736 .collect::<Vec<_>>()
2737 .join(", ");
2738 let sql = format!(
2739 "SELECT edge_id, seq FROM graph_edges_seq WHERE edge_id IN ({placeholders})"
2740 );
2741 let strings = chunk.iter().map(Uuid::to_string).collect::<Vec<_>>();
2742 let params = strings
2743 .iter()
2744 .map(|id| id as &dyn rusqlite::types::ToSql)
2745 .collect::<Vec<_>>();
2746 let mut stmt = conn.prepare(&sql)?;
2747 let rows = stmt.query_map(params.as_slice(), |row| {
2748 let id: String = row.get(0)?;
2749 Ok((parse_uuid(&id)?, row.get::<_, i64>(1)?))
2750 })?;
2751 resolved.extend(rows.collect::<Result<Vec<_>, _>>()?);
2752 }
2753 Ok(resolved)
2754 })
2755 .await
2756 }
2757
2758 async fn get_edge_by_natural_key_including_deleted(
2759 &self,
2760 namespace: &str,
2761 source_id: Uuid,
2762 target_id: Uuid,
2763 relation: EdgeRelation,
2764 ) -> Result<Option<Edge>, StorageError> {
2765 let namespace = namespace.to_string();
2766 let source_str = source_id.to_string();
2767 let target_str = target_id.to_string();
2768 let relation_str = relation.to_string();
2769
2770 self.with_reader("get_edge_by_natural_key_including_deleted", move |conn| {
2771 let mut stmt = conn.prepare(
2772 "SELECT namespace, id, source_id, target_id, relation, weight, \
2773 created_at, updated_at, deleted_at, metadata, target_backend \
2774 FROM graph_edges \
2775 WHERE namespace = ?1 AND source_id = ?2 AND target_id = ?3 AND relation = ?4",
2776 )?;
2777 let mut rows = stmt.query(rusqlite::params![
2778 namespace,
2779 source_str,
2780 target_str,
2781 relation_str
2782 ])?;
2783 match rows.next()? {
2784 Some(row) => Ok(Some(read_edge(row)?)),
2785 None => Ok(None),
2786 }
2787 })
2788 .await
2789 }
2790
2791 async fn get_edges(&self, ids: &[LinkId]) -> Result<Vec<Edge>, StorageError> {
2792 if ids.is_empty() {
2793 return Ok(Vec::new());
2794 }
2795 const CHUNK: usize = 900;
2797 let id_strs: Vec<String> = ids.iter().map(|id| Uuid::from(*id).to_string()).collect();
2798
2799 let mut result: Vec<Edge> = Vec::with_capacity(ids.len());
2800 for chunk in id_strs.chunks(CHUNK) {
2801 let chunk_owned: Vec<String> = chunk.to_vec();
2802 let edges = self
2803 .with_reader("get_edges", move |conn| {
2804 let placeholders: Vec<String> =
2805 (1..=chunk_owned.len()).map(|i| format!("?{}", i)).collect();
2806 let sql = format!(
2807 "SELECT namespace, id, source_id, target_id, relation, weight, \
2808 created_at, updated_at, deleted_at, metadata, target_backend \
2809 FROM graph_edges WHERE id IN ({}) AND deleted_at IS NULL",
2810 placeholders.join(",")
2811 );
2812 let mut stmt = conn.prepare(&sql)?;
2813 let params: Vec<&dyn rusqlite::types::ToSql> = chunk_owned
2814 .iter()
2815 .map(|s| s as &dyn rusqlite::types::ToSql)
2816 .collect();
2817 let rows = stmt.query_map(params.as_slice(), read_edge)?;
2818 let mut edges = Vec::new();
2819 for row in rows {
2820 edges.push(row?);
2821 }
2822 Ok(edges)
2823 })
2824 .await?;
2825 result.extend(edges);
2826 }
2827 Ok(result)
2828 }
2829
2830 async fn get_edge_read_outcomes(
2831 &self,
2832 ids: &[LinkId],
2833 ) -> Result<Vec<Result<Option<Edge>, StorageError>>, StorageError> {
2834 const CHUNK: usize = 900;
2835 let mut outcomes = Vec::with_capacity(ids.len());
2836 for chunk in ids.chunks(CHUNK) {
2837 let requested: Vec<String> =
2838 chunk.iter().map(|id| Uuid::from(*id).to_string()).collect();
2839 let expected = requested.len();
2840 let rows = self
2841 .with_reader("get_edge_read_outcomes", move |conn| {
2842 let requested_json = serde_json::to_string(&requested)
2843 .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))?;
2844 let mut stmt = conn.prepare(
2845 "SELECT e.namespace, e.id, e.source_id, e.target_id, e.relation, e.weight, \
2846 e.created_at, e.updated_at, e.deleted_at, e.metadata, e.target_backend, \
2847 requested.key \
2848 FROM json_each(?1) AS requested \
2849 LEFT JOIN graph_edges AS e \
2850 ON e.id = requested.value AND e.deleted_at IS NULL \
2851 ORDER BY CAST(requested.key AS INTEGER) ASC",
2852 )?;
2853 let mut rows = stmt.query([requested_json])?;
2854 let mut outcomes = Vec::with_capacity(expected);
2855 while let Some(row) = rows.next()? {
2856 let ordinal: i64 = row.get(11)?;
2857 let ordinal =
2858 usize::try_from(ordinal).map_err(|_| rusqlite::Error::InvalidQuery)?;
2859 if ordinal != outcomes.len() {
2860 return Err(rusqlite::Error::InvalidQuery);
2861 }
2862 let outcome = match row.get_ref(1)? {
2863 rusqlite::types::ValueRef::Null => Ok(None),
2864 _ => read_edge(row).map(Some).map_err(|e| map_err(e, "get_edge")),
2865 };
2866 outcomes.push(outcome);
2867 }
2868 if outcomes.len() != expected {
2869 return Err(rusqlite::Error::InvalidQuery);
2870 }
2871 Ok(outcomes)
2872 })
2873 .await?;
2874 outcomes.extend(rows);
2875 }
2876 Ok(outcomes)
2877 }
2878
2879 async fn batch_neighbors(
2880 &self,
2881 sources: &[Uuid],
2882 query: NeighborQuery,
2883 ) -> Result<Vec<(Uuid, NeighborHit)>, StorageError> {
2884 use khive_storage::types::Direction;
2885
2886 if sources.is_empty() {
2887 return Ok(Vec::new());
2888 }
2889 let mut seen_sources = HashSet::with_capacity(sources.len());
2890 let unique_sources: Vec<Uuid> = sources
2891 .iter()
2892 .copied()
2893 .filter(|source| seen_sources.insert(*source))
2894 .collect();
2895 const CHUNK_SIZE: usize = 880;
2896
2897 let namespace = self.namespace.clone();
2898 let mut result: Vec<(Uuid, NeighborHit)> = Vec::new();
2899 let counted_queries = Arc::new(std::sync::atomic::AtomicU64::new(0));
2902 let counted_rows = Arc::new(std::sync::atomic::AtomicU64::new(0));
2903
2904 for chunk in unique_sources.chunks(CHUNK_SIZE) {
2905 let chunk_owned: Vec<Uuid> = chunk.to_vec();
2906 let query_clone = query.clone();
2907 let ns = namespace.clone();
2908 let closure_queries = Arc::clone(&counted_queries);
2909 let closure_rows = Arc::clone(&counted_rows);
2910
2911 let chunk_result = self
2912 .with_reader("batch_neighbors", move |conn| {
2913 let src_strs: Vec<String> = chunk_owned.iter().map(|u| u.to_string()).collect();
2914
2915 let sources_json = serde_json::to_string(&src_strs).map_err(|error| {
2916 rusqlite::Error::ToSqlConversionFailure(Box::new(error))
2917 })?;
2918
2919 let build_inner_sql =
2920 |direction_out: bool,
2921 q: &NeighborQuery|
2922 -> (String, Vec<String>, Option<f64>) {
2923 let (filter_col, node_col) = if direction_out {
2924 ("source_id", "target_id")
2925 } else {
2926 ("target_id", "source_id")
2927 };
2928
2929 let mut rel_params: Vec<String> = Vec::new();
2930 let mut conditions: Vec<String> = Vec::new();
2931 let mut param_idx = 3;
2932
2933 if let Some(ref rels) = q.relations {
2934 if !rels.is_empty() {
2935 let ps: Vec<String> = rels
2936 .iter()
2937 .map(|r| {
2938 rel_params.push(r.to_string());
2939 let p = format!("?{param_idx}");
2940 param_idx += 1;
2941 p
2942 })
2943 .collect();
2944 conditions
2945 .push(format!("edges.relation IN ({})", ps.join(",")));
2946 }
2947 }
2948
2949 let min_weight_val = if let Some(min_w) = q.min_weight {
2952 conditions.push(format!("edges.weight >= ?{param_idx}"));
2953 Some(min_w)
2954 } else {
2955 None
2956 };
2957
2958 let where_extra = if conditions.is_empty() {
2959 String::new()
2960 } else {
2961 format!(" AND {}", conditions.join(" AND "))
2962 };
2963
2964 let sql = format!(
2965 "SELECT requested.origin_id, edges.{node_col} AS node_id, \
2966 edges.id AS edge_id, edges.relation, edges.weight \
2967 FROM requested CROSS JOIN graph_edges AS edges \
2968 ON edges.{filter_col} = requested.origin_id \
2969 WHERE edges.namespace = ?1 \
2970 AND edges.deleted_at IS NULL{where_extra}",
2971 );
2972 (sql, rel_params, min_weight_val)
2973 };
2974
2975 let mut all_params: Vec<Box<dyn rusqlite::types::ToSql>> = Vec::new();
2976 all_params.push(Box::new(ns.to_string()));
2977 all_params.push(Box::new(sources_json));
2978
2979 let (combined_inner, rel_params, min_weight_val) = match query_clone.direction {
2980 Direction::Out => build_inner_sql(true, &query_clone),
2981 Direction::In => build_inner_sql(false, &query_clone),
2982 Direction::Both => {
2983 let (out_sql, rel_params, min_weight_val) =
2984 build_inner_sql(true, &query_clone);
2985 let (in_sql, _, _) = build_inner_sql(false, &query_clone);
2986 (
2987 format!("{out_sql} UNION ALL {in_sql}"),
2988 rel_params,
2989 min_weight_val,
2990 )
2991 }
2992 };
2993
2994 for relation in rel_params {
2995 all_params.push(Box::new(relation));
2996 }
2997 if let Some(min_weight) = min_weight_val {
2998 all_params.push(Box::new(min_weight));
2999 }
3000 let limit_param_idx = all_params.len() + 1;
3001
3002 let full_sql = if let Some(lim) = query_clone.limit {
3010 all_params.push(Box::new(lim as i64));
3011 format!(
3012 "WITH requested(origin_id) AS (\
3013 SELECT value FROM json_each(?2)\
3014 ) SELECT origin_id, node_id, edge_id, relation, weight \
3015 FROM (SELECT *, ROW_NUMBER() OVER (PARTITION BY origin_id \
3016 ORDER BY weight DESC, node_id ASC) AS rn \
3017 FROM ({combined_inner})) WHERE rn <= ?{limit_param_idx}",
3018 )
3019 } else {
3020 format!(
3021 "WITH requested(origin_id) AS (\
3022 SELECT value FROM json_each(?2)\
3023 ) SELECT origin_id, node_id, edge_id, relation, weight \
3024 FROM ({combined_inner})",
3025 )
3026 };
3027
3028 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
3029 all_params.iter().map(|p| p.as_ref()).collect();
3030
3031 let mut stmt = conn.prepare(&full_sql)?;
3032 closure_queries.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3033 let rows = stmt.query_map(param_refs.as_slice(), |row| {
3034 let origin_str: String = row.get(0)?;
3035 let nid_str: String = row.get(1)?;
3036 let eid_str: String = row.get(2)?;
3037 let relation_str: String = row.get(3)?;
3038 let weight: f64 = row.get(4)?;
3039 Ok((origin_str, nid_str, eid_str, relation_str, weight))
3040 })?;
3041
3042 let mut pairs = Vec::new();
3043 for row in rows {
3044 let (origin_str, nid_str, eid_str, relation_str, weight) = row?;
3045 closure_rows.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3046 let origin = parse_uuid(&origin_str)?;
3047 let node_id = parse_uuid(&nid_str)?;
3048 let edge_id = parse_uuid(&eid_str)?;
3049 let relation = relation_str.parse::<EdgeRelation>().map_err(|e| {
3050 rusqlite::Error::FromSqlConversionFailure(
3051 3,
3052 rusqlite::types::Type::Text,
3053 Box::new(e),
3054 )
3055 })?;
3056 pairs.push((
3057 origin,
3058 NeighborHit {
3059 node_id,
3060 edge_id,
3061 relation,
3062 weight,
3063 name: None,
3064 kind: None,
3065 entity_type: None,
3066 },
3067 ));
3068 }
3069 Ok(pairs)
3070 })
3071 .await;
3072 let pairs = match chunk_result {
3073 Ok(pairs) => pairs,
3074 Err(e) => {
3075 report_graph_usage(&counted_queries, &counted_rows);
3076 return Err(e);
3077 }
3078 };
3079 result.extend(pairs);
3080 }
3081 report_graph_usage(&counted_queries, &counted_rows);
3082
3083 let requested: HashSet<Uuid> = unique_sources.iter().copied().collect();
3084 let mut grouped: HashMap<Uuid, Vec<NeighborHit>> =
3085 HashMap::with_capacity(unique_sources.len());
3086 for (origin, hit) in result {
3087 if !requested.contains(&origin) {
3088 return Err(StorageError::Internal(format!(
3089 "batch_neighbors returned unrequested origin {origin}"
3090 )));
3091 }
3092 grouped.entry(origin).or_default().push(hit);
3093 }
3094
3095 for hits in grouped.values_mut() {
3096 hits.sort_by(|a, b| {
3097 b.weight
3098 .partial_cmp(&a.weight)
3099 .unwrap_or(std::cmp::Ordering::Equal)
3100 .then(a.node_id.cmp(&b.node_id))
3101 .then(a.edge_id.cmp(&b.edge_id))
3102 });
3103 }
3104
3105 let mut ordered = Vec::new();
3106 for &source in sources {
3107 if let Some(hits) = grouped.get(&source) {
3108 ordered.extend(hits.iter().cloned().map(|hit| (source, hit)));
3109 }
3110 }
3111 Ok(ordered)
3112 }
3113
3114 async fn delete_edge(&self, id: LinkId, mode: DeleteMode) -> Result<bool, StorageError> {
3115 let id = Uuid::from(id);
3116 let statement = match mode {
3117 DeleteMode::Soft => {
3118 edge_soft_delete_statement(id, chrono::Utc::now().timestamp_micros())
3119 }
3120 DeleteMode::Hard => edge_hard_delete_statement(id),
3121 };
3122 self.with_writer("delete_edge", move |conn| {
3123 let mut stmt = conn.prepare(&statement.sql)?;
3124 bind_params(&mut stmt, &statement.params)?;
3125 Ok(stmt.raw_execute()? > 0)
3126 })
3127 .await
3128 }
3129
3130 async fn query_edges(
3131 &self,
3132 filter: EdgeFilter,
3133 sort: Vec<SortOrder<EdgeSortField>>,
3134 page: PageRequest,
3135 ) -> Result<Page<Edge>, StorageError> {
3136 let namespace = self.namespace.clone();
3137 let limit_i64 = i64::from(page.limit);
3138 let offset_i64 = i64::try_from(page.offset).map_err(|_| StorageError::InvalidInput {
3139 capability: StorageCapability::Graph,
3140 operation: "query_edges".into(),
3141 message: format!(
3142 "PageRequest: offset must be <= i64::MAX, got {}",
3143 page.offset
3144 ),
3145 })?;
3146 self.with_reader("query_edges", move |conn| {
3147 let (where_clause, mut all_params) = build_edge_filter_sql(&namespace, &filter);
3148 let order_clause = edge_order_clause(&sort);
3149 all_params.push(Box::new(limit_i64));
3150 all_params.push(Box::new(offset_i64));
3151
3152 let limit_idx = all_params.len() - 1;
3153 let offset_idx = all_params.len();
3154
3155 let data_sql = format!(
3156 "SELECT namespace, id, source_id, target_id, relation, weight, \
3157 created_at, updated_at, deleted_at, metadata, target_backend \
3158 FROM graph_edges{}{} LIMIT ?{} OFFSET ?{}",
3159 where_clause, order_clause, limit_idx, offset_idx,
3160 );
3161
3162 let mut stmt = conn.prepare(&data_sql)?;
3163 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
3164 all_params.iter().map(|p| p.as_ref()).collect();
3165 let rows = stmt.query_map(param_refs.as_slice(), read_edge)?;
3166
3167 let mut items = Vec::new();
3168 for row in rows {
3169 items.push(row?);
3170 }
3171
3172 Ok(Page { items, total: None })
3173 })
3174 .await
3175 }
3176
3177 async fn count_edges(&self, filter: EdgeFilter) -> Result<u64, StorageError> {
3178 let namespace = self.namespace.clone();
3179 self.with_reader("count_edges", move |conn| {
3180 let (where_clause, params) = build_edge_filter_sql(&namespace, &filter);
3181 let sql = format!(
3182 "SELECT COUNT(*) FROM graph_edges{}",
3183 with_live_endpoints(&where_clause)
3184 );
3185 let mut stmt = conn.prepare(&sql)?;
3186 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
3187 params.iter().map(|p| p.as_ref()).collect();
3188 let count: i64 = stmt.query_row(param_refs.as_slice(), |row| row.get(0))?;
3189 Ok(count as u64)
3190 })
3191 .await
3192 }
3193
3194 async fn count_edges_in_namespaces(
3195 &self,
3196 namespaces: &[String],
3197 filter: EdgeFilter,
3198 ) -> Result<u64, StorageError> {
3199 let namespaces: Vec<String> = namespaces
3200 .iter()
3201 .cloned()
3202 .collect::<HashSet<_>>()
3203 .into_iter()
3204 .collect();
3205 self.with_reader("count_edges_in_namespaces", move |conn| {
3206 let mut total = 0;
3207 for chunk in namespaces.chunks(NAMESPACE_COUNT_CHUNK_SIZE) {
3208 let (where_clause, params) = build_edge_filter_sql_for_namespaces(chunk, &filter);
3209 let sql = format!(
3210 "SELECT COUNT(*) FROM graph_edges{}",
3211 with_live_endpoints(&where_clause)
3212 );
3213 let mut stmt = conn.prepare(&sql)?;
3214 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
3215 params.iter().map(|p| p.as_ref()).collect();
3216 let count: i64 = stmt.query_row(param_refs.as_slice(), |row| row.get(0))?;
3217 total += count as u64;
3218 }
3219 Ok(total)
3220 })
3221 .await
3222 }
3223
3224 async fn query_edges_in_namespaces(
3225 &self,
3226 namespaces: &[String],
3227 filter: EdgeFilter,
3228 sort: Vec<SortOrder<EdgeSortField>>,
3229 page: PageRequest,
3230 ) -> Result<Page<Edge>, StorageError> {
3231 let namespaces: Vec<String> = {
3246 let mut seen = HashSet::new();
3247 namespaces
3248 .iter()
3249 .filter(|ns| seen.insert((*ns).clone()))
3250 .cloned()
3251 .collect()
3252 };
3253 let limit_i64 = i64::from(page.limit);
3254 let offset_i64 = i64::try_from(page.offset).map_err(|_| StorageError::InvalidInput {
3255 capability: StorageCapability::Graph,
3256 operation: "query_edges_in_namespaces".into(),
3257 message: format!(
3258 "PageRequest: offset must be <= i64::MAX, got {}",
3259 page.offset
3260 ),
3261 })?;
3262 self.with_reader("query_edges_in_namespaces", move |conn| {
3263 let namespaces_json = serde_json::to_string(&namespaces)
3264 .map_err(|error| rusqlite::Error::ToSqlConversionFailure(Box::new(error)))?;
3265
3266 let (where_clause, mut all_params) =
3267 build_edge_filter_sql_for_namespaces_json(&namespaces_json, &filter);
3268 let order_clause = edge_order_clause(&sort);
3269 all_params.push(Box::new(limit_i64));
3270 all_params.push(Box::new(offset_i64));
3271
3272 let limit_idx = all_params.len() - 1;
3273 let offset_idx = all_params.len();
3274
3275 let data_sql = format!(
3276 "SELECT namespace, id, source_id, target_id, relation, weight, \
3277 created_at, updated_at, deleted_at, metadata, target_backend \
3278 FROM graph_edges{}{} LIMIT ?{} OFFSET ?{}",
3279 where_clause, order_clause, limit_idx, offset_idx,
3280 );
3281
3282 let mut stmt = conn.prepare(&data_sql)?;
3283 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
3284 all_params.iter().map(|p| p.as_ref()).collect();
3285 let rows = stmt.query_map(param_refs.as_slice(), read_edge)?;
3286
3287 let mut items = Vec::new();
3288 for row in rows {
3289 items.push(row?);
3290 }
3291
3292 Ok(Page { items, total: None })
3293 })
3294 .await
3295 }
3296
3297 async fn count_edges_by_relation(&self) -> Result<Vec<(EdgeRelation, u64)>, StorageError> {
3298 let namespace = self.namespace.clone();
3299 self.with_reader("count_edges_by_relation", move |conn| {
3300 let sql = format!(
3301 "SELECT relation, COUNT(*) FROM graph_edges \
3302 WHERE namespace = ?1 AND deleted_at IS NULL AND {LIVE_ENDPOINTS_CONDITION} \
3303 GROUP BY relation"
3304 );
3305 let mut stmt = conn.prepare(&sql)?;
3306 let rows = stmt.query_map([&namespace], |row| {
3307 let relation_str: String = row.get(0)?;
3308 let count: i64 = row.get(1)?;
3309 Ok((relation_str, count))
3310 })?;
3311 let mut out = Vec::new();
3312 for row in rows {
3313 let (relation_str, count) = row?;
3314 let relation = relation_str.parse::<EdgeRelation>().map_err(|e| {
3315 rusqlite::Error::FromSqlConversionFailure(
3316 0,
3317 rusqlite::types::Type::Text,
3318 Box::new(e),
3319 )
3320 })?;
3321 out.push((relation, count as u64));
3322 }
3323 Ok(out)
3324 })
3325 .await
3326 }
3327
3328 async fn count_edges_by_relation_in_namespaces(
3329 &self,
3330 namespaces: &[String],
3331 ) -> Result<Vec<(EdgeRelation, u64)>, StorageError> {
3332 let namespaces: Vec<String> = namespaces
3333 .iter()
3334 .cloned()
3335 .collect::<HashSet<_>>()
3336 .into_iter()
3337 .collect();
3338 self.with_reader("count_edges_by_relation_in_namespaces", move |conn| {
3339 let mut totals = HashMap::new();
3340 for chunk in namespaces.chunks(NAMESPACE_COUNT_CHUNK_SIZE) {
3341 let (where_clause, params) =
3342 build_edge_filter_sql_for_namespaces(chunk, &EdgeFilter::default());
3343 let sql = format!(
3344 "SELECT relation, COUNT(*) FROM graph_edges{} GROUP BY relation",
3345 with_live_endpoints(&where_clause)
3346 );
3347 let mut stmt = conn.prepare(&sql)?;
3348 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
3349 params.iter().map(|p| p.as_ref()).collect();
3350 let rows = stmt.query_map(param_refs.as_slice(), |row| {
3351 let relation_str: String = row.get(0)?;
3352 let count: i64 = row.get(1)?;
3353 Ok((relation_str, count))
3354 })?;
3355 for row in rows {
3356 let (relation_str, count) = row?;
3357 let relation = relation_str.parse::<EdgeRelation>().map_err(|e| {
3358 rusqlite::Error::FromSqlConversionFailure(
3359 0,
3360 rusqlite::types::Type::Text,
3361 Box::new(e),
3362 )
3363 })?;
3364 *totals.entry(relation).or_insert(0) += count as u64;
3365 }
3366 }
3367 Ok(totals.into_iter().collect())
3368 })
3369 .await
3370 }
3371
3372 async fn count_edges_by_endpoint_base(&self) -> Result<EdgeEndpointBaseCounts, StorageError> {
3373 let namespace = self.namespace.clone();
3374 self.with_reader("count_edges_by_endpoint_base", move |conn| {
3375 let source_case = endpoint_base_case("source_id");
3376 let target_case = endpoint_base_case("target_id");
3377 let sql = format!(
3378 "SELECT {source_case} AS source_base, {target_case} AS target_base, COUNT(*) \
3379 FROM graph_edges \
3380 WHERE namespace = ?1 AND deleted_at IS NULL AND {LIVE_ENDPOINTS_CONDITION} \
3381 GROUP BY source_base, target_base"
3382 );
3383 let mut stmt = conn.prepare(&sql)?;
3384 let rows = stmt.query_map([&namespace], |row| {
3385 let source: String = row.get(0)?;
3386 let target: String = row.get(1)?;
3387 let count: i64 = row.get(2)?;
3388 Ok((source, target, count))
3389 })?;
3390 let mut counts = EdgeEndpointBaseCounts::default();
3391 for row in rows {
3392 let (source, target, count) = row?;
3393 fold_endpoint_base_row(&mut counts, &source, &target, count as u64);
3394 }
3395 Ok(counts)
3396 })
3397 .await
3398 }
3399
3400 async fn count_edges_by_endpoint_base_in_namespaces(
3401 &self,
3402 namespaces: &[String],
3403 ) -> Result<EdgeEndpointBaseCounts, StorageError> {
3404 let namespaces: Vec<String> = namespaces
3405 .iter()
3406 .cloned()
3407 .collect::<HashSet<_>>()
3408 .into_iter()
3409 .collect();
3410 self.with_reader("count_edges_by_endpoint_base_in_namespaces", move |conn| {
3411 let source_case = endpoint_base_case("source_id");
3412 let target_case = endpoint_base_case("target_id");
3413 let mut counts = EdgeEndpointBaseCounts::default();
3414 for chunk in namespaces.chunks(NAMESPACE_COUNT_CHUNK_SIZE) {
3415 let (where_clause, params) =
3416 build_edge_filter_sql_for_namespaces(chunk, &EdgeFilter::default());
3417 let sql = format!(
3418 "SELECT {source_case} AS source_base, {target_case} AS target_base, COUNT(*) \
3419 FROM graph_edges{} GROUP BY source_base, target_base",
3420 with_live_endpoints(&where_clause)
3421 );
3422 let mut stmt = conn.prepare(&sql)?;
3423 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
3424 params.iter().map(|p| p.as_ref()).collect();
3425 let rows = stmt.query_map(param_refs.as_slice(), |row| {
3426 let source: String = row.get(0)?;
3427 let target: String = row.get(1)?;
3428 let count: i64 = row.get(2)?;
3429 Ok((source, target, count))
3430 })?;
3431 for row in rows {
3432 let (source, target, count) = row?;
3433 fold_endpoint_base_row(&mut counts, &source, &target, count as u64);
3434 }
3435 }
3436 Ok(counts)
3437 })
3438 .await
3439 }
3440
3441 async fn count_edges_by_relation_and_endpoint_base_in_namespaces(
3442 &self,
3443 namespaces: &[String],
3444 ) -> Result<Vec<(EdgeRelation, EdgeEndpointBaseCounts)>, StorageError> {
3445 let namespaces: Vec<String> = namespaces
3446 .iter()
3447 .cloned()
3448 .collect::<HashSet<_>>()
3449 .into_iter()
3450 .collect();
3451 self.with_reader(
3452 "count_edges_by_relation_and_endpoint_base_in_namespaces",
3453 move |conn| {
3454 let mut totals: HashMap<EdgeRelation, EdgeEndpointBaseCounts> = HashMap::new();
3455 for chunk in namespaces.chunks(NAMESPACE_COUNT_CHUNK_SIZE) {
3456 let (where_clause, params) =
3457 build_edge_filter_sql_for_namespaces(chunk, &EdgeFilter::default());
3458 let sql = format!(
3468 "SELECT e.relation, \
3469 CASE WHEN se.id IS NOT NULL THEN 'entity' \
3470 WHEN sn.id IS NOT NULL THEN 'note' ELSE 'none' END AS source_base, \
3471 CASE WHEN te.id IS NOT NULL THEN 'entity' \
3472 WHEN tn.id IS NOT NULL THEN 'note' ELSE 'none' END AS target_base, \
3473 COUNT(*) \
3474 FROM (SELECT relation, source_id, target_id \
3475 FROM graph_edges{where_clause}) e \
3476 LEFT JOIN entities se ON se.id = e.source_id \
3477 LEFT JOIN notes sn ON sn.id = e.source_id \
3478 LEFT JOIN entities te ON te.id = e.target_id \
3479 LEFT JOIN notes tn ON tn.id = e.target_id \
3480 WHERE se.deleted_at IS NULL AND sn.deleted_at IS NULL \
3481 AND te.deleted_at IS NULL AND tn.deleted_at IS NULL \
3482 GROUP BY e.relation, source_base, target_base"
3483 );
3484 let mut stmt = conn.prepare(&sql)?;
3485 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
3486 params.iter().map(|p| p.as_ref()).collect();
3487 let rows = stmt.query_map(param_refs.as_slice(), |row| {
3488 let relation: String = row.get(0)?;
3489 let source: String = row.get(1)?;
3490 let target: String = row.get(2)?;
3491 let count: i64 = row.get(3)?;
3492 Ok((relation, source, target, count))
3493 })?;
3494 for row in rows {
3495 let (relation_str, source, target, count) = row?;
3496 let relation = relation_str.parse::<EdgeRelation>().map_err(|e| {
3497 rusqlite::Error::FromSqlConversionFailure(
3498 0,
3499 rusqlite::types::Type::Text,
3500 Box::new(e),
3501 )
3502 })?;
3503 fold_endpoint_base_row(
3504 totals.entry(relation).or_default(),
3505 &source,
3506 &target,
3507 count as u64,
3508 );
3509 }
3510 }
3511 Ok(totals.into_iter().collect())
3512 },
3513 )
3514 .await
3515 }
3516
3517 async fn query_edges_after(
3518 &self,
3519 filter: EdgeFilter,
3520 after: Option<Uuid>,
3521 limit: u32,
3522 ) -> Result<EdgeSeekPage, StorageError> {
3523 let namespace = self.namespace.clone();
3524 let limit_usize = limit as usize;
3525 let probe_limit_i64 = i64::from(limit) + 1;
3526 self.with_reader("query_edges_after", move |conn| {
3527 let (mut where_clause, mut params) = build_edge_filter_sql(&namespace, &filter);
3528 if let Some(cursor) = after {
3529 params.push(Box::new(cursor.to_string()));
3530 where_clause.push_str(&format!(" AND id > ?{}", params.len()));
3531 }
3532 params.push(Box::new(probe_limit_i64));
3533 let limit_idx = params.len();
3534
3535 let data_sql = format!(
3541 "SELECT namespace, id, source_id, target_id, relation, weight, \
3542 created_at, updated_at, deleted_at, metadata, target_backend \
3543 FROM graph_edges{} ORDER BY id ASC LIMIT ?{}",
3544 where_clause, limit_idx,
3545 );
3546
3547 let mut stmt = conn.prepare(&data_sql)?;
3548 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
3549 params.iter().map(|p| p.as_ref()).collect();
3550 let rows = stmt.query_map(param_refs.as_slice(), read_edge)?;
3551
3552 let mut items = Vec::new();
3553 for row in rows {
3554 items.push(row?);
3555 }
3556 let has_more = items.len() > limit_usize;
3557 if has_more {
3558 items.truncate(limit_usize);
3559 }
3560 let next_after = if has_more {
3561 items.last().map(|e| Uuid::from(e.id))
3562 } else {
3563 None
3564 };
3565
3566 Ok(EdgeSeekPage { items, next_after })
3567 })
3568 .await
3569 }
3570
3571 async fn query_edges_sequence_after(
3572 &self,
3573 filter: EdgeFilter,
3574 after: Option<SeekCursor>,
3575 limit: u32,
3576 ) -> Result<SeekPage<Edge>, StorageError> {
3577 if limit == 0 {
3578 return Ok(SeekPage::default());
3579 }
3580 let namespace = self.namespace.clone();
3581 let limit_usize = limit as usize;
3582 let probe_limit_i64 = i64::from(limit) + 1;
3583 self.with_reader("query_edges_sequence_after", move |conn| {
3584 let (mut where_clause, mut params) = build_edge_filter_sql(&namespace, &filter);
3585 if let Some(cursor) = after {
3586 params.push(Box::new(cursor.sequence));
3587 where_clause.push_str(&format!(" AND graph_edges_seq.seq > ?{}", params.len()));
3588 }
3589 params.push(Box::new(probe_limit_i64));
3590 let limit_idx = params.len();
3591 let sql = format!(
3594 "SELECT namespace, id, source_id, target_id, relation, weight, \
3595 created_at, updated_at, deleted_at, metadata, target_backend, \
3596 graph_edges_seq.seq \
3597 FROM graph_edges_seq CROSS JOIN graph_edges \
3598 ON graph_edges.id = graph_edges_seq.edge_id{where_clause} \
3599 ORDER BY graph_edges_seq.seq ASC LIMIT ?{limit_idx}"
3600 );
3601 let mut stmt = conn.prepare(&sql)?;
3602 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
3603 params.iter().map(|param| param.as_ref()).collect();
3604 let rows = stmt.query_map(param_refs.as_slice(), |row| {
3605 Ok((read_edge(row)?, row.get::<_, i64>(11)?))
3606 })?;
3607 let mut entries = rows.collect::<Result<Vec<_>, _>>()?;
3608 let has_more = entries.len() > limit_usize;
3609 if has_more {
3610 entries.truncate(limit_usize);
3611 }
3612 let next_after = if has_more {
3613 entries.last().map(|(edge, sequence)| SeekCursor {
3614 sequence: *sequence,
3615 id: Uuid::from(edge.id),
3616 })
3617 } else {
3618 None
3619 };
3620 let items = entries.into_iter().map(|(edge, _)| edge).collect();
3621 Ok(SeekPage { items, next_after })
3622 })
3623 .await
3624 }
3625
3626 async fn neighbors(
3627 &self,
3628 node_id: Uuid,
3629 query: NeighborQuery,
3630 ) -> Result<Vec<NeighborHit>, StorageError> {
3631 self.query_neighbors_page("neighbors", node_id, query, None, None)
3632 .await
3633 }
3634
3635 async fn neighbors_page(
3636 &self,
3637 node_id: Uuid,
3638 query: NeighborQuery,
3639 after: Option<NeighborCursor>,
3640 neighbor_kinds: Option<Vec<String>>,
3641 ) -> Result<Vec<NeighborHit>, StorageError> {
3642 self.query_neighbors_page("neighbors_page", node_id, query, after, neighbor_kinds)
3643 .await
3644 }
3645
3646 async fn neighbors_both_directions(
3651 &self,
3652 node_id: Uuid,
3653 query: NeighborQuery,
3654 ) -> Result<Vec<DirectedNeighborHit>, StorageError> {
3655 count_neighbor_select();
3656
3657 let namespace = self.namespace.clone();
3658 let node_str = node_id.to_string();
3659
3660 let counted_queries = Arc::new(std::sync::atomic::AtomicU64::new(0));
3661 let counted_rows = Arc::new(std::sync::atomic::AtomicU64::new(0));
3662 let closure_queries = Arc::clone(&counted_queries);
3663 let closure_rows = Arc::clone(&counted_rows);
3664 let result = self
3665 .with_reader("neighbors_both_directions", move |conn| {
3666 let base_out = "SELECT target_id AS node_id, id AS edge_id, relation, weight, \
3667 'out' AS dir \
3668 FROM graph_edges \
3669 WHERE namespace = ?1 AND source_id = ?2 AND deleted_at IS NULL";
3670 let base_in = "SELECT source_id AS node_id, id AS edge_id, relation, weight, \
3671 'in' AS dir \
3672 FROM graph_edges \
3673 WHERE namespace = ?1 AND target_id = ?2 AND deleted_at IS NULL";
3674 let sql = format!("{} UNION ALL {}", base_out, base_in);
3675
3676 let (where_extra, limit_clause, extra_params) =
3677 neighbor_extra_clause(&query, 3, None, None);
3678
3679 let full_sql = format!(
3688 "SELECT node_id, edge_id, relation, weight, dir FROM ({}){} \
3689 ORDER BY weight DESC, node_id ASC, \
3690 CASE dir WHEN 'out' THEN 0 ELSE 1 END ASC, edge_id ASC{}",
3691 sql, where_extra, limit_clause
3692 );
3693
3694 let mut stmt = conn.prepare(&full_sql)?;
3695
3696 let mut all_params: Vec<Box<dyn rusqlite::types::ToSql>> = Vec::new();
3697 all_params.push(Box::new(namespace.clone()));
3698 all_params.push(Box::new(node_str.clone()));
3699 all_params.extend(extra_params);
3700
3701 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
3702 all_params.iter().map(|p| p.as_ref()).collect();
3703
3704 closure_queries.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3705 let rows = stmt.query_map(param_refs.as_slice(), |row| {
3706 let nid_str: String = row.get(0)?;
3707 let eid_str: String = row.get(1)?;
3708 let relation_str: String = row.get(2)?;
3709 let weight: f64 = row.get(3)?;
3710 let dir_str: String = row.get(4)?;
3711 Ok((nid_str, eid_str, relation_str, weight, dir_str))
3712 })?;
3713
3714 let mut hits = Vec::new();
3715 for row in rows {
3716 let (nid_str, eid_str, relation_str, weight, dir_str) = row?;
3717 closure_rows.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3718 let relation = relation_str.parse::<EdgeRelation>().map_err(|e| {
3719 rusqlite::Error::FromSqlConversionFailure(
3720 2,
3721 rusqlite::types::Type::Text,
3722 Box::new(e),
3723 )
3724 })?;
3725 let direction = if dir_str == "out" {
3726 Direction::Out
3727 } else {
3728 Direction::In
3729 };
3730 hits.push(DirectedNeighborHit {
3731 hit: NeighborHit {
3732 node_id: parse_uuid(&nid_str)?,
3733 edge_id: parse_uuid(&eid_str)?,
3734 relation,
3735 weight,
3736 name: None,
3737 kind: None,
3738 entity_type: None,
3739 },
3740 direction,
3741 });
3742 }
3743
3744 Ok(hits)
3745 })
3746 .await;
3747
3748 report_graph_usage(&counted_queries, &counted_rows);
3749 result
3750 }
3751
3752 async fn traverse(&self, request: TraversalRequest) -> Result<Vec<GraphPath>, StorageError> {
3753 request
3754 .validate()
3755 .map_err(|message| StorageError::InvalidInput {
3756 capability: StorageCapability::Graph,
3757 operation: "traverse".into(),
3758 message,
3759 })?;
3760 if request.roots.is_empty() {
3761 return Ok(Vec::new());
3762 }
3763
3764 let budget = request.execution_budget;
3765 let outer_context = khive_storage::capture_request_read_context();
3766 let execution_deadline =
3767 khive_storage::RequestReadDeadline::after(budget.remaining_duration());
3768 khive_storage::scope_request_read_deadline_at(execution_deadline, async {
3769 if outer_context.stop_reason().is_some() {
3770 return Err(StorageError::Timeout {
3771 operation: "traverse".into(),
3772 });
3773 }
3774 if budget.is_expired() {
3775 return Err(traversal_timeout_error(&budget));
3776 }
3777 let mut distinct_roots = HashSet::with_capacity(request.roots.len());
3778 let roots = request
3779 .roots
3780 .iter()
3781 .copied()
3782 .filter(|root| distinct_roots.insert(*root))
3783 .collect::<Vec<_>>();
3784 let opts = request.options;
3785 if self
3786 .index_repair
3787 .as_ref()
3788 .is_some_and(super::index_repair::IndexRepairContext::is_writable)
3789 {
3790 let directions = match opts.direction {
3796 Direction::Out => vec![Direction::Out],
3797 Direction::In => vec![Direction::In],
3798 Direction::Both => vec![Direction::Out, Direction::In],
3799 };
3800 let relation_count = opts.relations.as_ref().map_or(0, Vec::len);
3801 let statements = directions
3802 .into_iter()
3803 .map(|direction| {
3804 traversal_neighbor_sql(direction, relation_count, opts.min_weight.is_some())
3805 })
3806 .collect::<Vec<_>>();
3807 let preflight = self.with_indexed_reader("traverse", move |conn| {
3810 for sql in &statements {
3811 let mut statement = conn.prepare(sql)?;
3812 statement.raw_bind_parameter(3, 0_i64)?;
3813 let _ = statement.raw_query().next()?;
3814 }
3815 Ok(())
3816 });
3817 let preflight_result = tokio::select! {
3818 biased;
3819 _ = outer_context.clone().wait_for_stop() => Err(StorageError::Timeout {
3820 operation: "traverse".into(),
3821 }),
3822 result = tokio::time::timeout_at(execution_deadline.async_at(), preflight) => {
3823 match result {
3824 Ok(result) => result,
3825 Err(_) => Err(traversal_timeout_error(&budget)),
3826 }
3827 },
3828 };
3829 preflight_result.map_err(|error| match error {
3830 StorageError::Timeout { .. }
3831 if budget.is_expired() && outer_context.stop_reason().is_none() =>
3832 {
3833 traversal_timeout_error(&budget)
3834 }
3835 error => error,
3836 })?;
3837 if budget.is_expired() {
3838 return Err(traversal_timeout_error(&budget));
3839 }
3840 }
3841 let include_roots = request.include_roots;
3842 let namespace = self.namespace.clone();
3843 let origin = self.pool.origin();
3844 let closure_budget = budget.clone();
3845 let counted_rows = Arc::new(std::sync::atomic::AtomicU64::new(0));
3851 let counted_queries = Arc::new(std::sync::atomic::AtomicU64::new(0));
3852 let closure_rows = Arc::clone(&counted_rows);
3853 let closure_queries = Arc::clone(&counted_queries);
3854 let result = self
3855 .with_reader("traverse", move |conn| {
3856 Ok(run_bounded_traversal(
3857 conn,
3858 roots,
3859 opts,
3860 include_roots,
3861 namespace,
3862 origin,
3863 closure_budget,
3864 closure_rows.as_ref(),
3865 closure_queries.as_ref(),
3866 ))
3867 })
3868 .await
3869 .and_then(|inner| inner)
3870 .map_err(|error| match error {
3871 StorageError::Timeout { .. }
3872 if budget.is_expired() && outer_context.stop_reason().is_none() =>
3873 {
3874 traversal_timeout_error(&budget)
3875 }
3876 error => error,
3877 });
3878
3879 khive_storage::usage::count(
3888 khive_storage::usage::UsageUnit::DbRoundTrips,
3889 counted_queries.load(std::sync::atomic::Ordering::Relaxed),
3890 );
3891 khive_storage::usage::count(
3892 khive_storage::usage::UsageUnit::GraphHops,
3893 counted_rows.load(std::sync::atomic::Ordering::Relaxed),
3894 );
3895
3896 result
3897 })
3898 .await
3899 }
3900
3901 async fn purge_incident_edges(&self, node_id: Uuid) -> Result<u64, StorageError> {
3902 let statement = purge_incident_edges_statement(node_id);
3906 self.with_writer("purge_incident_edges", move |conn| {
3907 let mut stmt = conn.prepare(&statement.sql)?;
3908 bind_params(&mut stmt, &statement.params)?;
3909 Ok(stmt.raw_execute()? as u64)
3910 })
3911 .await
3912 }
3913}
3914
3915const GRAPH_DDL: &str = include_str!("../../sql/graph-ddl.sql");
3920
3921pub(crate) fn ensure_graph_schema(conn: &rusqlite::Connection) -> Result<(), rusqlite::Error> {
3922 conn.execute_batch(GRAPH_DDL)
3923}
3924
3925#[cfg(test)]
3926#[path = "graph_annotation_tests.rs"]
3927mod annotation_tests;
3928
3929#[cfg(test)]
3930#[path = "graph_tests.rs"]
3931mod tests;
3932
3933#[cfg(test)]
3934#[path = "incident_edge_plan_tests.rs"]
3935mod incident_edge_plan_tests;
3936
3937#[cfg(test)]
3938#[path = "graph_index_repair_tests.rs"]
3939mod index_repair_tests;
3940
3941#[cfg(test)]
3942#[path = "graph_busy_tests.rs"]
3943mod direct_busy_tests;