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_merge_rewire_statement(
556 namespace: &str,
557 id: Uuid,
558 source_id: Uuid,
559 target_id: Uuid,
560 updated_at: i64,
561) -> SqlStatement {
562 SqlStatement::new(
563 "UPDATE graph_edges SET \
564 source_id = ?1, target_id = ?2, updated_at = ?3 \
565 WHERE namespace = ?4 AND id = ?5",
566 vec![
567 SqlValue::Text(source_id.to_string()),
568 SqlValue::Text(target_id.to_string()),
569 SqlValue::Integer(updated_at),
570 SqlValue::Text(namespace.to_string()),
571 SqlValue::Text(id.to_string()),
572 ],
573 )
574}
575
576pub fn edge_hard_delete_statement(id: Uuid) -> SqlStatement {
578 SqlStatement {
579 sql: "DELETE FROM graph_edges WHERE id = ?1".to_string(),
580 params: vec![SqlValue::Text(id.to_string())],
581 label: Some("edge-delete-hard".to_string()),
582 }
583}
584
585pub(crate) const EDGE_NAMESPACES_CTE: &str = "WITH RECURSIVE edge_ns(namespace) AS ( \
596 SELECT (SELECT min(namespace) FROM graph_edges) \
597 UNION ALL \
598 SELECT (SELECT min(namespace) FROM graph_edges WHERE namespace > edge_ns.namespace) \
599 FROM edge_ns WHERE edge_ns.namespace IS NOT NULL)";
600
601pub fn purge_incident_edges_statement(node_id: Uuid) -> SqlStatement {
607 SqlStatement {
608 sql: format!(
609 "DELETE FROM graph_edges WHERE id IN ({EDGE_NAMESPACES_CTE} \
610 SELECT e.id FROM edge_ns CROSS JOIN graph_edges AS e \
611 ON e.namespace = edge_ns.namespace AND e.source_id = ?1 \
612 UNION ALL \
613 SELECT e.id FROM edge_ns CROSS JOIN graph_edges AS e \
614 ON e.namespace = edge_ns.namespace AND e.target_id = ?1)"
615 ),
616 params: vec![SqlValue::Text(node_id.to_string())],
617 label: Some("edge-purge-incident".to_string()),
618 }
619}
620
621#[allow(clippy::too_many_arguments)]
715pub fn edge_symmetric_delete_if_conflict_statement(
716 namespace: &str,
717 id: Uuid,
718 canon_src: Uuid,
719 canon_tgt: Uuid,
720 relation: EdgeRelation,
721 expected_updated_at_micros: i64,
722 expected_deleted_at_micros: Option<i64>,
723) -> SqlStatement {
724 SqlStatement {
725 sql: "DELETE FROM graph_edges \
726 WHERE namespace = ?1 AND id = ?2 \
727 AND updated_at = ?6 AND deleted_at IS ?7 \
728 AND EXISTS ( \
729 SELECT 1 FROM graph_edges \
730 WHERE namespace = ?1 AND source_id = ?3 AND target_id = ?4 \
731 AND relation = ?5 AND id != ?2 \
732 )"
733 .to_string(),
734 params: vec![
735 SqlValue::Text(namespace.to_string()),
736 SqlValue::Text(id.to_string()),
737 SqlValue::Text(canon_src.to_string()),
738 SqlValue::Text(canon_tgt.to_string()),
739 SqlValue::Text(relation.to_string()),
740 SqlValue::Integer(expected_updated_at_micros),
741 match expected_deleted_at_micros {
742 Some(value) => SqlValue::Integer(value),
743 None => SqlValue::Null,
744 },
745 ],
746 label: Some("edge-symmetric-delete-if-conflict".to_string()),
747 }
748}
749
750#[allow(clippy::too_many_arguments)]
759pub fn edge_symmetric_absorb_or_update_inplace_statement(
760 namespace: &str,
761 id: Uuid,
762 canon_src: Uuid,
763 canon_tgt: Uuid,
764 relation: EdgeRelation,
765 weight: f64,
766 updated_at_micros: i64,
767 metadata: Option<&str>,
768 target_backend: Option<&str>,
769 expected_updated_at_micros: i64,
770 expected_deleted_at_micros: Option<i64>,
771) -> SqlStatement {
772 SqlStatement {
773 sql: "UPDATE graph_edges SET \
774 source_id = CASE WHEN id = ?2 THEN ?3 ELSE source_id END, \
775 target_id = CASE WHEN id = ?2 THEN ?4 ELSE target_id END, \
776 relation = CASE WHEN id = ?2 THEN ?5 ELSE relation END, \
777 weight = CASE WHEN id = ?2 THEN ?6 ELSE weight END, \
778 updated_at = CASE WHEN id = ?2 THEN ?7 ELSE updated_at END, \
779 deleted_at = CASE WHEN id = ?2 THEN NULL ELSE deleted_at END, \
780 metadata = CASE WHEN id = ?2 THEN ?8 ELSE metadata END, \
781 target_backend = CASE WHEN id = ?2 THEN ?9 ELSE target_backend END \
782 WHERE namespace = ?1 \
783 AND ( \
784 (id = ?2 AND changes() = 0 AND updated_at = ?10 AND deleted_at IS ?11 \
785 AND ?7 > updated_at) \
786 OR (source_id = ?3 AND target_id = ?4 AND relation = ?5 \
787 AND id != ?2 AND changes() = 1) \
788 )"
789 .to_string(),
790 params: vec![
791 SqlValue::Text(namespace.to_string()),
792 SqlValue::Text(id.to_string()),
793 SqlValue::Text(canon_src.to_string()),
794 SqlValue::Text(canon_tgt.to_string()),
795 SqlValue::Text(relation.to_string()),
796 SqlValue::Float(weight),
797 SqlValue::Integer(updated_at_micros),
798 match metadata {
799 Some(m) => SqlValue::Text(m.to_string()),
800 None => SqlValue::Null,
801 },
802 match target_backend {
803 Some(b) => SqlValue::Text(b.to_string()),
804 None => SqlValue::Null,
805 },
806 SqlValue::Integer(expected_updated_at_micros),
807 match expected_deleted_at_micros {
808 Some(value) => SqlValue::Integer(value),
809 None => SqlValue::Null,
810 },
811 ],
812 label: Some("edge-symmetric-absorb-or-update-inplace".to_string()),
813 }
814}
815
816#[doc(hidden)]
820pub struct GraphDocumentGuard {
821 pub namespace: String,
822 pub id: Uuid,
823 pub expected_blob_ref: String,
824}
825
826#[doc(hidden)]
830pub struct GraphEdgeSnapshotGuard {
831 pub namespace: String,
832 pub source_id: Uuid,
833 pub target_id: Uuid,
834 pub relation: EdgeRelation,
835 pub expected: Option<Edge>,
836}
837
838#[doc(hidden)]
842#[derive(Default)]
843pub struct GraphMutationPreconditions {
844 pub document: Option<GraphDocumentGuard>,
845 pub edges: Vec<GraphEdgeSnapshotGuard>,
846}
847
848#[doc(hidden)]
852pub enum GraphMutationRequest {
853 Single {
854 request: EdgeUpsertRequest,
855 guard_endpoints: bool,
856 },
857 Batch {
858 requests: Vec<EdgeUpsertRequest>,
859 guard_endpoints: bool,
860 },
861 CommitAnnotation {
862 edge: Edge,
863 guard: CommitAnnotationGuard,
864 },
865}
866
867#[doc(hidden)]
871pub enum GraphMutationOutcome {
872 Single(GuardedEdgeUpsertOutcome),
873 Batch(GuardedEdgeBatchUpsertOutcome),
874 CommitAnnotation(CommitAnnotationInsertOutcome),
875}
876
877#[doc(hidden)]
881pub struct GraphMutationEventOutcome {
882 pub mutation: GraphMutationOutcome,
883 pub retired: Vec<Edge>,
884}
885
886const GRAPH_MUTATION_EVENTS_OP: &str = "compose_graph_mutation_events";
887
888fn count_graph_mutation_events(outcome: &GraphMutationEventOutcome) {
889 let written = match &outcome.mutation {
890 GraphMutationOutcome::Single(GuardedEdgeUpsertOutcome::Written(_)) => 1,
891 GraphMutationOutcome::Batch(outcome) if outcome.refusal.is_none() => outcome.rows.len(),
892 GraphMutationOutcome::CommitAnnotation(CommitAnnotationInsertOutcome::Created(_)) => 1,
893 _ => 0,
894 };
895 khive_storage::usage::count(
896 khive_storage::usage::UsageUnit::EventRows,
897 (written + outcome.retired.len()) as u64,
898 );
899}
900
901#[doc(hidden)]
905pub async fn compose_graph_mutation_events<F>(
906 backend: &crate::StorageBackend,
907 mutation: GraphMutationRequest,
908 preconditions: GraphMutationPreconditions,
909 retirements: Vec<Edge>,
910 make_events: F,
911) -> StorageResult<GraphMutationEventOutcome>
912where
913 F: FnOnce(&GraphMutationEventOutcome) -> StorageResult<Vec<Event>> + Send + 'static,
914{
915 if backend.is_read_only() {
916 return Err(StorageError::Pool {
917 operation: GRAPH_MUTATION_EVENTS_OP.into(),
918 message: "backend is read-only".into(),
919 });
920 }
921 backend
924 .graph()
925 .map_err(|error| map_sqlite_err(error, GRAPH_MUTATION_EVENTS_OP))?;
926 backend
927 .events()
928 .map_err(|error| map_sqlite_err(error, GRAPH_MUTATION_EVENTS_OP))?;
929 let pool = backend.pool_arc();
930 if let Some(writer_task) = pool.writer_task_for_write(None, GRAPH_MUTATION_EVENTS_OP)? {
931 return writer_task
932 .send_bounded_at(
933 move |conn| {
934 graph_mutation_events_enlisted(
935 conn,
936 mutation,
937 preconditions,
938 retirements,
939 make_events,
940 )
941 },
942 GRAPH_MUTATION_EVENTS_OP,
943 )
944 .await
945 .inspect(count_graph_mutation_events)
946 .inspect_err(|error| khive_storage::usage::account_event_write(Err(error)));
947 }
948 pool.record_direct_route(crate::timeout_sink::Site::DirectRouteGraphGeneralWrite);
949 let is_file_backed = backend.is_file_backed();
950 tokio::task::spawn_blocking(move || {
951 if is_file_backed {
952 let admission = pool.write_admission();
953 let _lease = admission
954 .acquire()
955 .map_err(|error| map_sqlite_err(error, GRAPH_MUTATION_EVENTS_OP))?;
956 let conn = pool
957 .open_standalone_writer_for_admitted_operation()
958 .map_err(|error| map_sqlite_err(error, GRAPH_MUTATION_EVENTS_OP))?;
959 run_graph_mutation_transaction(&pool, &conn, false, move |conn| {
960 graph_mutation_events_enlisted(
961 conn,
962 mutation,
963 preconditions,
964 retirements,
965 make_events,
966 )
967 })
968 } else {
969 let guard = pool
970 .writer_for_admitted_operation()
971 .map_err(|error| map_sqlite_err(error, GRAPH_MUTATION_EVENTS_OP))?;
972 run_graph_mutation_transaction(&pool, guard.conn(), true, move |conn| {
973 graph_mutation_events_enlisted(
974 conn,
975 mutation,
976 preconditions,
977 retirements,
978 make_events,
979 )
980 })
981 }
982 })
983 .await
984 .map_err(|error| {
985 StorageError::driver(StorageCapability::Graph, GRAPH_MUTATION_EVENTS_OP, error)
986 })?
987 .inspect(count_graph_mutation_events)
988 .inspect_err(|error| khive_storage::usage::account_event_write(Err(error)))
989}
990
991fn graph_mutation_conflict(message: &'static str) -> StorageError {
992 StorageError::Conflict {
993 capability: StorageCapability::Graph,
994 operation: GRAPH_MUTATION_EVENTS_OP.into(),
995 message: message.into(),
996 }
997}
998
999fn edge_snapshot_matches(actual: &Edge, expected: &Edge) -> bool {
1000 actual.id == expected.id
1001 && actual.namespace == expected.namespace
1002 && actual.source_id == expected.source_id
1003 && actual.target_id == expected.target_id
1004 && actual.relation == expected.relation
1005 && actual.weight.to_bits() == expected.weight.to_bits()
1006 && actual.metadata == expected.metadata
1007 && actual.target_backend == expected.target_backend
1008 && actual.created_at == expected.created_at
1009 && actual.updated_at == expected.updated_at
1010 && actual.deleted_at == expected.deleted_at
1011}
1012
1013fn check_graph_mutation_preconditions(
1014 conn: &rusqlite::Connection,
1015 preconditions: &GraphMutationPreconditions,
1016 retirements: &[Edge],
1017) -> StorageResult<()> {
1018 if let Some(document) = &preconditions.document {
1019 let matches: bool = conn
1020 .query_row(
1021 "SELECT EXISTS(SELECT 1 FROM entities WHERE id=?1 AND namespace=?2 \
1022 AND deleted_at IS NULL AND json_type(properties, '$.blob_ref')='text' \
1023 AND json_extract(properties, '$.blob_ref')=?3 COLLATE BINARY)",
1024 rusqlite::params![
1025 document.id.to_string(),
1026 &document.namespace,
1027 &document.expected_blob_ref,
1028 ],
1029 |row| row.get(0),
1030 )
1031 .map_err(|error| map_err(error, GRAPH_MUTATION_EVENTS_OP))?;
1032 if !matches {
1033 return Err(graph_mutation_conflict(
1034 "source document body reference changed",
1035 ));
1036 }
1037 }
1038 for guard in &preconditions.edges {
1039 let actual = edge_by_natural_key_parts_including_deleted(
1040 conn,
1041 &guard.namespace,
1042 guard.source_id,
1043 guard.target_id,
1044 guard.relation,
1045 )
1046 .map_err(|error| map_err(error, GRAPH_MUTATION_EVENTS_OP))?;
1047 let matches = match (actual.as_ref(), guard.expected.as_ref()) {
1048 (None, None) => true,
1049 (Some(actual), Some(expected)) => edge_snapshot_matches(actual, expected),
1050 _ => false,
1051 };
1052 if !matches {
1053 return Err(graph_mutation_conflict("edge ownership snapshot changed"));
1054 }
1055 }
1056 for expected in retirements {
1057 let actual = edge_by_natural_key_including_deleted(conn, expected)
1058 .map_err(|error| map_err(error, GRAPH_MUTATION_EVENTS_OP))?;
1059 if expected.deleted_at.is_some()
1060 || !actual
1061 .as_ref()
1062 .is_some_and(|actual| edge_snapshot_matches(actual, expected))
1063 {
1064 return Err(graph_mutation_conflict("edge retirement snapshot changed"));
1065 }
1066 }
1067 Ok(())
1068}
1069
1070fn graph_mutation_events_enlisted<F>(
1072 conn: &rusqlite::Connection,
1073 mutation: GraphMutationRequest,
1074 preconditions: GraphMutationPreconditions,
1075 retirements: Vec<Edge>,
1076 make_events: F,
1077) -> StorageResult<GraphMutationEventOutcome>
1078where
1079 F: FnOnce(&GraphMutationEventOutcome) -> StorageResult<Vec<Event>>,
1080{
1081 if !retirements.is_empty() && !matches!(&mutation, GraphMutationRequest::Batch { .. }) {
1082 return Err(StorageError::InvalidInput {
1083 capability: StorageCapability::Graph,
1084 operation: GRAPH_MUTATION_EVENTS_OP.into(),
1085 message: "retirements require a batch mutation".into(),
1086 });
1087 }
1088 check_graph_mutation_preconditions(conn, &preconditions, &retirements)?;
1089 let mutation = match mutation {
1090 GraphMutationRequest::Single {
1091 request,
1092 guard_endpoints,
1093 } => GraphMutationOutcome::Single(
1094 observed_edge_upsert(conn, &request, guard_endpoints)
1095 .map_err(|error| map_err(error, GRAPH_MUTATION_EVENTS_OP))?,
1096 ),
1097 GraphMutationRequest::Batch {
1098 requests,
1099 guard_endpoints,
1100 } => {
1101 let outcome = observed_edge_batch_upsert(conn, &requests, guard_endpoints)
1102 .map_err(|error| map_err(error, GRAPH_MUTATION_EVENTS_OP))?;
1103 if outcome.refusal.is_none() && outcome.rows.len() != requests.len() {
1104 return Err(graph_mutation_conflict(
1105 "edge result count differs from request count",
1106 ));
1107 }
1108 GraphMutationOutcome::Batch(outcome)
1109 }
1110 GraphMutationRequest::CommitAnnotation { edge, guard } => {
1111 if edge.relation != EdgeRelation::Annotates || edge.deleted_at.is_some() {
1112 return Err(StorageError::InvalidInput {
1113 capability: StorageCapability::Graph,
1114 operation: GRAPH_MUTATION_EVENTS_OP.into(),
1115 message: "expected a live annotates edge".into(),
1116 });
1117 }
1118 GraphMutationOutcome::CommitAnnotation(
1119 conditional_commit_annotation_insert(conn, edge, &guard)
1120 .map_err(|error| map_err(error, GRAPH_MUTATION_EVENTS_OP))?,
1121 )
1122 }
1123 };
1124 let written = match &mutation {
1125 GraphMutationOutcome::Single(GuardedEdgeUpsertOutcome::Written(_)) => 1,
1126 GraphMutationOutcome::Batch(outcome) if outcome.refusal.is_none() => outcome.rows.len(),
1127 GraphMutationOutcome::CommitAnnotation(CommitAnnotationInsertOutcome::Created(_)) => 1,
1128 _ => {
1129 return Ok(GraphMutationEventOutcome {
1130 mutation,
1131 retired: Vec::new(),
1132 })
1133 }
1134 };
1135 let mut retired = Vec::with_capacity(retirements.len());
1136 for edge in retirements {
1137 let statement =
1138 edge_soft_delete_statement(Uuid::from(edge.id), Utc::now().timestamp_micros());
1139 let mut stmt = conn
1140 .prepare(&statement.sql)
1141 .map_err(|error| map_err(error, GRAPH_MUTATION_EVENTS_OP))?;
1142 bind_params(&mut stmt, &statement.params)
1143 .map_err(|error| map_err(error, GRAPH_MUTATION_EVENTS_OP))?;
1144 if stmt
1145 .raw_execute()
1146 .map_err(|error| map_err(error, GRAPH_MUTATION_EVENTS_OP))?
1147 != 1
1148 {
1149 return Err(graph_mutation_conflict(
1150 "edge retirement changed during mutation",
1151 ));
1152 }
1153 retired.push(edge);
1154 }
1155 let outcome = GraphMutationEventOutcome { mutation, retired };
1156 let expected_events = written + outcome.retired.len();
1157 if expected_events == 0 {
1158 return Ok(outcome);
1159 }
1160 let events = make_events(&outcome)?;
1161 if events.len() != expected_events {
1162 return Err(graph_mutation_conflict(
1163 "event result count differs from mutation count",
1164 ));
1165 }
1166 for event in &events {
1167 super::event::append_event_in_transaction(conn, event).map_err(|error| {
1168 StorageError::driver(StorageCapability::Events, GRAPH_MUTATION_EVENTS_OP, error)
1169 })?;
1170 }
1171 Ok(outcome)
1172}
1173
1174pub struct SqlGraphStore {
1176 pool: Arc<ConnectionPool>,
1177 index_repair: Option<super::index_repair::IndexRepairContext>,
1178 is_file_backed: bool,
1179 namespace: String,
1183 writer_task: Option<WriterTaskHandle>,
1184}
1185
1186impl SqlGraphStore {
1187 pub fn new_scoped(
1195 pool: Arc<ConnectionPool>,
1196 is_file_backed: bool,
1197 namespace: impl Into<String>,
1198 ) -> Self {
1199 let writer_task = pool.writer_task_handle().ok().flatten();
1205
1206 Self {
1207 pool,
1208 index_repair: None,
1209 is_file_backed,
1210 namespace: namespace.into(),
1211 writer_task,
1212 }
1213 }
1214
1215 pub(crate) fn with_index_repair(
1216 mut self,
1217 repair: super::index_repair::IndexRepairContext,
1218 ) -> Self {
1219 self.index_repair = Some(repair);
1220 self
1221 }
1222
1223 async fn with_indexed_reader<F, R>(&self, op: &'static str, read: F) -> Result<R, StorageError>
1224 where
1225 F: FnMut(&rusqlite::Connection) -> Result<R, rusqlite::Error> + Send + 'static,
1226 R: Send + 'static,
1227 {
1228 super::index_repair::run_indexed_read(
1229 Arc::clone(&self.pool),
1230 self.index_repair.clone(),
1231 StorageCapability::Graph,
1232 op,
1233 read,
1234 )
1235 .await
1236 }
1237
1238 fn current_writer_task(
1239 &self,
1240 operation: &'static str,
1241 ) -> Result<Option<WriterTaskHandle>, StorageError> {
1242 self.pool
1243 .writer_task_for_write(self.writer_task.as_ref(), operation)
1244 }
1245
1246 async fn with_writer<F, R>(&self, op: &'static str, f: F) -> Result<R, StorageError>
1253 where
1254 F: FnOnce(&rusqlite::Connection) -> Result<R, rusqlite::Error> + Send + 'static,
1255 R: Send + 'static,
1256 {
1257 if let Some(writer_task) = self.current_writer_task(op)? {
1258 return writer_task
1259 .send_bounded_at(move |conn| f(conn).map_err(|e| map_err(e, op)), op)
1260 .await;
1261 }
1262
1263 self.pool
1264 .record_direct_route(crate::timeout_sink::Site::DirectRouteGraphGeneralWrite);
1265 let pool = Arc::clone(&self.pool);
1266 let db = crate::timeout_sink::db_label(&pool);
1267 let is_file_backed = self.is_file_backed;
1268 tokio::task::spawn_blocking(move || {
1269 let result =
1270 pool.execute_direct_transaction(StorageCapability::Graph, op, move |conn| {
1271 f(conn).map_err(|error| map_err(error, op))
1272 });
1273 if is_file_backed {
1274 if let Err(error) = &result {
1275 crate::timeout_sink::maybe_emit_busy_storage_error(
1276 &db,
1277 crate::timeout_sink::Site::StandaloneGraph,
1278 error,
1279 );
1280 }
1281 }
1282 result
1283 })
1284 .await
1285 .map_err(|e| StorageError::driver(StorageCapability::Graph, op, e))?
1286 }
1287
1288 async fn observed_edge_write(
1289 &self,
1290 operation: &'static str,
1291 request: EdgeUpsertRequest,
1292 guard_endpoints: bool,
1293 ) -> Result<GuardedEdgeUpsertOutcome, 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_upsert(conn, &request, 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_upsert(conn, &request, guard_endpoints)
1308 })
1309 .await
1310 }
1311
1312 async fn observed_edge_batch_write(
1313 &self,
1314 operation: &'static str,
1315 requests: Vec<EdgeUpsertRequest>,
1316 guard_endpoints: bool,
1317 ) -> Result<GuardedEdgeBatchUpsertOutcome, StorageError> {
1318 if let Some(writer_task) = self.current_writer_task(operation)? {
1319 return writer_task
1320 .send_bounded_at(
1321 move |conn| {
1322 observed_edge_batch_upsert(conn, &requests, guard_endpoints)
1323 .map_err(|error| map_err(error, operation))
1324 },
1325 operation,
1326 )
1327 .await;
1328 }
1329
1330 self.with_writer(operation, move |conn| {
1331 observed_edge_batch_upsert(conn, &requests, guard_endpoints)
1332 })
1333 .await
1334 }
1335
1336 async fn with_reader<F, R>(&self, op: &'static str, f: F) -> Result<R, StorageError>
1337 where
1338 F: FnOnce(&rusqlite::Connection) -> Result<R, rusqlite::Error> + Send + 'static,
1339 R: Send + 'static,
1340 {
1341 super::run_pooled_store_read(
1342 Arc::clone(&self.pool),
1343 StorageCapability::Graph,
1344 op,
1345 move |conn| f(conn).map_err(|error| map_err(error, op)),
1346 )
1347 .await
1348 }
1349}
1350
1351fn report_graph_usage(queries: &std::sync::atomic::AtomicU64, rows: &std::sync::atomic::AtomicU64) {
1367 khive_storage::usage::count(
1368 khive_storage::usage::UsageUnit::DbRoundTrips,
1369 queries.load(std::sync::atomic::Ordering::Relaxed),
1370 );
1371 khive_storage::usage::count(
1372 khive_storage::usage::UsageUnit::GraphHops,
1373 rows.load(std::sync::atomic::Ordering::Relaxed),
1374 );
1375}
1376
1377fn read_edge(row: &rusqlite::Row<'_>) -> Result<Edge, rusqlite::Error> {
1378 let namespace: String = row.get(0)?;
1379 let id_str: String = row.get(1)?;
1380 let source_str: String = row.get(2)?;
1381 let target_str: String = row.get(3)?;
1382 let relation_str: String = row.get(4)?;
1383 let weight: f64 = row.get(5)?;
1384 let created_micros: i64 = row.get(6)?;
1385 let updated_micros: i64 = row.get(7)?;
1386 let deleted_micros: Option<i64> = row.get(8)?;
1387 let metadata_str: Option<String> = row.get(9)?;
1388 let target_backend: Option<String> = row.get(10)?;
1389
1390 let id = parse_uuid(&id_str)?;
1391 let source_id = parse_uuid(&source_str)?;
1392 let target_id = parse_uuid(&target_str)?;
1393 let created_at = micros_to_datetime(created_micros);
1394 let relation = relation_str.parse::<EdgeRelation>().map_err(|e| {
1395 rusqlite::Error::FromSqlConversionFailure(4, rusqlite::types::Type::Text, Box::new(e))
1396 })?;
1397 let metadata = match metadata_str {
1398 Some(s) => {
1399 let v = serde_json::from_str(&s).map_err(|e| {
1400 rusqlite::Error::FromSqlConversionFailure(
1401 9,
1402 rusqlite::types::Type::Text,
1403 Box::new(e),
1404 )
1405 })?;
1406 Some(v)
1407 }
1408 None => None,
1409 };
1410
1411 Ok(Edge {
1412 id: id.into(),
1413 namespace,
1414 source_id,
1415 target_id,
1416 relation,
1417 weight,
1418 created_at,
1419 updated_at: micros_to_datetime(updated_micros),
1420 deleted_at: deleted_micros.map(micros_to_datetime),
1421 metadata,
1422 target_backend,
1423 })
1424}
1425
1426fn parse_uuid(s: &str) -> Result<Uuid, rusqlite::Error> {
1427 Uuid::parse_str(s).map_err(|e| {
1428 rusqlite::Error::FromSqlConversionFailure(0, rusqlite::types::Type::Text, Box::new(e))
1429 })
1430}
1431
1432fn neighbor_extra_clause(
1438 query: &NeighborQuery,
1439 start_param_idx: usize,
1440 after: Option<&NeighborCursor>,
1441 neighbor_kinds: Option<&[String]>,
1442) -> (String, String, Vec<Box<dyn rusqlite::types::ToSql>>) {
1443 let mut conditions: Vec<String> = Vec::new();
1444 let mut extra_params: Vec<Box<dyn rusqlite::types::ToSql>> = Vec::new();
1445 let mut param_idx = start_param_idx;
1446
1447 if let Some(ref rels) = query.relations {
1448 if !rels.is_empty() {
1449 let placeholders: Vec<String> = rels
1450 .iter()
1451 .map(|r| {
1452 extra_params.push(Box::new(r.to_string()));
1453 let p = format!("?{}", param_idx);
1454 param_idx += 1;
1455 p
1456 })
1457 .collect();
1458 conditions.push(format!("relation IN ({})", placeholders.join(",")));
1459 }
1460 }
1461
1462 if let Some(min_w) = query.min_weight {
1463 extra_params.push(Box::new(min_w));
1464 conditions.push(format!("weight >= ?{}", param_idx));
1465 param_idx += 1;
1466 }
1467
1468 if let Some(cursor) = after {
1469 extra_params.push(Box::new(cursor.weight));
1470 let weight_idx = param_idx;
1471 param_idx += 1;
1472 extra_params.push(Box::new(cursor.node_id.to_string()));
1473 let node_idx = param_idx;
1474 param_idx += 1;
1475 extra_params.push(Box::new(cursor.edge_id.to_string()));
1476 let edge_idx = param_idx;
1477 param_idx += 1;
1478 conditions.push(format!(
1479 "(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}))"
1480 ));
1481 }
1482
1483 if let Some(kinds) = neighbor_kinds.filter(|kinds| !kinds.is_empty()) {
1484 let placeholders: Vec<String> = kinds
1485 .iter()
1486 .map(|kind| {
1487 extra_params.push(Box::new(kind.clone()));
1488 let p = format!("?{param_idx}");
1489 param_idx += 1;
1490 p
1491 })
1492 .collect();
1493 let entity_placeholders = placeholders.join(",");
1494 let note_placeholders: Vec<String> = kinds
1495 .iter()
1496 .map(|kind| {
1497 extra_params.push(Box::new(kind.clone()));
1498 let p = format!("?{param_idx}");
1499 param_idx += 1;
1500 p
1501 })
1502 .collect();
1503 conditions.push(format!(
1504 "(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 ({})))",
1505 note_placeholders.join(",")
1506 ));
1507 }
1508
1509 let where_extra = if conditions.is_empty() {
1510 String::new()
1511 } else {
1512 format!(" WHERE {}", conditions.join(" AND "))
1513 };
1514
1515 let limit_clause = if let Some(lim) = query.limit {
1516 extra_params.push(Box::new(lim as i64));
1517 format!(" LIMIT ?{}", param_idx)
1518 } else {
1519 String::new()
1520 };
1521
1522 (where_extra, limit_clause, extra_params)
1523}
1524
1525#[cfg(test)]
1533static NEIGHBOR_SELECT_COUNT: std::sync::atomic::AtomicUsize =
1534 std::sync::atomic::AtomicUsize::new(0);
1535
1536#[cfg(test)]
1537fn count_neighbor_select() {
1538 NEIGHBOR_SELECT_COUNT.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1539}
1540
1541#[cfg(not(test))]
1542fn count_neighbor_select() {}
1543
1544#[cfg(test)]
1545pub(crate) fn reset_neighbor_select_count() {
1546 NEIGHBOR_SELECT_COUNT.store(0, std::sync::atomic::Ordering::Relaxed);
1547}
1548
1549#[cfg(test)]
1550pub(crate) fn neighbor_select_count() -> usize {
1551 NEIGHBOR_SELECT_COUNT.load(std::sync::atomic::Ordering::Relaxed)
1552}
1553
1554fn micros_to_datetime(micros: i64) -> DateTime<Utc> {
1555 Utc.timestamp_micros(micros)
1556 .single()
1557 .unwrap_or_else(Utc::now)
1558}
1559
1560fn edge_order_clause(sort: &[SortOrder<EdgeSortField>]) -> String {
1567 if sort.is_empty() {
1568 return " ORDER BY created_at DESC, id DESC".to_string();
1569 }
1570 let mut parts: Vec<String> = sort
1571 .iter()
1572 .map(|s| {
1573 let dir = match s.direction {
1574 SortDirection::Asc => "ASC",
1575 SortDirection::Desc => "DESC",
1576 };
1577 format!("{} {}", edge_sort_col(&s.field), dir)
1578 })
1579 .collect();
1580 let dir = match sort.last().map(|s| &s.direction) {
1581 Some(SortDirection::Asc) => "ASC",
1582 _ => "DESC",
1583 };
1584 parts.push(format!("id {dir}"));
1585 format!(" ORDER BY {}", parts.join(", "))
1586}
1587
1588fn endpoint_base_case(column: &str) -> String {
1606 format!(
1607 "CASE WHEN EXISTS (SELECT 1 FROM entities be WHERE be.id = graph_edges.{column}) \
1608 THEN 'entity' \
1609 WHEN EXISTS (SELECT 1 FROM notes bn WHERE bn.id = graph_edges.{column}) \
1610 THEN 'note' ELSE 'none' END"
1611 )
1612}
1613
1614fn fold_endpoint_base_row(counts: &mut EdgeEndpointBaseCounts, source: &str, target: &str, n: u64) {
1618 let slot = match (source, target) {
1619 ("entity", "entity") => &mut counts.entity_entity,
1620 ("entity", "note") => &mut counts.entity_note,
1621 ("note", "entity") => &mut counts.note_entity,
1622 ("note", "note") => &mut counts.note_note,
1623 _ => &mut counts.unresolved,
1624 };
1625 *slot = slot.saturating_add(n);
1626}
1627
1628const LIVE_ENDPOINTS_CONDITION: &str = "NOT EXISTS (SELECT 1 FROM entities le \
1641 WHERE le.id = graph_edges.source_id AND le.deleted_at IS NOT NULL) \
1642 AND NOT EXISTS (SELECT 1 FROM entities le \
1643 WHERE le.id = graph_edges.target_id AND le.deleted_at IS NOT NULL) \
1644 AND NOT EXISTS (SELECT 1 FROM notes ln \
1645 WHERE ln.id = graph_edges.source_id AND ln.deleted_at IS NOT NULL) \
1646 AND NOT EXISTS (SELECT 1 FROM notes ln \
1647 WHERE ln.id = graph_edges.target_id AND ln.deleted_at IS NOT NULL)";
1648
1649fn with_live_endpoints(where_clause: &str) -> String {
1652 format!("{where_clause} AND {LIVE_ENDPOINTS_CONDITION}")
1653}
1654
1655fn build_edge_filter_sql(
1656 namespace: &str,
1657 filter: &EdgeFilter,
1658) -> (String, Vec<Box<dyn rusqlite::types::ToSql>>) {
1659 build_edge_filter_sql_for_namespaces(&[namespace.to_string()], filter)
1660}
1661
1662fn build_edge_filter_sql_for_namespaces(
1663 namespaces: &[String],
1664 filter: &EdgeFilter,
1665) -> (String, Vec<Box<dyn rusqlite::types::ToSql>>) {
1666 let params: Vec<Box<dyn rusqlite::types::ToSql>> = namespaces
1667 .iter()
1668 .map(|namespace| -> Box<dyn rusqlite::types::ToSql> { Box::new(namespace.clone()) })
1669 .collect();
1670 let namespace_condition = match namespaces.len() {
1671 0 => "0".to_string(),
1672 1 => "namespace = ?1".to_string(),
1673 _ => {
1674 let placeholders: Vec<String> =
1675 (1..=namespaces.len()).map(|i| format!("?{i}")).collect();
1676 format!("namespace IN ({})", placeholders.join(", "))
1677 }
1678 };
1679 build_edge_filter_conditions(namespace_condition, params, filter)
1680}
1681
1682fn build_edge_filter_sql_for_namespaces_json(
1697 namespaces_json: &str,
1698 filter: &EdgeFilter,
1699) -> (String, Vec<Box<dyn rusqlite::types::ToSql>>) {
1700 let params: Vec<Box<dyn rusqlite::types::ToSql>> = vec![Box::new(namespaces_json.to_string())];
1701 let namespace_condition = "namespace IN (SELECT value FROM json_each(?1))".to_string();
1702 build_edge_filter_conditions(namespace_condition, params, filter)
1703}
1704
1705fn build_edge_filter_conditions(
1706 namespace_condition: String,
1707 mut params: Vec<Box<dyn rusqlite::types::ToSql>>,
1708 filter: &EdgeFilter,
1709) -> (String, Vec<Box<dyn rusqlite::types::ToSql>>) {
1710 let mut conditions = vec![namespace_condition, "deleted_at IS NULL".to_string()];
1711
1712 if !filter.ids.is_empty() {
1713 let placeholders: Vec<String> = filter
1714 .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!("id IN ({})", placeholders.join(",")));
1722 }
1723
1724 if !filter.source_ids.is_empty() {
1725 let placeholders: Vec<String> = filter
1726 .source_ids
1727 .iter()
1728 .map(|id| {
1729 params.push(Box::new(id.to_string()));
1730 format!("?{}", params.len())
1731 })
1732 .collect();
1733 conditions.push(format!("source_id IN ({})", placeholders.join(",")));
1734 }
1735
1736 if !filter.target_ids.is_empty() {
1737 let placeholders: Vec<String> = filter
1738 .target_ids
1739 .iter()
1740 .map(|id| {
1741 params.push(Box::new(id.to_string()));
1742 format!("?{}", params.len())
1743 })
1744 .collect();
1745 conditions.push(format!("target_id IN ({})", placeholders.join(",")));
1746 }
1747
1748 if !filter.relations.is_empty() {
1749 let placeholders: Vec<String> = filter
1750 .relations
1751 .iter()
1752 .map(|r| {
1753 params.push(Box::new(r.to_string()));
1754 format!("?{}", params.len())
1755 })
1756 .collect();
1757 conditions.push(format!("relation IN ({})", placeholders.join(",")));
1758 }
1759
1760 if let Some(min_w) = filter.min_weight {
1761 params.push(Box::new(min_w));
1762 conditions.push(format!("weight >= ?{}", params.len()));
1763 }
1764
1765 if let Some(max_w) = filter.max_weight {
1766 params.push(Box::new(max_w));
1767 conditions.push(format!("weight <= ?{}", params.len()));
1768 }
1769
1770 if let Some(ref time_range) = filter.created_at {
1771 if let Some(start) = time_range.start {
1772 params.push(Box::new(start.timestamp_micros()));
1773 conditions.push(format!("created_at >= ?{}", params.len()));
1774 }
1775 if let Some(end) = time_range.end {
1776 params.push(Box::new(end.timestamp_micros()));
1777 conditions.push(format!("created_at < ?{}", params.len()));
1778 }
1779 }
1780
1781 let clause = format!(" WHERE {}", conditions.join(" AND "));
1782 (clause, params)
1783}
1784
1785fn edge_sort_col(field: &EdgeSortField) -> &'static str {
1786 match field {
1787 EdgeSortField::CreatedAt => "created_at",
1788 EdgeSortField::Weight => "weight",
1789 EdgeSortField::Relation => "relation",
1790 }
1791}
1792
1793fn canonical_edge_endpoints(
1802 relation: EdgeRelation,
1803 source_id: Uuid,
1804 target_id: Uuid,
1805) -> (Uuid, Uuid) {
1806 relation.canonical_endpoints(source_id, target_id)
1807}
1808
1809fn edge_endpoints_exist(
1817 conn: &rusqlite::Connection,
1818 source_id: Uuid,
1819 target_id: Uuid,
1820) -> Result<MissingEndpoints, rusqlite::Error> {
1821 let src_exists = endpoint_exists_clause("?1");
1822 let tgt_exists = endpoint_exists_clause("?2");
1823 let sql = format!("SELECT ({src_exists}), ({tgt_exists})");
1824 conn.query_row(
1825 &sql,
1826 rusqlite::params![source_id.to_string(), target_id.to_string()],
1827 |row| {
1828 let src_exists: bool = row.get(0)?;
1829 let tgt_exists: bool = row.get(1)?;
1830 Ok(MissingEndpoints {
1831 source: !src_exists,
1832 target: !tgt_exists,
1833 })
1834 },
1835 )
1836}
1837
1838fn edge_by_natural_key_including_deleted(
1839 conn: &rusqlite::Connection,
1840 edge: &Edge,
1841) -> Result<Option<Edge>, rusqlite::Error> {
1842 edge_by_natural_key_parts_including_deleted(
1843 conn,
1844 &edge.namespace,
1845 edge.source_id,
1846 edge.target_id,
1847 edge.relation,
1848 )
1849}
1850
1851fn edge_by_natural_key_parts_including_deleted(
1852 conn: &rusqlite::Connection,
1853 namespace: &str,
1854 source_id: Uuid,
1855 target_id: Uuid,
1856 relation: EdgeRelation,
1857) -> Result<Option<Edge>, rusqlite::Error> {
1858 let (source_id, target_id) = canonical_edge_endpoints(relation, source_id, target_id);
1859 conn.query_row(
1860 "SELECT namespace, id, source_id, target_id, relation, weight, \
1861 created_at, updated_at, deleted_at, metadata, target_backend \
1862 FROM graph_edges \
1863 WHERE namespace = ?1 AND source_id = ?2 AND target_id = ?3 AND relation = ?4",
1864 rusqlite::params![
1865 namespace,
1866 source_id.to_string(),
1867 target_id.to_string(),
1868 relation.as_str(),
1869 ],
1870 read_edge,
1871 )
1872 .optional()
1873}
1874
1875fn conditional_commit_annotation_insert(
1880 conn: &rusqlite::Connection,
1881 edge: Edge,
1882 guard: &CommitAnnotationGuard,
1883) -> Result<CommitAnnotationInsertOutcome, rusqlite::Error> {
1884 let cursor_matches: bool = conn.query_row(
1888 "SELECT (SELECT EXISTS(SELECT 1 FROM git_mirror_cursor \
1889 WHERE project_id=?1 AND kind='commits' \
1890 AND typeof(cursor_value)='text' \
1891 AND CAST(cursor_value AS BLOB)=?2 AND updated_at=?3)) \
1892 AND (SELECT EXISTS(SELECT 1 FROM git_mirror_cursor \
1893 WHERE project_id=?1 AND kind='commits_checkpoint' \
1894 AND typeof(cursor_value)='text' \
1895 AND CAST(cursor_value AS BLOB)=?4 AND updated_at=?5))",
1896 rusqlite::params![
1897 edge.target_id.to_string(),
1898 guard.commits.value,
1899 guard.commits.updated_at,
1900 guard.checkpoint.value,
1901 guard.checkpoint.updated_at,
1902 ],
1903 |row| row.get(0),
1904 )?;
1905 let metadata = edge
1906 .metadata
1907 .as_ref()
1908 .map(serde_json::to_string)
1909 .transpose()
1910 .map_err(|error| rusqlite::Error::ToSqlConversionFailure(Box::new(error)))?;
1911 let affected = conn.execute(
1912 include_str!("../../sql/commit-annotation-insert.sql"),
1913 rusqlite::params![
1914 edge.namespace,
1915 Uuid::from(edge.id).to_string(),
1916 edge.source_id.to_string(),
1917 edge.target_id.to_string(),
1918 edge.weight,
1919 edge.created_at.timestamp_micros(),
1920 edge.updated_at.timestamp_micros(),
1921 metadata,
1922 guard.expected_sha,
1923 guard.source_identity,
1924 cursor_matches,
1925 ],
1926 )?;
1927 if affected == 1 {
1928 return Ok(CommitAnnotationInsertOutcome::Created(edge));
1929 }
1930 let source_live: bool = conn.query_row(
1931 include_str!("../../sql/commit-annotation-source-live-select.sql"),
1932 rusqlite::params![
1933 edge.source_id.to_string(),
1934 edge.namespace,
1935 guard.expected_sha
1936 ],
1937 |row| row.get(0),
1938 )?;
1939 if !source_live {
1940 return Ok(CommitAnnotationInsertOutcome::SourceChanged);
1941 }
1942 let target_live: bool = conn.query_row(
1943 "SELECT EXISTS(SELECT 1 FROM entities WHERE id=?1 AND namespace=?2 \
1944 AND kind='project' AND deleted_at IS NULL \
1945 AND json_type(properties, '$.repo_slug')='text' \
1946 AND json_extract(properties, '$.repo_slug')=?3 COLLATE BINARY)",
1947 rusqlite::params![
1948 edge.target_id.to_string(),
1949 edge.namespace,
1950 guard.source_identity
1951 ],
1952 |row| row.get(0),
1953 )?;
1954 if !target_live {
1955 return Ok(CommitAnnotationInsertOutcome::TargetChanged);
1956 }
1957 if !cursor_matches {
1958 return Ok(CommitAnnotationInsertOutcome::CursorChanged);
1959 }
1960 match edge_by_natural_key_including_deleted(conn, &edge)? {
1961 Some(existing) if existing.deleted_at.is_some() => {
1962 Ok(CommitAnnotationInsertOutcome::Tombstoned)
1963 }
1964 Some(_) => Ok(CommitAnnotationInsertOutcome::ExistingLive),
1965 None => Err(rusqlite::Error::QueryReturnedNoRows),
1966 }
1967}
1968
1969fn observed_edge_upsert(
1972 conn: &rusqlite::Connection,
1973 request: &EdgeUpsertRequest,
1974 guard_endpoints: bool,
1975) -> Result<GuardedEdgeUpsertOutcome, rusqlite::Error> {
1976 let (source_id, target_id) = canonical_edge_endpoints(
1977 request.edge.relation,
1978 request.edge.source_id,
1979 request.edge.target_id,
1980 );
1981 if guard_endpoints {
1982 #[cfg(test)]
1985 tests::insert_probe_seam::hook((source_id, target_id));
1986 let missing = edge_endpoints_exist(conn, source_id, target_id)?;
1987 if missing.any() {
1988 return Ok(GuardedEdgeUpsertOutcome::Refused(
1989 EdgeUpsertRefusal::MissingEndpoints(missing),
1990 ));
1991 }
1992 }
1993
1994 let previous = edge_by_natural_key_including_deleted(conn, &request.edge)?;
1995 if let Some(edge) = previous.as_ref() {
1996 if edge.deleted_at.is_some() && !request.resurrect {
1997 return Ok(GuardedEdgeUpsertOutcome::Refused(
1998 EdgeUpsertRefusal::ResurrectionRequired { edge: edge.clone() },
1999 ));
2000 }
2001 }
2002
2003 let statement = edge_upsert_statement_with_resurrection(&request.edge, request.resurrect);
2004 let mut stmt = conn.prepare(&statement.sql)?;
2005 bind_params(&mut stmt, &statement.params)?;
2006 let affected = stmt.raw_execute()?;
2007 if affected == 0 {
2008 let edge = edge_by_natural_key_including_deleted(conn, &request.edge)?
2009 .ok_or(rusqlite::Error::QueryReturnedNoRows)?;
2010 return Ok(GuardedEdgeUpsertOutcome::Refused(
2011 EdgeUpsertRefusal::ResurrectionRequired { edge },
2012 ));
2013 }
2014
2015 let edge = edge_by_natural_key_including_deleted(conn, &request.edge)?
2016 .ok_or(rusqlite::Error::QueryReturnedNoRows)?;
2017 let disposition = match previous.as_ref().and_then(|edge| edge.deleted_at) {
2018 None if previous.is_none() => EdgeUpsertDisposition::Created,
2019 None => EdgeUpsertDisposition::Updated,
2020 Some(_) => EdgeUpsertDisposition::Resurrected,
2021 };
2022 Ok(GuardedEdgeUpsertOutcome::Written(EdgeUpsertResult {
2023 edge,
2024 disposition,
2025 previous,
2026 }))
2027}
2028
2029fn observed_edge_batch_upsert(
2032 conn: &rusqlite::Connection,
2033 requests: &[EdgeUpsertRequest],
2034 guard_endpoints: bool,
2035) -> Result<GuardedEdgeBatchUpsertOutcome, rusqlite::Error> {
2036 for (entry_index, request) in requests.iter().enumerate() {
2040 let (source_id, target_id) = canonical_edge_endpoints(
2041 request.edge.relation,
2042 request.edge.source_id,
2043 request.edge.target_id,
2044 );
2045 if guard_endpoints {
2046 let missing = edge_endpoints_exist(conn, source_id, target_id)?;
2047 if missing.any() {
2048 return Ok(GuardedEdgeBatchUpsertOutcome {
2049 rows: Vec::new(),
2050 refusal: Some(GuardedEdgeBatchRefusal {
2051 entry_index,
2052 reason: EdgeUpsertRefusal::MissingEndpoints(missing),
2053 }),
2054 });
2055 }
2056 }
2057 if let Some(edge) = edge_by_natural_key_including_deleted(conn, &request.edge)? {
2058 if edge.deleted_at.is_some() && !request.resurrect {
2059 return Ok(GuardedEdgeBatchUpsertOutcome {
2060 rows: Vec::new(),
2061 refusal: Some(GuardedEdgeBatchRefusal {
2062 entry_index,
2063 reason: EdgeUpsertRefusal::ResurrectionRequired { edge },
2064 }),
2065 });
2066 }
2067 }
2068 }
2069
2070 let mut rows = Vec::with_capacity(requests.len());
2071 for request in requests {
2072 match observed_edge_upsert(conn, request, false)? {
2073 GuardedEdgeUpsertOutcome::Written(row) => rows.push(row),
2074 GuardedEdgeUpsertOutcome::Refused(_) => {
2075 return Err(rusqlite::Error::ExecuteReturnedResults)
2076 }
2077 }
2078 }
2079 Ok(GuardedEdgeBatchUpsertOutcome {
2080 rows,
2081 refusal: None,
2082 })
2083}
2084
2085fn traversal_neighbor_sql(
2090 direction: Direction,
2091 relation_count: usize,
2092 has_min_weight: bool,
2093) -> String {
2094 let (node_column, endpoint_column, index) = match direction {
2095 Direction::Out => ("target_id", "source_id", "idx_graph_edges_ns_src_rel"),
2096 Direction::In => ("source_id", "target_id", "idx_graph_edges_ns_tgt_rel"),
2097 Direction::Both => unreachable!("Direction::Both is split into indexed Out/In seeks"),
2098 };
2099 let mut sql = format!(
2100 "SELECT {node_column}, id, weight \
2101 FROM graph_edges INDEXED BY {index} \
2102 WHERE namespace = ?1 AND {endpoint_column} = ?2 AND deleted_at IS NULL"
2103 );
2104 if relation_count > 0 {
2105 let placeholders = (0..relation_count)
2106 .map(|offset| format!("?{}", 4 + offset))
2107 .collect::<Vec<_>>()
2108 .join(",");
2109 sql.push_str(&format!(" AND relation IN ({placeholders})"));
2110 }
2111 if has_min_weight {
2112 sql.push_str(&format!(" AND weight >= ?{}", 4 + relation_count));
2113 }
2114 sql.push_str(" LIMIT ?3");
2115 sql
2116}
2117
2118fn traversal_timeout_error(budget: &TraversalExecutionBudget) -> StorageError {
2119 StorageError::Timeout {
2120 operation: format!(
2121 "traverse ({}ms execution budget)",
2122 budget.max_duration().as_millis()
2123 )
2124 .into(),
2125 }
2126}
2127
2128fn traversal_work_error(budget: &TraversalExecutionBudget) -> StorageError {
2129 StorageError::InvalidInput {
2130 capability: StorageCapability::Graph,
2131 operation: "traverse".into(),
2132 message: format!(
2133 "traversal work budget exceeded after {} adjacency rows; \
2134 narrow roots, depth, relations, or result limit",
2135 budget.work_limit()
2136 ),
2137 }
2138}
2139
2140#[derive(Clone, Copy)]
2141struct TraversalFrontierNode {
2142 node_id: Uuid,
2143 depth: usize,
2144 total_weight: f64,
2145}
2146
2147#[allow(clippy::too_many_arguments)]
2148fn run_bounded_traversal(
2149 conn: &rusqlite::Connection,
2150 roots: Vec<Uuid>,
2151 opts: TraversalOptions,
2152 include_roots: bool,
2153 namespace: String,
2154 origin: khive_storage::tx_registry::TxOrigin,
2155 budget: TraversalExecutionBudget,
2156 counted_rows: &std::sync::atomic::AtomicU64,
2157 counted_queries: &std::sync::atomic::AtomicU64,
2158) -> Result<Vec<GraphPath>, StorageError> {
2159 let progress_timed_out = Arc::new(std::sync::atomic::AtomicBool::new(false));
2160 let callback_timed_out = Arc::clone(&progress_timed_out);
2161 let callback_budget = budget.clone();
2162 #[cfg(test)]
2163 let progress_seam_root = roots.first().copied();
2164 conn.progress_handler(
2165 1_000,
2166 Some(move || {
2167 if crate::read_cancellation::current_read_should_interrupt() {
2168 return true;
2169 }
2170 #[cfg(test)]
2171 if tests::traverse_progress_seam::hook(progress_seam_root) {
2172 callback_timed_out.store(true, std::sync::atomic::Ordering::Relaxed);
2173 return true;
2174 }
2175 let expired = callback_budget.is_expired();
2176 if expired {
2177 callback_timed_out.store(true, std::sync::atomic::Ordering::Relaxed);
2178 }
2179 expired
2180 }),
2181 )
2182 .map_err(|e| map_err(e, "traverse_progress_handler"))?;
2183
2184 let result = (|| {
2185 let result_limit = opts.effective_limit() as usize;
2186 let relation_count = opts.relations.as_ref().map_or(0, Vec::len);
2187 let directions = match opts.direction {
2188 Direction::Out => vec![Direction::Out],
2189 Direction::In => vec![Direction::In],
2190 Direction::Both => vec![Direction::Out, Direction::In],
2191 };
2192 let statements = directions
2193 .into_iter()
2194 .map(|direction| {
2195 traversal_neighbor_sql(direction, relation_count, opts.min_weight.is_some())
2196 })
2197 .collect::<Vec<_>>();
2198 let map_sql_error = |error| {
2199 if progress_timed_out.load(std::sync::atomic::Ordering::Relaxed) {
2200 traversal_timeout_error(&budget)
2201 } else {
2202 map_err(error, "traverse")
2203 }
2204 };
2205
2206 let mut all_paths = Vec::with_capacity(roots.len());
2207 for root_id in roots {
2208 let mut seen = HashSet::new();
2209 seen.insert(root_id);
2210 let mut frontier = VecDeque::from([TraversalFrontierNode {
2211 node_id: root_id,
2212 depth: 0,
2213 total_weight: 0.0,
2214 }]);
2215 let mut nodes = Vec::with_capacity(result_limit + usize::from(include_roots));
2216 if include_roots {
2217 nodes.push(PathNode {
2218 node_id: root_id,
2219 via_edge: None,
2220 depth: 0,
2221 name: None,
2222 kind: None,
2223 properties: None,
2224 weight: 0.0,
2225 });
2226 }
2227 let mut non_root_count = 0usize;
2228
2229 'root_walk: while non_root_count < result_limit {
2230 let Some(current) = frontier.pop_front() else {
2231 break;
2232 };
2233 if current.depth >= opts.max_depth {
2234 continue;
2235 }
2236 if budget.is_expired() {
2237 return Err(traversal_timeout_error(&budget));
2238 }
2239
2240 for sql in &statements {
2241 let row_cap = budget.remaining_work().saturating_add(1);
2242 let mut params: Vec<Box<dyn rusqlite::types::ToSql>> = vec![
2243 Box::new(namespace.clone()),
2244 Box::new(current.node_id.to_string()),
2245 Box::new(row_cap as i64),
2246 ];
2247 if let Some(relations) = &opts.relations {
2248 params.extend(relations.iter().map(|relation| {
2249 Box::new(relation.to_string()) as Box<dyn rusqlite::types::ToSql>
2250 }));
2251 }
2252 if let Some(min_weight) = opts.min_weight {
2253 params.push(Box::new(min_weight));
2254 }
2255 let param_refs = params
2256 .iter()
2257 .map(|param| param.as_ref())
2258 .collect::<Vec<&dyn rusqlite::types::ToSql>>();
2259
2260 let _snapshot = khive_storage::tx_registry::register_scoped(
2261 Some("graph_traverse_read".to_string()),
2262 origin.clone(),
2263 );
2264 let mut stmt = conn.prepare(sql).map_err(&map_sql_error)?;
2265 counted_queries.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2266 let mut rows = stmt.query(param_refs.as_slice()).map_err(&map_sql_error)?;
2267 while let Some(row) = rows.next().map_err(&map_sql_error)? {
2268 #[cfg(test)]
2271 tests::traverse_snapshot_seam::hook(current.node_id);
2272 counted_rows.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2273 if budget.is_expired() {
2274 return Err(traversal_timeout_error(&budget));
2275 }
2276 if !budget.try_consume_row() {
2277 return Err(traversal_work_error(&budget));
2278 }
2279 let node_str: String = row.get(0).map_err(&map_sql_error)?;
2280 let edge_str: String = row.get(1).map_err(&map_sql_error)?;
2281 let edge_weight: f64 = row.get(2).map_err(&map_sql_error)?;
2282 let node_id = parse_uuid(&node_str).map_err(&map_sql_error)?;
2283 if !seen.insert(node_id) {
2284 continue;
2285 }
2286 let via_edge = parse_uuid(&edge_str).map_err(&map_sql_error)?;
2287 let depth = current.depth + 1;
2288 let total_weight = current.total_weight + edge_weight;
2289 nodes.push(PathNode {
2290 node_id,
2291 via_edge: Some(via_edge),
2292 depth,
2293 name: None,
2294 kind: None,
2295 properties: None,
2296 weight: total_weight,
2297 });
2298 non_root_count += 1;
2299 if depth < opts.max_depth {
2300 frontier.push_back(TraversalFrontierNode {
2301 node_id,
2302 depth,
2303 total_weight,
2304 });
2305 }
2306 if non_root_count == result_limit {
2307 break 'root_walk;
2308 }
2309 }
2310 }
2311 }
2312
2313 if !nodes.is_empty() {
2314 let total_weight = nodes.iter().map(|node| node.weight).fold(0.0_f64, f64::max);
2315 all_paths.push(GraphPath {
2316 root_id,
2317 nodes,
2318 total_weight,
2319 });
2320 }
2321 }
2322 Ok(all_paths)
2323 })();
2324
2325 conn.progress_handler(0, None::<fn() -> bool>)
2326 .map_err(|e| map_err(e, "traverse_progress_handler_clear"))?;
2327 result
2328}
2329
2330impl SqlGraphStore {
2331 async fn query_neighbors_page(
2332 &self,
2333 operation: &'static str,
2334 node_id: Uuid,
2335 query: NeighborQuery,
2336 after: Option<NeighborCursor>,
2337 neighbor_kinds: Option<Vec<String>>,
2338 ) -> Result<Vec<NeighborHit>, StorageError> {
2339 count_neighbor_select();
2340
2341 let namespace = self.namespace.clone();
2342 let node_str = node_id.to_string();
2343 let counted_queries = Arc::new(std::sync::atomic::AtomicU64::new(0));
2344 let counted_rows = Arc::new(std::sync::atomic::AtomicU64::new(0));
2345 let closure_queries = Arc::clone(&counted_queries);
2346 let closure_rows = Arc::clone(&counted_rows);
2347 let result = self
2348 .with_reader(operation, move |conn| {
2349 let base_out = "SELECT target_id AS node_id, id AS edge_id, relation, weight \
2350 FROM graph_edges \
2351 WHERE namespace = ?1 AND source_id = ?2 AND deleted_at IS NULL";
2352 let base_in = "SELECT source_id AS node_id, id AS edge_id, relation, weight \
2353 FROM graph_edges \
2354 WHERE namespace = ?1 AND target_id = ?2 AND deleted_at IS NULL";
2355 let sql = match query.direction {
2356 Direction::Out => base_out.to_string(),
2357 Direction::In => base_in.to_string(),
2358 Direction::Both => format!("{} UNION ALL {}", base_out, base_in),
2359 };
2360 let (where_extra, limit_clause, extra_params) =
2361 neighbor_extra_clause(&query, 3, after.as_ref(), neighbor_kinds.as_deref());
2362 let full_sql = format!(
2363 "SELECT node_id, edge_id, relation, weight FROM ({}){} \
2364 ORDER BY weight DESC, node_id ASC, edge_id ASC{}",
2365 sql, where_extra, limit_clause
2366 );
2367 let mut stmt = conn.prepare(&full_sql)?;
2368 let mut all_params: Vec<Box<dyn rusqlite::types::ToSql>> = Vec::new();
2369 all_params.push(Box::new(namespace.clone()));
2370 all_params.push(Box::new(node_str.clone()));
2371 all_params.extend(extra_params);
2372 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
2373 all_params.iter().map(|p| p.as_ref()).collect();
2374
2375 closure_queries.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2376 let rows = stmt.query_map(param_refs.as_slice(), |row| {
2377 let nid_str: String = row.get(0)?;
2378 let eid_str: String = row.get(1)?;
2379 let relation_str: String = row.get(2)?;
2380 let weight: f64 = row.get(3)?;
2381 Ok((nid_str, eid_str, relation_str, weight))
2382 })?;
2383 let mut hits = Vec::new();
2384 for row in rows {
2385 let (nid_str, eid_str, relation_str, weight) = row?;
2386 closure_rows.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2387 let relation = relation_str.parse::<EdgeRelation>().map_err(|e| {
2388 rusqlite::Error::FromSqlConversionFailure(
2389 2,
2390 rusqlite::types::Type::Text,
2391 Box::new(e),
2392 )
2393 })?;
2394 hits.push(NeighborHit {
2395 node_id: parse_uuid(&nid_str)?,
2396 edge_id: parse_uuid(&eid_str)?,
2397 relation,
2398 weight,
2399 name: None,
2400 kind: None,
2401 entity_type: None,
2402 });
2403 }
2404 Ok(hits)
2405 })
2406 .await;
2407
2408 report_graph_usage(&counted_queries, &counted_rows);
2409 result
2410 }
2411}
2412
2413#[async_trait]
2414impl GraphStore for SqlGraphStore {
2415 async fn update_symmetric_edge_if_unchanged(
2416 &self,
2417 request: SymmetricEdgeUpdateRequest,
2418 ) -> StorageResult<SymmetricEdgeUpdateOutcome> {
2419 self.update_symmetric_edge(request).await
2420 }
2421
2422 async fn latest_annotating_note(
2423 &self,
2424 node_id: Uuid,
2425 kind: &str,
2426 tag: &str,
2427 ) -> Result<Option<(Uuid, i64)>, StorageError> {
2428 let namespace = self.namespace.clone();
2429 let node_id = node_id.to_string();
2430 let kind = kind.to_owned();
2431 let tag = tag.to_owned();
2432 self.with_indexed_reader("latest_annotating_note", move |conn| {
2433 conn.query_row(
2434 LATEST_ANNOTATING_NOTE_SQL,
2435 rusqlite::params![namespace, node_id, kind, tag],
2436 |row| {
2437 let id: String = row.get(0)?;
2438 Ok((parse_uuid(&id)?, row.get(1)?))
2439 },
2440 )
2441 .optional()
2442 })
2443 .await
2444 }
2445
2446 async fn latest_annotating_note_with_property(
2447 &self,
2448 node_id: Uuid,
2449 kind: &str,
2450 tag: &str,
2451 property_key: &str,
2452 property_value: &str,
2453 ) -> Result<Option<(Uuid, i64)>, StorageError> {
2454 let namespace = self.namespace.clone();
2455 let node_id = node_id.to_string();
2456 let kind = kind.to_owned();
2457 let tag = tag.to_owned();
2458 let property_key = property_key.to_owned();
2459 let property_value = property_value.to_owned();
2460 self.with_indexed_reader("latest_annotating_note_with_property", move |conn| {
2461 conn.query_row(
2462 LATEST_ANNOTATING_NOTE_WITH_PROPERTY_SQL,
2463 rusqlite::params![namespace, node_id, kind, tag, property_key, property_value],
2464 |row| {
2465 let id: String = row.get(0)?;
2466 Ok((parse_uuid(&id)?, row.get(1)?))
2467 },
2468 )
2469 .optional()
2470 })
2471 .await
2472 }
2473
2474 async fn upsert_edge(&self, edge: Edge) -> Result<(), StorageError> {
2475 self.upsert_edge_observed(EdgeUpsertRequest {
2476 edge,
2477 resurrect: false,
2478 })
2479 .await
2480 .map(|_| ())
2481 }
2482
2483 async fn upsert_edge_observed(
2484 &self,
2485 request: EdgeUpsertRequest,
2486 ) -> Result<EdgeUpsertResult, StorageError> {
2487 match self
2488 .observed_edge_write("upsert_edge_observed", request, false)
2489 .await?
2490 {
2491 GuardedEdgeUpsertOutcome::Written(result) => Ok(result),
2492 GuardedEdgeUpsertOutcome::Refused(EdgeUpsertRefusal::ResurrectionRequired { edge }) => {
2493 Err(resurrection_required_error("upsert_edge_observed", &edge))
2494 }
2495 GuardedEdgeUpsertOutcome::Refused(EdgeUpsertRefusal::MissingEndpoints(_)) => {
2496 Err(StorageError::Conflict {
2497 capability: StorageCapability::Graph,
2498 operation: "upsert_edge_observed".into(),
2499 message: "unguarded edge upsert reported a missing-endpoint refusal".into(),
2500 })
2501 }
2502 }
2503 }
2504
2505 async fn insert_edge_if_absent(&self, edge: Edge) -> Result<bool, StorageError> {
2506 let statement = edge_insert_if_absent_statement(&edge);
2507 self.with_writer("insert_edge_if_absent", move |conn| {
2508 let mut stmt = conn.prepare(&statement.sql)?;
2509 bind_params(&mut stmt, &statement.params)?;
2510 Ok(stmt.raw_execute()? > 0)
2511 })
2512 .await
2513 }
2514
2515 async fn insert_commit_annotation_if_absent(
2516 &self,
2517 edge: Edge,
2518 guard: CommitAnnotationGuard,
2519 ) -> Result<CommitAnnotationInsertOutcome, StorageError> {
2520 const OP: &str = "insert_commit_annotation_if_absent";
2521 if edge.relation != EdgeRelation::Annotates || edge.deleted_at.is_some() {
2522 return Err(StorageError::InvalidInput {
2523 capability: StorageCapability::Graph,
2524 operation: OP.into(),
2525 message: "expected a live annotates edge".into(),
2526 });
2527 }
2528 if let Some(writer_task) = self.current_writer_task(OP)? {
2529 return writer_task
2530 .send_bounded_at(
2531 move |conn| {
2532 conditional_commit_annotation_insert(conn, edge, &guard)
2533 .map_err(|error| map_err(error, OP))
2534 },
2535 OP,
2536 )
2537 .await;
2538 }
2539 let origin = self.pool.origin();
2540 self.with_writer(OP, move |conn| {
2541 conn.execute_batch("BEGIN IMMEDIATE")?;
2542 let _tx_handle =
2543 khive_storage::tx_registry::register_scoped(Some(OP.to_string()), origin);
2544 let outcome = match conditional_commit_annotation_insert(conn, edge, &guard) {
2545 Ok(outcome) => outcome,
2546 Err(error) => {
2547 let _ = conn.execute_batch("ROLLBACK");
2548 return Err(error);
2549 }
2550 };
2551 if let Err(error) = conn.execute_batch("COMMIT") {
2552 let _ = conn.execute_batch("ROLLBACK");
2553 return Err(error);
2554 }
2555 Ok(outcome)
2556 })
2557 .await
2558 }
2559
2560 async fn replace_edge_if_unchanged(
2561 &self,
2562 edge: Edge,
2563 expected_updated_at: DateTime<Utc>,
2564 expected_deleted_at: Option<DateTime<Utc>>,
2565 ) -> Result<bool, StorageError> {
2566 let statement =
2567 edge_replace_if_unchanged_statement(&edge, expected_updated_at, expected_deleted_at);
2568 self.with_writer("replace_edge_if_unchanged", move |conn| {
2569 let mut stmt = conn.prepare(&statement.sql)?;
2570 bind_params(&mut stmt, &statement.params)?;
2571 Ok(stmt.raw_execute()? > 0)
2572 })
2573 .await
2574 }
2575
2576 async fn upsert_edges(&self, edges: Vec<Edge>) -> Result<BatchWriteSummary, StorageError> {
2577 let attempted = edges.len() as u64;
2578 let requests = edges
2579 .into_iter()
2580 .map(|edge| EdgeUpsertRequest {
2581 edge,
2582 resurrect: false,
2583 })
2584 .collect();
2585 let outcome = self
2586 .observed_edge_batch_write("upsert_edges", requests, false)
2587 .await?;
2588 if let Some(refusal) = outcome.refusal {
2589 return match refusal.reason {
2590 EdgeUpsertRefusal::ResurrectionRequired { edge } => {
2591 Err(resurrection_required_error("upsert_edges", &edge))
2592 }
2593 EdgeUpsertRefusal::MissingEndpoints(_) => Err(StorageError::Conflict {
2594 capability: StorageCapability::Graph,
2595 operation: "upsert_edges".into(),
2596 message: "unguarded edge batch reported a missing-endpoint refusal".into(),
2597 }),
2598 };
2599 }
2600 Ok(BatchWriteSummary {
2601 attempted,
2602 affected: outcome.rows.len() as u64,
2603 ..BatchWriteSummary::default()
2604 })
2605 }
2606
2607 async fn upsert_edge_guarded(&self, edge: Edge) -> Result<GuardedWriteOutcome, StorageError> {
2608 match self
2609 .upsert_edge_guarded_observed(EdgeUpsertRequest {
2610 edge,
2611 resurrect: false,
2612 })
2613 .await?
2614 {
2615 GuardedEdgeUpsertOutcome::Written(_) => Ok(GuardedWriteOutcome::Written),
2616 GuardedEdgeUpsertOutcome::Refused(EdgeUpsertRefusal::MissingEndpoints(missing)) => {
2617 Ok(GuardedWriteOutcome::Refused(missing))
2618 }
2619 GuardedEdgeUpsertOutcome::Refused(EdgeUpsertRefusal::ResurrectionRequired { edge }) => {
2620 Err(resurrection_required_error("upsert_edge_guarded", &edge))
2621 }
2622 }
2623 }
2624
2625 async fn upsert_edge_guarded_observed(
2626 &self,
2627 request: EdgeUpsertRequest,
2628 ) -> Result<GuardedEdgeUpsertOutcome, StorageError> {
2629 self.observed_edge_write("upsert_edge_guarded_observed", request, true)
2630 .await
2631 }
2632
2633 async fn upsert_edges_guarded(
2634 &self,
2635 edges: Vec<Edge>,
2636 ) -> Result<GuardedBatchOutcome, StorageError> {
2637 let attempted = edges.len() as u64;
2638 let requests = edges
2639 .iter()
2640 .cloned()
2641 .map(|edge| EdgeUpsertRequest {
2642 edge,
2643 resurrect: false,
2644 })
2645 .collect();
2646 let outcome = self.upsert_edges_guarded_observed(requests).await?;
2647 match outcome.refusal {
2648 None => Ok(GuardedBatchOutcome {
2649 summary: BatchWriteSummary {
2650 attempted,
2651 affected: outcome.rows.len() as u64,
2652 ..BatchWriteSummary::default()
2653 },
2654 refused: None,
2655 }),
2656 Some(refusal) => match refusal.reason {
2657 EdgeUpsertRefusal::MissingEndpoints(missing) => {
2658 let index = refusal.entry_index;
2659 let edge = &edges[index];
2660 let (source_id, target_id) =
2661 canonical_edge_endpoints(edge.relation, edge.source_id, edge.target_id);
2662 let message = format!(
2663 "batch entry {index}: edge endpoint no longer exists at write time: source \
2664 {source_id} or target {target_id}"
2665 );
2666 let mut summary = BatchWriteSummary {
2667 attempted,
2668 ..BatchWriteSummary::default()
2669 };
2670 summary.first_error = message.clone();
2673 let refusal = GuardedBatchRefusal {
2674 entry_index: index,
2675 missing,
2676 };
2677 for (failed_index, failed_edge) in edges.iter().enumerate() {
2678 refusal.record_failure(&mut summary, failed_index, failed_edge, &message);
2679 }
2680 Ok(GuardedBatchOutcome {
2681 summary,
2682 refused: Some(refusal),
2683 })
2684 }
2685 EdgeUpsertRefusal::ResurrectionRequired { edge } => {
2686 Err(resurrection_required_error("upsert_edges_guarded", &edge))
2687 }
2688 },
2689 }
2690 }
2691
2692 async fn upsert_edges_guarded_observed(
2693 &self,
2694 requests: Vec<EdgeUpsertRequest>,
2695 ) -> Result<GuardedEdgeBatchUpsertOutcome, StorageError> {
2696 self.observed_edge_batch_write("upsert_edges_guarded_observed", requests, true)
2697 .await
2698 }
2699
2700 async fn get_edge(&self, id: LinkId) -> Result<Option<Edge>, StorageError> {
2701 let id_str = Uuid::from(id).to_string();
2702
2703 self.with_reader("get_edge", move |conn| {
2704 let mut stmt = conn.prepare(
2705 "SELECT namespace, id, source_id, target_id, relation, weight, \
2706 created_at, updated_at, deleted_at, metadata, target_backend \
2707 FROM graph_edges WHERE id = ?1 AND deleted_at IS NULL",
2708 )?;
2709 let mut rows = stmt.query(rusqlite::params![id_str])?;
2710 match rows.next()? {
2711 Some(row) => Ok(Some(read_edge(row)?)),
2712 None => Ok(None),
2713 }
2714 })
2715 .await
2716 }
2717
2718 async fn get_edge_including_deleted(&self, id: LinkId) -> Result<Option<Edge>, StorageError> {
2719 let id_str = Uuid::from(id).to_string();
2720
2721 self.with_reader("get_edge_including_deleted", move |conn| {
2722 let mut stmt = conn.prepare(
2723 "SELECT namespace, id, source_id, target_id, relation, weight, \
2724 created_at, updated_at, deleted_at, metadata, target_backend \
2725 FROM graph_edges WHERE id = ?1",
2726 )?;
2727 let mut rows = stmt.query(rusqlite::params![id_str])?;
2728 match rows.next()? {
2729 Some(row) => Ok(Some(read_edge(row)?)),
2730 None => Ok(None),
2731 }
2732 })
2733 .await
2734 }
2735
2736 async fn edge_sequence(&self, id: Uuid) -> Result<Option<i64>, StorageError> {
2737 let id = id.to_string();
2738 self.with_reader("edge_sequence", move |conn| {
2739 conn.query_row(
2740 "SELECT seq FROM graph_edges_seq WHERE edge_id = ?1",
2741 rusqlite::params![id],
2742 |row| row.get(0),
2743 )
2744 .optional()
2745 })
2746 .await
2747 }
2748
2749 async fn edge_sequences(&self, ids: &[Uuid]) -> Result<Vec<(Uuid, i64)>, StorageError> {
2750 if ids.is_empty() {
2751 return Ok(Vec::new());
2752 }
2753 let ids = ids.to_vec();
2754 self.with_reader("edge_sequences", move |conn| {
2755 const CHUNK: usize = 900;
2756 let mut resolved = Vec::with_capacity(ids.len());
2757 for chunk in ids.chunks(CHUNK) {
2758 let placeholders = (1..=chunk.len())
2759 .map(|index| format!("?{index}"))
2760 .collect::<Vec<_>>()
2761 .join(", ");
2762 let sql = format!(
2763 "SELECT edge_id, seq FROM graph_edges_seq WHERE edge_id IN ({placeholders})"
2764 );
2765 let strings = chunk.iter().map(Uuid::to_string).collect::<Vec<_>>();
2766 let params = strings
2767 .iter()
2768 .map(|id| id as &dyn rusqlite::types::ToSql)
2769 .collect::<Vec<_>>();
2770 let mut stmt = conn.prepare(&sql)?;
2771 let rows = stmt.query_map(params.as_slice(), |row| {
2772 let id: String = row.get(0)?;
2773 Ok((parse_uuid(&id)?, row.get::<_, i64>(1)?))
2774 })?;
2775 resolved.extend(rows.collect::<Result<Vec<_>, _>>()?);
2776 }
2777 Ok(resolved)
2778 })
2779 .await
2780 }
2781
2782 async fn get_edge_by_natural_key_including_deleted(
2783 &self,
2784 namespace: &str,
2785 source_id: Uuid,
2786 target_id: Uuid,
2787 relation: EdgeRelation,
2788 ) -> Result<Option<Edge>, StorageError> {
2789 let namespace = namespace.to_string();
2790 let source_str = source_id.to_string();
2791 let target_str = target_id.to_string();
2792 let relation_str = relation.to_string();
2793
2794 self.with_reader("get_edge_by_natural_key_including_deleted", move |conn| {
2795 let mut stmt = conn.prepare(
2796 "SELECT namespace, id, source_id, target_id, relation, weight, \
2797 created_at, updated_at, deleted_at, metadata, target_backend \
2798 FROM graph_edges \
2799 WHERE namespace = ?1 AND source_id = ?2 AND target_id = ?3 AND relation = ?4",
2800 )?;
2801 let mut rows = stmt.query(rusqlite::params![
2802 namespace,
2803 source_str,
2804 target_str,
2805 relation_str
2806 ])?;
2807 match rows.next()? {
2808 Some(row) => Ok(Some(read_edge(row)?)),
2809 None => Ok(None),
2810 }
2811 })
2812 .await
2813 }
2814
2815 async fn get_edges(&self, ids: &[LinkId]) -> Result<Vec<Edge>, StorageError> {
2816 if ids.is_empty() {
2817 return Ok(Vec::new());
2818 }
2819 const CHUNK: usize = 900;
2821 let id_strs: Vec<String> = ids.iter().map(|id| Uuid::from(*id).to_string()).collect();
2822
2823 let mut result: Vec<Edge> = Vec::with_capacity(ids.len());
2824 for chunk in id_strs.chunks(CHUNK) {
2825 let chunk_owned: Vec<String> = chunk.to_vec();
2826 let edges = self
2827 .with_reader("get_edges", move |conn| {
2828 let placeholders: Vec<String> =
2829 (1..=chunk_owned.len()).map(|i| format!("?{}", i)).collect();
2830 let sql = format!(
2831 "SELECT namespace, id, source_id, target_id, relation, weight, \
2832 created_at, updated_at, deleted_at, metadata, target_backend \
2833 FROM graph_edges WHERE id IN ({}) AND deleted_at IS NULL",
2834 placeholders.join(",")
2835 );
2836 let mut stmt = conn.prepare(&sql)?;
2837 let params: Vec<&dyn rusqlite::types::ToSql> = chunk_owned
2838 .iter()
2839 .map(|s| s as &dyn rusqlite::types::ToSql)
2840 .collect();
2841 let rows = stmt.query_map(params.as_slice(), read_edge)?;
2842 let mut edges = Vec::new();
2843 for row in rows {
2844 edges.push(row?);
2845 }
2846 Ok(edges)
2847 })
2848 .await?;
2849 result.extend(edges);
2850 }
2851 Ok(result)
2852 }
2853
2854 async fn get_edge_read_outcomes(
2855 &self,
2856 ids: &[LinkId],
2857 ) -> Result<Vec<Result<Option<Edge>, StorageError>>, StorageError> {
2858 const CHUNK: usize = 900;
2859 let mut outcomes = Vec::with_capacity(ids.len());
2860 for chunk in ids.chunks(CHUNK) {
2861 let requested: Vec<String> =
2862 chunk.iter().map(|id| Uuid::from(*id).to_string()).collect();
2863 let expected = requested.len();
2864 let rows = self
2865 .with_reader("get_edge_read_outcomes", move |conn| {
2866 let requested_json = serde_json::to_string(&requested)
2867 .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))?;
2868 let mut stmt = conn.prepare(
2869 "SELECT e.namespace, e.id, e.source_id, e.target_id, e.relation, e.weight, \
2870 e.created_at, e.updated_at, e.deleted_at, e.metadata, e.target_backend, \
2871 requested.key \
2872 FROM json_each(?1) AS requested \
2873 LEFT JOIN graph_edges AS e \
2874 ON e.id = requested.value AND e.deleted_at IS NULL \
2875 ORDER BY CAST(requested.key AS INTEGER) ASC",
2876 )?;
2877 let mut rows = stmt.query([requested_json])?;
2878 let mut outcomes = Vec::with_capacity(expected);
2879 while let Some(row) = rows.next()? {
2880 let ordinal: i64 = row.get(11)?;
2881 let ordinal =
2882 usize::try_from(ordinal).map_err(|_| rusqlite::Error::InvalidQuery)?;
2883 if ordinal != outcomes.len() {
2884 return Err(rusqlite::Error::InvalidQuery);
2885 }
2886 let outcome = match row.get_ref(1)? {
2887 rusqlite::types::ValueRef::Null => Ok(None),
2888 _ => read_edge(row).map(Some).map_err(|e| map_err(e, "get_edge")),
2889 };
2890 outcomes.push(outcome);
2891 }
2892 if outcomes.len() != expected {
2893 return Err(rusqlite::Error::InvalidQuery);
2894 }
2895 Ok(outcomes)
2896 })
2897 .await?;
2898 outcomes.extend(rows);
2899 }
2900 Ok(outcomes)
2901 }
2902
2903 async fn batch_neighbors(
2904 &self,
2905 sources: &[Uuid],
2906 query: NeighborQuery,
2907 ) -> Result<Vec<(Uuid, NeighborHit)>, StorageError> {
2908 use khive_storage::types::Direction;
2909
2910 if sources.is_empty() {
2911 return Ok(Vec::new());
2912 }
2913 let mut seen_sources = HashSet::with_capacity(sources.len());
2914 let unique_sources: Vec<Uuid> = sources
2915 .iter()
2916 .copied()
2917 .filter(|source| seen_sources.insert(*source))
2918 .collect();
2919 const CHUNK_SIZE: usize = 880;
2920
2921 let namespace = self.namespace.clone();
2922 let mut result: Vec<(Uuid, NeighborHit)> = Vec::new();
2923 let counted_queries = Arc::new(std::sync::atomic::AtomicU64::new(0));
2926 let counted_rows = Arc::new(std::sync::atomic::AtomicU64::new(0));
2927
2928 for chunk in unique_sources.chunks(CHUNK_SIZE) {
2929 let chunk_owned: Vec<Uuid> = chunk.to_vec();
2930 let query_clone = query.clone();
2931 let ns = namespace.clone();
2932 let closure_queries = Arc::clone(&counted_queries);
2933 let closure_rows = Arc::clone(&counted_rows);
2934
2935 let chunk_result = self
2936 .with_reader("batch_neighbors", move |conn| {
2937 let src_strs: Vec<String> = chunk_owned.iter().map(|u| u.to_string()).collect();
2938
2939 let sources_json = serde_json::to_string(&src_strs).map_err(|error| {
2940 rusqlite::Error::ToSqlConversionFailure(Box::new(error))
2941 })?;
2942
2943 let build_inner_sql =
2944 |direction_out: bool,
2945 q: &NeighborQuery|
2946 -> (String, Vec<String>, Option<f64>) {
2947 let (filter_col, node_col) = if direction_out {
2948 ("source_id", "target_id")
2949 } else {
2950 ("target_id", "source_id")
2951 };
2952
2953 let mut rel_params: Vec<String> = Vec::new();
2954 let mut conditions: Vec<String> = Vec::new();
2955 let mut param_idx = 3;
2956
2957 if let Some(ref rels) = q.relations {
2958 if !rels.is_empty() {
2959 let ps: Vec<String> = rels
2960 .iter()
2961 .map(|r| {
2962 rel_params.push(r.to_string());
2963 let p = format!("?{param_idx}");
2964 param_idx += 1;
2965 p
2966 })
2967 .collect();
2968 conditions
2969 .push(format!("edges.relation IN ({})", ps.join(",")));
2970 }
2971 }
2972
2973 let min_weight_val = if let Some(min_w) = q.min_weight {
2976 conditions.push(format!("edges.weight >= ?{param_idx}"));
2977 Some(min_w)
2978 } else {
2979 None
2980 };
2981
2982 let where_extra = if conditions.is_empty() {
2983 String::new()
2984 } else {
2985 format!(" AND {}", conditions.join(" AND "))
2986 };
2987
2988 let sql = format!(
2989 "SELECT requested.origin_id, edges.{node_col} AS node_id, \
2990 edges.id AS edge_id, edges.relation, edges.weight \
2991 FROM requested CROSS JOIN graph_edges AS edges \
2992 ON edges.{filter_col} = requested.origin_id \
2993 WHERE edges.namespace = ?1 \
2994 AND edges.deleted_at IS NULL{where_extra}",
2995 );
2996 (sql, rel_params, min_weight_val)
2997 };
2998
2999 let mut all_params: Vec<Box<dyn rusqlite::types::ToSql>> = Vec::new();
3000 all_params.push(Box::new(ns.to_string()));
3001 all_params.push(Box::new(sources_json));
3002
3003 let (combined_inner, rel_params, min_weight_val) = match query_clone.direction {
3004 Direction::Out => build_inner_sql(true, &query_clone),
3005 Direction::In => build_inner_sql(false, &query_clone),
3006 Direction::Both => {
3007 let (out_sql, rel_params, min_weight_val) =
3008 build_inner_sql(true, &query_clone);
3009 let (in_sql, _, _) = build_inner_sql(false, &query_clone);
3010 (
3011 format!("{out_sql} UNION ALL {in_sql}"),
3012 rel_params,
3013 min_weight_val,
3014 )
3015 }
3016 };
3017
3018 for relation in rel_params {
3019 all_params.push(Box::new(relation));
3020 }
3021 if let Some(min_weight) = min_weight_val {
3022 all_params.push(Box::new(min_weight));
3023 }
3024 let limit_param_idx = all_params.len() + 1;
3025
3026 let full_sql = if let Some(lim) = query_clone.limit {
3034 all_params.push(Box::new(lim as i64));
3035 format!(
3036 "WITH requested(origin_id) AS (\
3037 SELECT value FROM json_each(?2)\
3038 ) SELECT origin_id, node_id, edge_id, relation, weight \
3039 FROM (SELECT *, ROW_NUMBER() OVER (PARTITION BY origin_id \
3040 ORDER BY weight DESC, node_id ASC) AS rn \
3041 FROM ({combined_inner})) WHERE rn <= ?{limit_param_idx}",
3042 )
3043 } else {
3044 format!(
3045 "WITH requested(origin_id) AS (\
3046 SELECT value FROM json_each(?2)\
3047 ) SELECT origin_id, node_id, edge_id, relation, weight \
3048 FROM ({combined_inner})",
3049 )
3050 };
3051
3052 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
3053 all_params.iter().map(|p| p.as_ref()).collect();
3054
3055 let mut stmt = conn.prepare(&full_sql)?;
3056 closure_queries.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3057 let rows = stmt.query_map(param_refs.as_slice(), |row| {
3058 let origin_str: String = row.get(0)?;
3059 let nid_str: String = row.get(1)?;
3060 let eid_str: String = row.get(2)?;
3061 let relation_str: String = row.get(3)?;
3062 let weight: f64 = row.get(4)?;
3063 Ok((origin_str, nid_str, eid_str, relation_str, weight))
3064 })?;
3065
3066 let mut pairs = Vec::new();
3067 for row in rows {
3068 let (origin_str, nid_str, eid_str, relation_str, weight) = row?;
3069 closure_rows.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3070 let origin = parse_uuid(&origin_str)?;
3071 let node_id = parse_uuid(&nid_str)?;
3072 let edge_id = parse_uuid(&eid_str)?;
3073 let relation = relation_str.parse::<EdgeRelation>().map_err(|e| {
3074 rusqlite::Error::FromSqlConversionFailure(
3075 3,
3076 rusqlite::types::Type::Text,
3077 Box::new(e),
3078 )
3079 })?;
3080 pairs.push((
3081 origin,
3082 NeighborHit {
3083 node_id,
3084 edge_id,
3085 relation,
3086 weight,
3087 name: None,
3088 kind: None,
3089 entity_type: None,
3090 },
3091 ));
3092 }
3093 Ok(pairs)
3094 })
3095 .await;
3096 let pairs = match chunk_result {
3097 Ok(pairs) => pairs,
3098 Err(e) => {
3099 report_graph_usage(&counted_queries, &counted_rows);
3100 return Err(e);
3101 }
3102 };
3103 result.extend(pairs);
3104 }
3105 report_graph_usage(&counted_queries, &counted_rows);
3106
3107 let requested: HashSet<Uuid> = unique_sources.iter().copied().collect();
3108 let mut grouped: HashMap<Uuid, Vec<NeighborHit>> =
3109 HashMap::with_capacity(unique_sources.len());
3110 for (origin, hit) in result {
3111 if !requested.contains(&origin) {
3112 return Err(StorageError::Internal(format!(
3113 "batch_neighbors returned unrequested origin {origin}"
3114 )));
3115 }
3116 grouped.entry(origin).or_default().push(hit);
3117 }
3118
3119 for hits in grouped.values_mut() {
3120 hits.sort_by(|a, b| {
3121 b.weight
3122 .partial_cmp(&a.weight)
3123 .unwrap_or(std::cmp::Ordering::Equal)
3124 .then(a.node_id.cmp(&b.node_id))
3125 .then(a.edge_id.cmp(&b.edge_id))
3126 });
3127 }
3128
3129 let mut ordered = Vec::new();
3130 for &source in sources {
3131 if let Some(hits) = grouped.get(&source) {
3132 ordered.extend(hits.iter().cloned().map(|hit| (source, hit)));
3133 }
3134 }
3135 Ok(ordered)
3136 }
3137
3138 async fn delete_edge(&self, id: LinkId, mode: DeleteMode) -> Result<bool, StorageError> {
3139 let id = Uuid::from(id);
3140 let statement = match mode {
3141 DeleteMode::Soft => {
3142 edge_soft_delete_statement(id, chrono::Utc::now().timestamp_micros())
3143 }
3144 DeleteMode::Hard => edge_hard_delete_statement(id),
3145 };
3146 self.with_writer("delete_edge", move |conn| {
3147 let mut stmt = conn.prepare(&statement.sql)?;
3148 bind_params(&mut stmt, &statement.params)?;
3149 Ok(stmt.raw_execute()? > 0)
3150 })
3151 .await
3152 }
3153
3154 async fn query_edges(
3155 &self,
3156 filter: EdgeFilter,
3157 sort: Vec<SortOrder<EdgeSortField>>,
3158 page: PageRequest,
3159 ) -> Result<Page<Edge>, StorageError> {
3160 let namespace = self.namespace.clone();
3161 let limit_i64 = i64::from(page.limit);
3162 let offset_i64 = i64::try_from(page.offset).map_err(|_| StorageError::InvalidInput {
3163 capability: StorageCapability::Graph,
3164 operation: "query_edges".into(),
3165 message: format!(
3166 "PageRequest: offset must be <= i64::MAX, got {}",
3167 page.offset
3168 ),
3169 })?;
3170 self.with_reader("query_edges", move |conn| {
3171 let (where_clause, mut all_params) = build_edge_filter_sql(&namespace, &filter);
3172 let order_clause = edge_order_clause(&sort);
3173 all_params.push(Box::new(limit_i64));
3174 all_params.push(Box::new(offset_i64));
3175
3176 let limit_idx = all_params.len() - 1;
3177 let offset_idx = all_params.len();
3178
3179 let data_sql = format!(
3180 "SELECT namespace, id, source_id, target_id, relation, weight, \
3181 created_at, updated_at, deleted_at, metadata, target_backend \
3182 FROM graph_edges{}{} LIMIT ?{} OFFSET ?{}",
3183 where_clause, order_clause, limit_idx, offset_idx,
3184 );
3185
3186 let mut stmt = conn.prepare(&data_sql)?;
3187 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
3188 all_params.iter().map(|p| p.as_ref()).collect();
3189 let rows = stmt.query_map(param_refs.as_slice(), read_edge)?;
3190
3191 let mut items = Vec::new();
3192 for row in rows {
3193 items.push(row?);
3194 }
3195
3196 Ok(Page { items, total: None })
3197 })
3198 .await
3199 }
3200
3201 async fn count_edges(&self, filter: EdgeFilter) -> Result<u64, StorageError> {
3202 let namespace = self.namespace.clone();
3203 self.with_reader("count_edges", move |conn| {
3204 let (where_clause, params) = build_edge_filter_sql(&namespace, &filter);
3205 let sql = format!(
3206 "SELECT COUNT(*) FROM graph_edges{}",
3207 with_live_endpoints(&where_clause)
3208 );
3209 let mut stmt = conn.prepare(&sql)?;
3210 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
3211 params.iter().map(|p| p.as_ref()).collect();
3212 let count: i64 = stmt.query_row(param_refs.as_slice(), |row| row.get(0))?;
3213 Ok(count as u64)
3214 })
3215 .await
3216 }
3217
3218 async fn count_edges_in_namespaces(
3219 &self,
3220 namespaces: &[String],
3221 filter: EdgeFilter,
3222 ) -> Result<u64, StorageError> {
3223 let namespaces: Vec<String> = namespaces
3224 .iter()
3225 .cloned()
3226 .collect::<HashSet<_>>()
3227 .into_iter()
3228 .collect();
3229 self.with_reader("count_edges_in_namespaces", move |conn| {
3230 let mut total = 0;
3231 for chunk in namespaces.chunks(NAMESPACE_COUNT_CHUNK_SIZE) {
3232 let (where_clause, params) = build_edge_filter_sql_for_namespaces(chunk, &filter);
3233 let sql = format!(
3234 "SELECT COUNT(*) FROM graph_edges{}",
3235 with_live_endpoints(&where_clause)
3236 );
3237 let mut stmt = conn.prepare(&sql)?;
3238 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
3239 params.iter().map(|p| p.as_ref()).collect();
3240 let count: i64 = stmt.query_row(param_refs.as_slice(), |row| row.get(0))?;
3241 total += count as u64;
3242 }
3243 Ok(total)
3244 })
3245 .await
3246 }
3247
3248 async fn query_edges_in_namespaces(
3249 &self,
3250 namespaces: &[String],
3251 filter: EdgeFilter,
3252 sort: Vec<SortOrder<EdgeSortField>>,
3253 page: PageRequest,
3254 ) -> Result<Page<Edge>, StorageError> {
3255 let namespaces: Vec<String> = {
3270 let mut seen = HashSet::new();
3271 namespaces
3272 .iter()
3273 .filter(|ns| seen.insert((*ns).clone()))
3274 .cloned()
3275 .collect()
3276 };
3277 let limit_i64 = i64::from(page.limit);
3278 let offset_i64 = i64::try_from(page.offset).map_err(|_| StorageError::InvalidInput {
3279 capability: StorageCapability::Graph,
3280 operation: "query_edges_in_namespaces".into(),
3281 message: format!(
3282 "PageRequest: offset must be <= i64::MAX, got {}",
3283 page.offset
3284 ),
3285 })?;
3286 self.with_reader("query_edges_in_namespaces", move |conn| {
3287 let namespaces_json = serde_json::to_string(&namespaces)
3288 .map_err(|error| rusqlite::Error::ToSqlConversionFailure(Box::new(error)))?;
3289
3290 let (where_clause, mut all_params) =
3291 build_edge_filter_sql_for_namespaces_json(&namespaces_json, &filter);
3292 let order_clause = edge_order_clause(&sort);
3293 all_params.push(Box::new(limit_i64));
3294 all_params.push(Box::new(offset_i64));
3295
3296 let limit_idx = all_params.len() - 1;
3297 let offset_idx = all_params.len();
3298
3299 let data_sql = format!(
3300 "SELECT namespace, id, source_id, target_id, relation, weight, \
3301 created_at, updated_at, deleted_at, metadata, target_backend \
3302 FROM graph_edges{}{} LIMIT ?{} OFFSET ?{}",
3303 where_clause, order_clause, limit_idx, offset_idx,
3304 );
3305
3306 let mut stmt = conn.prepare(&data_sql)?;
3307 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
3308 all_params.iter().map(|p| p.as_ref()).collect();
3309 let rows = stmt.query_map(param_refs.as_slice(), read_edge)?;
3310
3311 let mut items = Vec::new();
3312 for row in rows {
3313 items.push(row?);
3314 }
3315
3316 Ok(Page { items, total: None })
3317 })
3318 .await
3319 }
3320
3321 async fn count_edges_by_relation(&self) -> Result<Vec<(EdgeRelation, u64)>, StorageError> {
3322 let namespace = self.namespace.clone();
3323 self.with_reader("count_edges_by_relation", move |conn| {
3324 let sql = format!(
3325 "SELECT relation, COUNT(*) FROM graph_edges \
3326 WHERE namespace = ?1 AND deleted_at IS NULL AND {LIVE_ENDPOINTS_CONDITION} \
3327 GROUP BY relation"
3328 );
3329 let mut stmt = conn.prepare(&sql)?;
3330 let rows = stmt.query_map([&namespace], |row| {
3331 let relation_str: String = row.get(0)?;
3332 let count: i64 = row.get(1)?;
3333 Ok((relation_str, count))
3334 })?;
3335 let mut out = Vec::new();
3336 for row in rows {
3337 let (relation_str, count) = row?;
3338 let relation = relation_str.parse::<EdgeRelation>().map_err(|e| {
3339 rusqlite::Error::FromSqlConversionFailure(
3340 0,
3341 rusqlite::types::Type::Text,
3342 Box::new(e),
3343 )
3344 })?;
3345 out.push((relation, count as u64));
3346 }
3347 Ok(out)
3348 })
3349 .await
3350 }
3351
3352 async fn count_edges_by_relation_in_namespaces(
3353 &self,
3354 namespaces: &[String],
3355 ) -> Result<Vec<(EdgeRelation, u64)>, StorageError> {
3356 let namespaces: Vec<String> = namespaces
3357 .iter()
3358 .cloned()
3359 .collect::<HashSet<_>>()
3360 .into_iter()
3361 .collect();
3362 self.with_reader("count_edges_by_relation_in_namespaces", move |conn| {
3363 let mut totals = HashMap::new();
3364 for chunk in namespaces.chunks(NAMESPACE_COUNT_CHUNK_SIZE) {
3365 let (where_clause, params) =
3366 build_edge_filter_sql_for_namespaces(chunk, &EdgeFilter::default());
3367 let sql = format!(
3368 "SELECT relation, COUNT(*) FROM graph_edges{} GROUP BY relation",
3369 with_live_endpoints(&where_clause)
3370 );
3371 let mut stmt = conn.prepare(&sql)?;
3372 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
3373 params.iter().map(|p| p.as_ref()).collect();
3374 let rows = stmt.query_map(param_refs.as_slice(), |row| {
3375 let relation_str: String = row.get(0)?;
3376 let count: i64 = row.get(1)?;
3377 Ok((relation_str, count))
3378 })?;
3379 for row in rows {
3380 let (relation_str, count) = row?;
3381 let relation = relation_str.parse::<EdgeRelation>().map_err(|e| {
3382 rusqlite::Error::FromSqlConversionFailure(
3383 0,
3384 rusqlite::types::Type::Text,
3385 Box::new(e),
3386 )
3387 })?;
3388 *totals.entry(relation).or_insert(0) += count as u64;
3389 }
3390 }
3391 Ok(totals.into_iter().collect())
3392 })
3393 .await
3394 }
3395
3396 async fn count_edges_by_endpoint_base(&self) -> Result<EdgeEndpointBaseCounts, StorageError> {
3397 let namespace = self.namespace.clone();
3398 self.with_reader("count_edges_by_endpoint_base", move |conn| {
3399 let source_case = endpoint_base_case("source_id");
3400 let target_case = endpoint_base_case("target_id");
3401 let sql = format!(
3402 "SELECT {source_case} AS source_base, {target_case} AS target_base, COUNT(*) \
3403 FROM graph_edges \
3404 WHERE namespace = ?1 AND deleted_at IS NULL AND {LIVE_ENDPOINTS_CONDITION} \
3405 GROUP BY source_base, target_base"
3406 );
3407 let mut stmt = conn.prepare(&sql)?;
3408 let rows = stmt.query_map([&namespace], |row| {
3409 let source: String = row.get(0)?;
3410 let target: String = row.get(1)?;
3411 let count: i64 = row.get(2)?;
3412 Ok((source, target, count))
3413 })?;
3414 let mut counts = EdgeEndpointBaseCounts::default();
3415 for row in rows {
3416 let (source, target, count) = row?;
3417 fold_endpoint_base_row(&mut counts, &source, &target, count as u64);
3418 }
3419 Ok(counts)
3420 })
3421 .await
3422 }
3423
3424 async fn count_edges_by_endpoint_base_in_namespaces(
3425 &self,
3426 namespaces: &[String],
3427 ) -> Result<EdgeEndpointBaseCounts, StorageError> {
3428 let namespaces: Vec<String> = namespaces
3429 .iter()
3430 .cloned()
3431 .collect::<HashSet<_>>()
3432 .into_iter()
3433 .collect();
3434 self.with_reader("count_edges_by_endpoint_base_in_namespaces", move |conn| {
3435 let source_case = endpoint_base_case("source_id");
3436 let target_case = endpoint_base_case("target_id");
3437 let mut counts = EdgeEndpointBaseCounts::default();
3438 for chunk in namespaces.chunks(NAMESPACE_COUNT_CHUNK_SIZE) {
3439 let (where_clause, params) =
3440 build_edge_filter_sql_for_namespaces(chunk, &EdgeFilter::default());
3441 let sql = format!(
3442 "SELECT {source_case} AS source_base, {target_case} AS target_base, COUNT(*) \
3443 FROM graph_edges{} GROUP BY source_base, target_base",
3444 with_live_endpoints(&where_clause)
3445 );
3446 let mut stmt = conn.prepare(&sql)?;
3447 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
3448 params.iter().map(|p| p.as_ref()).collect();
3449 let rows = stmt.query_map(param_refs.as_slice(), |row| {
3450 let source: String = row.get(0)?;
3451 let target: String = row.get(1)?;
3452 let count: i64 = row.get(2)?;
3453 Ok((source, target, count))
3454 })?;
3455 for row in rows {
3456 let (source, target, count) = row?;
3457 fold_endpoint_base_row(&mut counts, &source, &target, count as u64);
3458 }
3459 }
3460 Ok(counts)
3461 })
3462 .await
3463 }
3464
3465 async fn count_edges_by_relation_and_endpoint_base_in_namespaces(
3466 &self,
3467 namespaces: &[String],
3468 ) -> Result<Vec<(EdgeRelation, EdgeEndpointBaseCounts)>, StorageError> {
3469 let namespaces: Vec<String> = namespaces
3470 .iter()
3471 .cloned()
3472 .collect::<HashSet<_>>()
3473 .into_iter()
3474 .collect();
3475 self.with_reader(
3476 "count_edges_by_relation_and_endpoint_base_in_namespaces",
3477 move |conn| {
3478 let mut totals: HashMap<EdgeRelation, EdgeEndpointBaseCounts> = HashMap::new();
3479 for chunk in namespaces.chunks(NAMESPACE_COUNT_CHUNK_SIZE) {
3480 let (where_clause, params) =
3481 build_edge_filter_sql_for_namespaces(chunk, &EdgeFilter::default());
3482 let sql = format!(
3492 "SELECT e.relation, \
3493 CASE WHEN se.id IS NOT NULL THEN 'entity' \
3494 WHEN sn.id IS NOT NULL THEN 'note' ELSE 'none' END AS source_base, \
3495 CASE WHEN te.id IS NOT NULL THEN 'entity' \
3496 WHEN tn.id IS NOT NULL THEN 'note' ELSE 'none' END AS target_base, \
3497 COUNT(*) \
3498 FROM (SELECT relation, source_id, target_id \
3499 FROM graph_edges{where_clause}) e \
3500 LEFT JOIN entities se ON se.id = e.source_id \
3501 LEFT JOIN notes sn ON sn.id = e.source_id \
3502 LEFT JOIN entities te ON te.id = e.target_id \
3503 LEFT JOIN notes tn ON tn.id = e.target_id \
3504 WHERE se.deleted_at IS NULL AND sn.deleted_at IS NULL \
3505 AND te.deleted_at IS NULL AND tn.deleted_at IS NULL \
3506 GROUP BY e.relation, source_base, target_base"
3507 );
3508 let mut stmt = conn.prepare(&sql)?;
3509 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
3510 params.iter().map(|p| p.as_ref()).collect();
3511 let rows = stmt.query_map(param_refs.as_slice(), |row| {
3512 let relation: String = row.get(0)?;
3513 let source: String = row.get(1)?;
3514 let target: String = row.get(2)?;
3515 let count: i64 = row.get(3)?;
3516 Ok((relation, source, target, count))
3517 })?;
3518 for row in rows {
3519 let (relation_str, source, target, count) = row?;
3520 let relation = relation_str.parse::<EdgeRelation>().map_err(|e| {
3521 rusqlite::Error::FromSqlConversionFailure(
3522 0,
3523 rusqlite::types::Type::Text,
3524 Box::new(e),
3525 )
3526 })?;
3527 fold_endpoint_base_row(
3528 totals.entry(relation).or_default(),
3529 &source,
3530 &target,
3531 count as u64,
3532 );
3533 }
3534 }
3535 Ok(totals.into_iter().collect())
3536 },
3537 )
3538 .await
3539 }
3540
3541 async fn query_edges_after(
3542 &self,
3543 filter: EdgeFilter,
3544 after: Option<Uuid>,
3545 limit: u32,
3546 ) -> Result<EdgeSeekPage, StorageError> {
3547 let namespace = self.namespace.clone();
3548 let limit_usize = limit as usize;
3549 let probe_limit_i64 = i64::from(limit) + 1;
3550 self.with_reader("query_edges_after", move |conn| {
3551 let (mut where_clause, mut params) = build_edge_filter_sql(&namespace, &filter);
3552 if let Some(cursor) = after {
3553 params.push(Box::new(cursor.to_string()));
3554 where_clause.push_str(&format!(" AND id > ?{}", params.len()));
3555 }
3556 params.push(Box::new(probe_limit_i64));
3557 let limit_idx = params.len();
3558
3559 let data_sql = format!(
3565 "SELECT namespace, id, source_id, target_id, relation, weight, \
3566 created_at, updated_at, deleted_at, metadata, target_backend \
3567 FROM graph_edges{} ORDER BY id ASC LIMIT ?{}",
3568 where_clause, limit_idx,
3569 );
3570
3571 let mut stmt = conn.prepare(&data_sql)?;
3572 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
3573 params.iter().map(|p| p.as_ref()).collect();
3574 let rows = stmt.query_map(param_refs.as_slice(), read_edge)?;
3575
3576 let mut items = Vec::new();
3577 for row in rows {
3578 items.push(row?);
3579 }
3580 let has_more = items.len() > limit_usize;
3581 if has_more {
3582 items.truncate(limit_usize);
3583 }
3584 let next_after = if has_more {
3585 items.last().map(|e| Uuid::from(e.id))
3586 } else {
3587 None
3588 };
3589
3590 Ok(EdgeSeekPage { items, next_after })
3591 })
3592 .await
3593 }
3594
3595 async fn query_edges_sequence_after(
3596 &self,
3597 filter: EdgeFilter,
3598 after: Option<SeekCursor>,
3599 limit: u32,
3600 ) -> Result<SeekPage<Edge>, StorageError> {
3601 if limit == 0 {
3602 return Ok(SeekPage::default());
3603 }
3604 let namespace = self.namespace.clone();
3605 let limit_usize = limit as usize;
3606 let probe_limit_i64 = i64::from(limit) + 1;
3607 self.with_reader("query_edges_sequence_after", move |conn| {
3608 let (mut where_clause, mut params) = build_edge_filter_sql(&namespace, &filter);
3609 if let Some(cursor) = after {
3610 params.push(Box::new(cursor.sequence));
3611 where_clause.push_str(&format!(" AND graph_edges_seq.seq > ?{}", params.len()));
3612 }
3613 params.push(Box::new(probe_limit_i64));
3614 let limit_idx = params.len();
3615 let sql = format!(
3618 "SELECT namespace, id, source_id, target_id, relation, weight, \
3619 created_at, updated_at, deleted_at, metadata, target_backend, \
3620 graph_edges_seq.seq \
3621 FROM graph_edges_seq CROSS JOIN graph_edges \
3622 ON graph_edges.id = graph_edges_seq.edge_id{where_clause} \
3623 ORDER BY graph_edges_seq.seq ASC LIMIT ?{limit_idx}"
3624 );
3625 let mut stmt = conn.prepare(&sql)?;
3626 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
3627 params.iter().map(|param| param.as_ref()).collect();
3628 let rows = stmt.query_map(param_refs.as_slice(), |row| {
3629 Ok((read_edge(row)?, row.get::<_, i64>(11)?))
3630 })?;
3631 let mut entries = rows.collect::<Result<Vec<_>, _>>()?;
3632 let has_more = entries.len() > limit_usize;
3633 if has_more {
3634 entries.truncate(limit_usize);
3635 }
3636 let next_after = if has_more {
3637 entries.last().map(|(edge, sequence)| SeekCursor {
3638 sequence: *sequence,
3639 id: Uuid::from(edge.id),
3640 })
3641 } else {
3642 None
3643 };
3644 let items = entries.into_iter().map(|(edge, _)| edge).collect();
3645 Ok(SeekPage { items, next_after })
3646 })
3647 .await
3648 }
3649
3650 async fn neighbors(
3651 &self,
3652 node_id: Uuid,
3653 query: NeighborQuery,
3654 ) -> Result<Vec<NeighborHit>, StorageError> {
3655 self.query_neighbors_page("neighbors", node_id, query, None, None)
3656 .await
3657 }
3658
3659 async fn neighbors_page(
3660 &self,
3661 node_id: Uuid,
3662 query: NeighborQuery,
3663 after: Option<NeighborCursor>,
3664 neighbor_kinds: Option<Vec<String>>,
3665 ) -> Result<Vec<NeighborHit>, StorageError> {
3666 self.query_neighbors_page("neighbors_page", node_id, query, after, neighbor_kinds)
3667 .await
3668 }
3669
3670 async fn neighbors_both_directions(
3675 &self,
3676 node_id: Uuid,
3677 query: NeighborQuery,
3678 ) -> Result<Vec<DirectedNeighborHit>, StorageError> {
3679 count_neighbor_select();
3680
3681 let namespace = self.namespace.clone();
3682 let node_str = node_id.to_string();
3683
3684 let counted_queries = Arc::new(std::sync::atomic::AtomicU64::new(0));
3685 let counted_rows = Arc::new(std::sync::atomic::AtomicU64::new(0));
3686 let closure_queries = Arc::clone(&counted_queries);
3687 let closure_rows = Arc::clone(&counted_rows);
3688 let result = self
3689 .with_reader("neighbors_both_directions", move |conn| {
3690 let base_out = "SELECT target_id AS node_id, id AS edge_id, relation, weight, \
3691 'out' AS dir \
3692 FROM graph_edges \
3693 WHERE namespace = ?1 AND source_id = ?2 AND deleted_at IS NULL";
3694 let base_in = "SELECT source_id AS node_id, id AS edge_id, relation, weight, \
3695 'in' AS dir \
3696 FROM graph_edges \
3697 WHERE namespace = ?1 AND target_id = ?2 AND deleted_at IS NULL";
3698 let sql = format!("{} UNION ALL {}", base_out, base_in);
3699
3700 let (where_extra, limit_clause, extra_params) =
3701 neighbor_extra_clause(&query, 3, None, None);
3702
3703 let full_sql = format!(
3712 "SELECT node_id, edge_id, relation, weight, dir FROM ({}){} \
3713 ORDER BY weight DESC, node_id ASC, \
3714 CASE dir WHEN 'out' THEN 0 ELSE 1 END ASC, edge_id ASC{}",
3715 sql, where_extra, limit_clause
3716 );
3717
3718 let mut stmt = conn.prepare(&full_sql)?;
3719
3720 let mut all_params: Vec<Box<dyn rusqlite::types::ToSql>> = Vec::new();
3721 all_params.push(Box::new(namespace.clone()));
3722 all_params.push(Box::new(node_str.clone()));
3723 all_params.extend(extra_params);
3724
3725 let param_refs: Vec<&dyn rusqlite::types::ToSql> =
3726 all_params.iter().map(|p| p.as_ref()).collect();
3727
3728 closure_queries.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3729 let rows = stmt.query_map(param_refs.as_slice(), |row| {
3730 let nid_str: String = row.get(0)?;
3731 let eid_str: String = row.get(1)?;
3732 let relation_str: String = row.get(2)?;
3733 let weight: f64 = row.get(3)?;
3734 let dir_str: String = row.get(4)?;
3735 Ok((nid_str, eid_str, relation_str, weight, dir_str))
3736 })?;
3737
3738 let mut hits = Vec::new();
3739 for row in rows {
3740 let (nid_str, eid_str, relation_str, weight, dir_str) = row?;
3741 closure_rows.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3742 let relation = relation_str.parse::<EdgeRelation>().map_err(|e| {
3743 rusqlite::Error::FromSqlConversionFailure(
3744 2,
3745 rusqlite::types::Type::Text,
3746 Box::new(e),
3747 )
3748 })?;
3749 let direction = if dir_str == "out" {
3750 Direction::Out
3751 } else {
3752 Direction::In
3753 };
3754 hits.push(DirectedNeighborHit {
3755 hit: NeighborHit {
3756 node_id: parse_uuid(&nid_str)?,
3757 edge_id: parse_uuid(&eid_str)?,
3758 relation,
3759 weight,
3760 name: None,
3761 kind: None,
3762 entity_type: None,
3763 },
3764 direction,
3765 });
3766 }
3767
3768 Ok(hits)
3769 })
3770 .await;
3771
3772 report_graph_usage(&counted_queries, &counted_rows);
3773 result
3774 }
3775
3776 async fn traverse(&self, request: TraversalRequest) -> Result<Vec<GraphPath>, StorageError> {
3777 request
3778 .validate()
3779 .map_err(|message| StorageError::InvalidInput {
3780 capability: StorageCapability::Graph,
3781 operation: "traverse".into(),
3782 message,
3783 })?;
3784 if request.roots.is_empty() {
3785 return Ok(Vec::new());
3786 }
3787
3788 let budget = request.execution_budget;
3789 let outer_context = khive_storage::capture_request_read_context();
3790 let execution_deadline =
3791 khive_storage::RequestReadDeadline::after(budget.remaining_duration());
3792 khive_storage::scope_request_read_deadline_at(execution_deadline, async {
3793 if outer_context.stop_reason().is_some() {
3794 return Err(StorageError::Timeout {
3795 operation: "traverse".into(),
3796 });
3797 }
3798 if budget.is_expired() {
3799 return Err(traversal_timeout_error(&budget));
3800 }
3801 let mut distinct_roots = HashSet::with_capacity(request.roots.len());
3802 let roots = request
3803 .roots
3804 .iter()
3805 .copied()
3806 .filter(|root| distinct_roots.insert(*root))
3807 .collect::<Vec<_>>();
3808 let opts = request.options;
3809 if self
3810 .index_repair
3811 .as_ref()
3812 .is_some_and(super::index_repair::IndexRepairContext::is_writable)
3813 {
3814 let directions = match opts.direction {
3820 Direction::Out => vec![Direction::Out],
3821 Direction::In => vec![Direction::In],
3822 Direction::Both => vec![Direction::Out, Direction::In],
3823 };
3824 let relation_count = opts.relations.as_ref().map_or(0, Vec::len);
3825 let statements = directions
3826 .into_iter()
3827 .map(|direction| {
3828 traversal_neighbor_sql(direction, relation_count, opts.min_weight.is_some())
3829 })
3830 .collect::<Vec<_>>();
3831 let preflight = self.with_indexed_reader("traverse", move |conn| {
3834 for sql in &statements {
3835 let mut statement = conn.prepare(sql)?;
3836 statement.raw_bind_parameter(3, 0_i64)?;
3837 let _ = statement.raw_query().next()?;
3838 }
3839 Ok(())
3840 });
3841 let preflight_result = tokio::select! {
3842 biased;
3843 _ = outer_context.clone().wait_for_stop() => Err(StorageError::Timeout {
3844 operation: "traverse".into(),
3845 }),
3846 result = tokio::time::timeout_at(execution_deadline.async_at(), preflight) => {
3847 match result {
3848 Ok(result) => result,
3849 Err(_) => Err(traversal_timeout_error(&budget)),
3850 }
3851 },
3852 };
3853 preflight_result.map_err(|error| match error {
3854 StorageError::Timeout { .. }
3855 if budget.is_expired() && outer_context.stop_reason().is_none() =>
3856 {
3857 traversal_timeout_error(&budget)
3858 }
3859 error => error,
3860 })?;
3861 if budget.is_expired() {
3862 return Err(traversal_timeout_error(&budget));
3863 }
3864 }
3865 let include_roots = request.include_roots;
3866 let namespace = self.namespace.clone();
3867 let origin = self.pool.origin();
3868 let closure_budget = budget.clone();
3869 let counted_rows = Arc::new(std::sync::atomic::AtomicU64::new(0));
3875 let counted_queries = Arc::new(std::sync::atomic::AtomicU64::new(0));
3876 let closure_rows = Arc::clone(&counted_rows);
3877 let closure_queries = Arc::clone(&counted_queries);
3878 let result = self
3879 .with_reader("traverse", move |conn| {
3880 Ok(run_bounded_traversal(
3881 conn,
3882 roots,
3883 opts,
3884 include_roots,
3885 namespace,
3886 origin,
3887 closure_budget,
3888 closure_rows.as_ref(),
3889 closure_queries.as_ref(),
3890 ))
3891 })
3892 .await
3893 .and_then(|inner| inner)
3894 .map_err(|error| match error {
3895 StorageError::Timeout { .. }
3896 if budget.is_expired() && outer_context.stop_reason().is_none() =>
3897 {
3898 traversal_timeout_error(&budget)
3899 }
3900 error => error,
3901 });
3902
3903 khive_storage::usage::count(
3912 khive_storage::usage::UsageUnit::DbRoundTrips,
3913 counted_queries.load(std::sync::atomic::Ordering::Relaxed),
3914 );
3915 khive_storage::usage::count(
3916 khive_storage::usage::UsageUnit::GraphHops,
3917 counted_rows.load(std::sync::atomic::Ordering::Relaxed),
3918 );
3919
3920 result
3921 })
3922 .await
3923 }
3924
3925 async fn purge_incident_edges(&self, node_id: Uuid) -> Result<u64, StorageError> {
3926 let statement = purge_incident_edges_statement(node_id);
3930 self.with_writer("purge_incident_edges", move |conn| {
3931 let mut stmt = conn.prepare(&statement.sql)?;
3932 bind_params(&mut stmt, &statement.params)?;
3933 Ok(stmt.raw_execute()? as u64)
3934 })
3935 .await
3936 }
3937}
3938
3939const GRAPH_DDL: &str = include_str!("../../sql/graph-ddl.sql");
3944
3945pub(crate) fn ensure_graph_schema(conn: &rusqlite::Connection) -> Result<(), rusqlite::Error> {
3946 conn.execute_batch(GRAPH_DDL)
3947}
3948
3949#[cfg(test)]
3950#[path = "graph_annotation_tests.rs"]
3951mod annotation_tests;
3952
3953#[cfg(test)]
3954#[path = "graph_tests.rs"]
3955mod tests;
3956
3957#[cfg(test)]
3958#[path = "incident_edge_plan_tests.rs"]
3959mod incident_edge_plan_tests;
3960
3961#[cfg(test)]
3962#[path = "graph_index_repair_tests.rs"]
3963mod index_repair_tests;
3964
3965#[cfg(test)]
3966#[path = "graph_busy_tests.rs"]
3967mod direct_busy_tests;