1use std::{
2 collections::{BTreeMap, BTreeSet},
3 path::Path,
4 sync::{Arc, Mutex, TryLockError},
5 time::{Duration, Instant},
6};
7
8mod code;
9
10use rusqlite::{Connection, OptionalExtension, params};
11
12mod canvas;
13mod canvas_code;
14mod code_graph;
15mod file_index;
16mod helpers;
17mod indexing;
18mod operations;
19mod read_pool;
20mod retrieval;
21mod retry;
22mod schema_columns;
23mod schema_migration;
24mod store_impls;
25
26use crate::{
27 domain::{CommitReceipt, GraphMutationBatch, GraphVersion, SourceScope},
28 storage::{GraphInspection, StorageError, StorageFuture},
29};
30
31#[cfg(test)]
32use crate::{
33 domain::IndexKind,
34 storage::{
35 CodeGraphStore, GraphSearchRequest, GraphStore, IndexStore, MutationLogStore, RetrievalHit,
36 },
37};
38use helpers::{count_rows, source_hash_for_evidence, stable_id, storage_version_range};
39use read_pool::{
40 ReadConnectionPool, lock_any_read_connection, lock_any_read_connection_until,
41 lock_connection_until, try_lock_any_read_connection,
42};
43
44const SQLITE_BUSY_TIMEOUT: Duration = Duration::from_secs(30);
45
46#[derive(Debug, Clone)]
48pub struct SqliteGraphStore {
49 connection: Arc<Mutex<Connection>>,
50 read_pool: Option<Arc<ReadConnectionPool>>,
51}
52
53impl SqliteGraphStore {
54 pub fn open(path: impl AsRef<Path>) -> Result<Self, StorageError> {
56 let path = path.as_ref().to_path_buf();
57 if let Some(parent) = path.parent() {
58 std::fs::create_dir_all(parent)?;
59 }
60
61 let connection = Connection::open(&path)?;
62 configure_connection(&connection)?;
63 schema_migration::prepare_existing_database(&connection)?;
64 initialize_schema(&connection)?;
65 let read_pool = ReadConnectionPool::open(&path)?;
66
67 Ok(Self {
68 connection: Arc::new(Mutex::new(connection)),
69 read_pool: Some(Arc::new(read_pool)),
70 })
71 }
72
73 pub fn open_in_memory() -> Result<Self, StorageError> {
75 let connection = Connection::open_in_memory()?;
76 configure_connection(&connection)?;
77 initialize_schema(&connection)?;
78
79 Ok(Self {
80 connection: Arc::new(Mutex::new(connection)),
81 read_pool: None,
82 })
83 }
84
85 pub(super) fn run<T, F>(&self, operation: F) -> StorageFuture<'_, T>
86 where
87 T: Send + 'static,
88 F: FnOnce(&mut Connection) -> Result<T, StorageError> + Send + 'static,
89 {
90 let connection = Arc::clone(&self.connection);
91
92 Box::pin(async move {
93 tokio::task::spawn_blocking(move || {
94 let mut guard = connection.lock().map_err(|_| StorageError::LockPoisoned)?;
95
96 operation(&mut guard)
97 })
98 .await?
99 })
100 }
101
102 pub(super) fn run_read<T, F>(&self, operation: F) -> StorageFuture<'_, T>
103 where
104 T: Send + 'static,
105 F: FnOnce(&mut Connection) -> Result<T, StorageError> + Send + 'static,
106 {
107 if let Some(read_pool) = &self.read_pool {
108 let connections = read_pool.connections();
109 return Box::pin(async move {
110 tokio::task::spawn_blocking(move || {
111 let mut guard = lock_any_read_connection(&connections)?;
112
113 operation(&mut guard)
114 })
115 .await?
116 });
117 }
118
119 self.run(operation)
120 }
121
122 pub(super) fn run_read_until<T, F>(
123 &self,
124 deadline: Instant,
125 timeout_message: &'static str,
126 operation: F,
127 ) -> StorageFuture<'_, T>
128 where
129 T: Send + 'static,
130 F: FnOnce(&mut Connection) -> Result<T, StorageError> + Send + 'static,
131 {
132 if let Some(read_pool) = &self.read_pool {
133 let connections = read_pool.connections();
134 return Box::pin(async move {
135 tokio::task::spawn_blocking(move || {
136 let mut guard =
137 lock_any_read_connection_until(&connections, deadline, timeout_message)?;
138
139 operation(&mut guard)
140 })
141 .await?
142 });
143 }
144
145 let connection = Arc::clone(&self.connection);
146 Box::pin(async move {
147 tokio::task::spawn_blocking(move || {
148 let mut guard = lock_connection_until(&connection, deadline, timeout_message)?;
149
150 operation(&mut guard)
151 })
152 .await?
153 })
154 }
155
156 pub(super) fn try_run_read<T, F>(&self, operation: F) -> StorageFuture<'_, T>
157 where
158 T: Send + 'static,
159 F: FnOnce(&mut Connection) -> Result<T, StorageError> + Send + 'static,
160 {
161 if let Some(read_pool) = &self.read_pool {
162 let connections = read_pool.connections();
163 return Box::pin(async move {
164 tokio::task::spawn_blocking(move || {
165 let mut guard = try_lock_any_read_connection(&connections)?;
166
167 operation(&mut guard)
168 })
169 .await?
170 });
171 }
172
173 let connection = Arc::clone(&self.connection);
174 Box::pin(async move {
175 tokio::task::spawn_blocking(move || {
176 let mut guard = match connection.try_lock() {
177 Ok(guard) => guard,
178 Err(TryLockError::Poisoned(_)) => return Err(StorageError::LockPoisoned),
179 Err(TryLockError::WouldBlock) => {
180 return Err(StorageError::Busy(
181 "sqlite write connection is currently occupied".to_owned(),
182 ));
183 }
184 };
185
186 operation(&mut guard)
187 })
188 .await?
189 })
190 }
191}
192
193fn configure_connection(connection: &Connection) -> Result<(), StorageError> {
194 connection.busy_timeout(SQLITE_BUSY_TIMEOUT)?;
195
196 Ok(())
197}
198
199fn initialize_schema(connection: &Connection) -> Result<(), StorageError> {
200 retry::retry_sqlite_transient(|| initialize_schema_once(connection))
201}
202
203fn initialize_schema_once(connection: &Connection) -> Result<(), StorageError> {
204 connection.execute_batch(
205 "
206 PRAGMA foreign_keys = ON;
207 PRAGMA journal_mode = WAL;
208
209 CREATE TABLE IF NOT EXISTS graph_state (
210 id INTEGER PRIMARY KEY CHECK (id = 1),
211 graph_version INTEGER NOT NULL
212 );
213
214 INSERT OR IGNORE INTO graph_state (id, graph_version) VALUES (1, 0);
215
216 CREATE TABLE IF NOT EXISTS entities (
217 id TEXT PRIMARY KEY,
218 label TEXT NOT NULL,
219 created_graph_version INTEGER NOT NULL
220 );
221
222 CREATE TABLE IF NOT EXISTS evidence (
223 id TEXT PRIMARY KEY,
224 source_scope TEXT NOT NULL,
225 source_path TEXT,
226 span_start_byte INTEGER,
227 span_end_byte INTEGER,
228 span_start_line INTEGER,
229 span_end_line INTEGER,
230 content TEXT NOT NULL,
231 confidence_basis_points INTEGER NOT NULL DEFAULT 10000,
232 status TEXT NOT NULL DEFAULT 'accepted',
233 modality TEXT NOT NULL DEFAULT 'text_span',
234 source_uri TEXT,
235 source_hash TEXT,
236 media_hash TEXT,
237 extractor TEXT,
238 extractor_version TEXT,
239 observed_at TEXT,
240 parent_evidence_id TEXT,
241 layout_page_number INTEGER,
242 layout_x INTEGER,
243 layout_y INTEGER,
244 layout_width INTEGER,
245 layout_height INTEGER,
246 embedding_model TEXT,
247 embedding_dimension INTEGER,
248 extraction_status TEXT NOT NULL DEFAULT 'succeeded',
249 extraction_message TEXT,
250 created_graph_version INTEGER NOT NULL
251 );
252
253 CREATE TABLE IF NOT EXISTS evidence_entities (
254 evidence_id TEXT NOT NULL,
255 entity_id TEXT NOT NULL,
256 PRIMARY KEY (evidence_id, entity_id),
257 FOREIGN KEY (evidence_id) REFERENCES evidence(id) ON DELETE CASCADE,
258 FOREIGN KEY (entity_id) REFERENCES entities(id) ON DELETE CASCADE
259 );
260
261 CREATE TABLE IF NOT EXISTS graph_mutations (
262 graph_version INTEGER PRIMARY KEY,
263 evidence_count INTEGER NOT NULL,
264 entity_count INTEGER NOT NULL,
265 relation_count INTEGER NOT NULL DEFAULT 0,
266 claim_count INTEGER NOT NULL DEFAULT 0,
267 event_count INTEGER NOT NULL DEFAULT 0,
268 affected_scopes_json TEXT NOT NULL DEFAULT '[]',
269 affected_entity_ids_json TEXT NOT NULL DEFAULT '[]',
270 evidence_ids_json TEXT NOT NULL DEFAULT '[]',
271 source_hashes_json TEXT NOT NULL DEFAULT '[]'
272 );
273
274 CREATE TABLE IF NOT EXISTS graph_relations (
275 id TEXT PRIMARY KEY,
276 source_entity_id TEXT NOT NULL,
277 relation_type TEXT NOT NULL,
278 target_entity_id TEXT NOT NULL,
279 evidence_ids_json TEXT NOT NULL,
280 confidence_basis_points INTEGER NOT NULL,
281 status TEXT NOT NULL,
282 valid_from_graph_version INTEGER NOT NULL,
283 valid_until_graph_version INTEGER,
284 created_graph_version INTEGER NOT NULL,
285 FOREIGN KEY (source_entity_id) REFERENCES entities(id),
286 FOREIGN KEY (target_entity_id) REFERENCES entities(id)
287 );
288
289 CREATE TABLE IF NOT EXISTS graph_claims (
290 id TEXT PRIMARY KEY,
291 subject_entity_id TEXT NOT NULL,
292 predicate TEXT NOT NULL,
293 object TEXT NOT NULL,
294 evidence_ids_json TEXT NOT NULL,
295 confidence_basis_points INTEGER NOT NULL,
296 status TEXT NOT NULL,
297 valid_from_graph_version INTEGER NOT NULL,
298 valid_until_graph_version INTEGER,
299 created_graph_version INTEGER NOT NULL,
300 FOREIGN KEY (subject_entity_id) REFERENCES entities(id)
301 );
302
303 CREATE TABLE IF NOT EXISTS graph_events (
304 id TEXT PRIMARY KEY,
305 event_type TEXT NOT NULL,
306 occurred_at TEXT,
307 evidence_ids_json TEXT NOT NULL,
308 confidence_basis_points INTEGER NOT NULL,
309 status TEXT NOT NULL,
310 valid_from_graph_version INTEGER NOT NULL,
311 valid_until_graph_version INTEGER,
312 created_graph_version INTEGER NOT NULL
313 );
314
315 CREATE TABLE IF NOT EXISTS graph_event_entities (
316 event_id TEXT NOT NULL,
317 entity_id TEXT NOT NULL,
318 PRIMARY KEY (event_id, entity_id),
319 FOREIGN KEY (event_id) REFERENCES graph_events(id) ON DELETE CASCADE,
320 FOREIGN KEY (entity_id) REFERENCES entities(id)
321 );
322
323 CREATE TABLE IF NOT EXISTS graph_fact_evidence (
324 fact_kind TEXT NOT NULL,
325 fact_id TEXT NOT NULL,
326 evidence_id TEXT NOT NULL,
327 PRIMARY KEY (fact_kind, fact_id, evidence_id),
328 FOREIGN KEY (evidence_id) REFERENCES evidence(id) ON DELETE CASCADE
329 );
330
331 CREATE INDEX IF NOT EXISTS graph_fact_evidence_by_evidence
332 ON graph_fact_evidence(evidence_id, fact_kind);
333 ",
334 )?;
335 schema_columns::ensure_core_schema_columns(connection)?;
336 code::initialize_code_schema(connection)?;
337 indexing::initialize_schema(connection)?;
338 code_graph::initialize_schema(connection)?;
339 operations::initialize_schema(connection)?;
340 file_index::initialize_schema(connection)?;
341 backfill_fact_evidence_links(connection)?;
342 retrieval::initialize_schema(connection)?;
343
344 Ok(())
345}
346
347fn backfill_fact_evidence_links(connection: &Connection) -> Result<(), StorageError> {
348 backfill_fact_evidence_kind(connection, "relation", "graph_relations")?;
349 backfill_fact_evidence_kind(connection, "claim", "graph_claims")?;
350 backfill_fact_evidence_kind(connection, "event", "graph_events")?;
351
352 Ok(())
353}
354
355fn backfill_fact_evidence_kind(
356 connection: &Connection,
357 fact_kind: &'static str,
358 table: &'static str,
359) -> Result<(), StorageError> {
360 let mut statement =
361 connection.prepare(&format!("SELECT id, evidence_ids_json FROM {table}"))?;
362 let rows = statement.query_map([], |row| {
363 Ok((row.get::<_, String>(0)?, row.get::<_, String>(1)?))
364 })?;
365 let facts = rows
366 .collect::<Result<Vec<_>, _>>()
367 .map_err(StorageError::from)?;
368 drop(statement);
369
370 for (fact_id, evidence_json) in facts {
371 let evidence_ids: Vec<String> = serde_json::from_str(&evidence_json)
372 .map_err(|error| StorageError::InvalidInput(error.to_string()))?;
373 for evidence_id in evidence_ids {
374 connection.execute(
375 "
376 INSERT OR IGNORE INTO graph_fact_evidence (fact_kind, fact_id, evidence_id)
377 SELECT ?1, ?2, e.id
378 FROM evidence e
379 WHERE e.id = ?3
380 ",
381 params![fact_kind, fact_id, evidence_id],
382 )?;
383 }
384 }
385
386 Ok(())
387}
388
389fn commit_batch(
390 connection: &mut Connection,
391 batch: GraphMutationBatch,
392) -> Result<CommitReceipt, StorageError> {
393 let transaction = connection.transaction()?;
394 let current = current_graph_version_in_transaction(&transaction)?;
395 let next = GraphVersion::new(current.get() + 1);
396 let evidence_count = batch.evidence.len();
397 let relation_count = batch.relations.len();
398 let claim_count = batch.claims.len();
399 let event_count = batch.events.len();
400 let mut affected_entity_ids = BTreeSet::new();
401 let mut affected_scopes = BTreeSet::new();
402 let mut evidence_ids = BTreeSet::new();
403 let mut source_hashes = BTreeSet::new();
404 let batch_evidence_scopes = batch
405 .evidence
406 .iter()
407 .map(|evidence| {
408 (
409 evidence.id.clone(),
410 evidence.source_scope.as_str().to_owned(),
411 )
412 })
413 .collect::<BTreeMap<_, _>>();
414
415 for evidence in batch.evidence {
416 let evidence_id = evidence.id;
417 let source_scope = evidence.source_scope;
418 let source_scope_text = source_scope.as_str().to_owned();
419 let source_path = evidence.source_path;
420 let span = evidence.span;
421 let content = evidence.content;
422 let entity_labels = evidence.entity_labels;
423 let extraction = evidence.extraction;
424 let derived_source_hash = source_hash_for_evidence(
425 &extraction,
426 &source_scope_text,
427 source_path.as_deref(),
428 &content,
429 );
430 if let Some(previous_scope) = evidence_scope(&transaction, &evidence_id)? {
431 affected_scopes.insert(previous_scope);
432 }
433 affected_scopes.insert(source_scope_text.clone());
434 evidence_ids.insert(evidence_id.clone());
435 source_hashes.insert(derived_source_hash.clone());
436 if let Some(media_hash) = &extraction.media_hash {
437 source_hashes.insert(media_hash.clone());
438 }
439 if let Some(parent_evidence_id) = extraction.parent_evidence_id.as_deref() {
440 validate_parent_evidence(
441 &transaction,
442 &batch_evidence_scopes,
443 &evidence_id,
444 &source_scope_text,
445 parent_evidence_id,
446 )?;
447 }
448 let layout_region = extraction.layout_region;
449 transaction.execute(
450 "INSERT INTO evidence (
451 id, source_scope, source_path, span_start_byte, span_end_byte,
452 span_start_line, span_end_line, content, confidence_basis_points,
453 status, modality, source_uri, source_hash, media_hash, extractor,
454 extractor_version, observed_at, parent_evidence_id, layout_page_number,
455 layout_x, layout_y, layout_width, layout_height, embedding_model,
456 embedding_dimension, extraction_status, extraction_message, created_graph_version
457 )
458 VALUES (
459 ?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14,
460 ?15, ?16, ?17, ?18, ?19, ?20, ?21, ?22, ?23, ?24, ?25, ?26,
461 ?27, ?28
462 )
463 ON CONFLICT(id) DO UPDATE SET
464 source_scope = excluded.source_scope,
465 source_path = excluded.source_path,
466 span_start_byte = excluded.span_start_byte,
467 span_end_byte = excluded.span_end_byte,
468 span_start_line = excluded.span_start_line,
469 span_end_line = excluded.span_end_line,
470 content = excluded.content,
471 confidence_basis_points = excluded.confidence_basis_points,
472 status = excluded.status,
473 modality = excluded.modality,
474 source_uri = excluded.source_uri,
475 source_hash = excluded.source_hash,
476 media_hash = excluded.media_hash,
477 extractor = excluded.extractor,
478 extractor_version = excluded.extractor_version,
479 observed_at = excluded.observed_at,
480 parent_evidence_id = excluded.parent_evidence_id,
481 layout_page_number = excluded.layout_page_number,
482 layout_x = excluded.layout_x,
483 layout_y = excluded.layout_y,
484 layout_width = excluded.layout_width,
485 layout_height = excluded.layout_height,
486 embedding_model = excluded.embedding_model,
487 embedding_dimension = excluded.embedding_dimension,
488 extraction_status = excluded.extraction_status,
489 extraction_message = excluded.extraction_message,
490 created_graph_version = excluded.created_graph_version",
491 params![
492 &evidence_id,
493 &source_scope_text,
494 source_path.as_deref(),
495 span.map(|value| value.start_byte),
496 span.map(|value| value.end_byte),
497 span.map(|value| value.start_line),
498 span.map(|value| value.end_line),
499 &content,
500 evidence.confidence.basis_points,
501 evidence.status.as_str(),
502 extraction.modality.as_str(),
503 extraction.source_uri.as_deref(),
504 &derived_source_hash,
505 extraction.media_hash.as_deref(),
506 extraction.extractor.as_deref(),
507 extraction.extractor_version.as_deref(),
508 extraction.observed_at.as_deref(),
509 extraction.parent_evidence_id.as_deref(),
510 layout_region.map(|region| region.page_number),
511 layout_region.map(|region| region.x),
512 layout_region.map(|region| region.y),
513 layout_region.map(|region| region.width),
514 layout_region.map(|region| region.height),
515 extraction.embedding_model.as_deref(),
516 extraction.embedding_dimension.map(i64::from),
517 extraction.diagnostic.status.as_str(),
518 extraction.diagnostic.message.as_deref(),
519 next.get()
520 ],
521 )?;
522
523 transaction.execute(
524 "DELETE FROM evidence_entities WHERE evidence_id = ?1",
525 params![&evidence_id],
526 )?;
527
528 for label in &entity_labels {
529 let entity_id = upsert_entity(&transaction, label, next)?;
530 transaction.execute(
531 "INSERT OR IGNORE INTO evidence_entities (evidence_id, entity_id)
532 VALUES (?1, ?2)",
533 params![evidence_id, entity_id],
534 )?;
535 affected_entity_ids.insert(entity_id);
536 }
537 retrieval::replace_evidence_document(
538 &transaction,
539 retrieval::EvidenceDocumentInput {
540 evidence_id: &evidence_id,
541 source_scope: &source_scope_text,
542 source_path: source_path.as_deref(),
543 entity_labels: &entity_labels,
544 content: &content,
545 status: evidence.status,
546 extraction: &extraction,
547 source_hash: &derived_source_hash,
548 graph_version: next.get(),
549 },
550 )?;
551 }
552
553 for relation in batch.relations {
554 validate_evidence_references(&transaction, &relation.source_scope, &relation.evidence_ids)?;
555 evidence_ids.extend(relation.evidence_ids.iter().cloned());
556 let source_entity_id = upsert_entity(&transaction, &relation.source_entity_label, next)?;
557 let target_entity_id = upsert_entity(&transaction, &relation.target_entity_label, next)?;
558 let version_range = storage_version_range(relation.version_range, next);
559 affected_entity_ids.insert(source_entity_id.clone());
560 affected_entity_ids.insert(target_entity_id.clone());
561 transaction.execute(
562 "
563 INSERT INTO graph_relations (
564 id, source_entity_id, relation_type, target_entity_id,
565 evidence_ids_json, confidence_basis_points, status,
566 valid_from_graph_version, valid_until_graph_version, created_graph_version
567 )
568 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)
569 ON CONFLICT(id) DO UPDATE SET
570 source_entity_id = excluded.source_entity_id,
571 relation_type = excluded.relation_type,
572 target_entity_id = excluded.target_entity_id,
573 evidence_ids_json = excluded.evidence_ids_json,
574 confidence_basis_points = excluded.confidence_basis_points,
575 status = excluded.status,
576 valid_from_graph_version = excluded.valid_from_graph_version,
577 valid_until_graph_version = excluded.valid_until_graph_version,
578 created_graph_version = excluded.created_graph_version
579 ",
580 params![
581 relation.id.as_str(),
582 source_entity_id,
583 relation.relation_type,
584 target_entity_id,
585 evidence_ids_json(&relation.evidence_ids)?,
586 relation.confidence.basis_points,
587 relation.status.as_str(),
588 version_range.valid_from.get(),
589 version_range.valid_until.map(GraphVersion::get),
590 next.get(),
591 ],
592 )?;
593 replace_fact_evidence_links(
594 &transaction,
595 "relation",
596 &relation.id,
597 &relation.evidence_ids,
598 )?;
599 }
600
601 for claim in batch.claims {
602 validate_evidence_references(&transaction, &claim.source_scope, &claim.evidence_ids)?;
603 evidence_ids.extend(claim.evidence_ids.iter().cloned());
604 let subject_entity_id = upsert_entity(&transaction, &claim.subject_entity_label, next)?;
605 let version_range = storage_version_range(claim.version_range, next);
606 affected_entity_ids.insert(subject_entity_id.clone());
607 transaction.execute(
608 "
609 INSERT INTO graph_claims (
610 id, subject_entity_id, predicate, object, evidence_ids_json,
611 confidence_basis_points, status, valid_from_graph_version,
612 valid_until_graph_version, created_graph_version
613 )
614 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)
615 ON CONFLICT(id) DO UPDATE SET
616 subject_entity_id = excluded.subject_entity_id,
617 predicate = excluded.predicate,
618 object = excluded.object,
619 evidence_ids_json = excluded.evidence_ids_json,
620 confidence_basis_points = excluded.confidence_basis_points,
621 status = excluded.status,
622 valid_from_graph_version = excluded.valid_from_graph_version,
623 valid_until_graph_version = excluded.valid_until_graph_version,
624 created_graph_version = excluded.created_graph_version
625 ",
626 params![
627 claim.id.as_str(),
628 subject_entity_id,
629 claim.predicate,
630 claim.object,
631 evidence_ids_json(&claim.evidence_ids)?,
632 claim.confidence.basis_points,
633 claim.status.as_str(),
634 version_range.valid_from.get(),
635 version_range.valid_until.map(GraphVersion::get),
636 next.get(),
637 ],
638 )?;
639 replace_fact_evidence_links(&transaction, "claim", &claim.id, &claim.evidence_ids)?;
640 }
641
642 for event in batch.events {
643 validate_evidence_references(&transaction, &event.source_scope, &event.evidence_ids)?;
644 evidence_ids.extend(event.evidence_ids.iter().cloned());
645 let version_range = storage_version_range(event.version_range, next);
646 transaction.execute(
647 "
648 INSERT INTO graph_events (
649 id, event_type, occurred_at, evidence_ids_json,
650 confidence_basis_points, status, valid_from_graph_version,
651 valid_until_graph_version, created_graph_version
652 )
653 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9)
654 ON CONFLICT(id) DO UPDATE SET
655 event_type = excluded.event_type,
656 occurred_at = excluded.occurred_at,
657 evidence_ids_json = excluded.evidence_ids_json,
658 confidence_basis_points = excluded.confidence_basis_points,
659 status = excluded.status,
660 valid_from_graph_version = excluded.valid_from_graph_version,
661 valid_until_graph_version = excluded.valid_until_graph_version,
662 created_graph_version = excluded.created_graph_version
663 ",
664 params![
665 event.id.as_str(),
666 event.event_type,
667 event.occurred_at,
668 evidence_ids_json(&event.evidence_ids)?,
669 event.confidence.basis_points,
670 event.status.as_str(),
671 version_range.valid_from.get(),
672 version_range.valid_until.map(GraphVersion::get),
673 next.get(),
674 ],
675 )?;
676 replace_fact_evidence_links(&transaction, "event", &event.id, &event.evidence_ids)?;
677 transaction.execute(
678 "DELETE FROM graph_event_entities WHERE event_id = ?1",
679 params![event.id],
680 )?;
681 for label in event.entity_labels {
682 let entity_id = upsert_entity(&transaction, &label, next)?;
683 affected_entity_ids.insert(entity_id.clone());
684 transaction.execute(
685 "INSERT OR IGNORE INTO graph_event_entities (event_id, entity_id)
686 VALUES (?1, ?2)",
687 params![event.id, entity_id],
688 )?;
689 }
690 }
691
692 transaction.execute(
693 "
694 DELETE FROM entities
695 WHERE id NOT IN (SELECT entity_id FROM evidence_entities)
696 AND id NOT IN (SELECT source_entity_id FROM graph_relations)
697 AND id NOT IN (SELECT target_entity_id FROM graph_relations)
698 AND id NOT IN (SELECT subject_entity_id FROM graph_claims)
699 AND id NOT IN (SELECT entity_id FROM graph_event_entities)
700 ",
701 [],
702 )?;
703
704 let entity_count = affected_entity_ids.len();
705 add_scopes_for_evidence_ids(&transaction, &evidence_ids, &mut affected_scopes)?;
706 if affected_scopes.is_empty()
707 && (evidence_count > 0 || relation_count > 0 || claim_count > 0 || event_count > 0)
708 {
709 affected_scopes.insert(indexing::DEFAULT_SCOPE.to_owned());
710 }
711 let affected_scopes = affected_scopes.into_iter().collect::<Vec<_>>();
712 let affected_entity_ids = affected_entity_ids.into_iter().collect::<Vec<_>>();
713 let affected_scopes_json = indexing::json_array(affected_scopes.clone())?;
714 let affected_entity_ids_json = indexing::json_array(affected_entity_ids.clone())?;
715 let evidence_ids_json = indexing::json_array(evidence_ids)?;
716 let source_hashes_json = indexing::json_array(source_hashes)?;
717 transaction.execute(
718 "INSERT INTO graph_mutations (
719 graph_version, evidence_count, entity_count, relation_count, claim_count, event_count,
720 affected_scopes_json, affected_entity_ids_json, evidence_ids_json, source_hashes_json
721 )
722 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)",
723 params![
724 next.get(),
725 evidence_count,
726 entity_count,
727 relation_count,
728 claim_count,
729 event_count,
730 affected_scopes_json,
731 affected_entity_ids_json,
732 evidence_ids_json,
733 source_hashes_json
734 ],
735 )?;
736 indexing::mark_mutation_cursors_stale(&transaction, &affected_scopes)?;
737 transaction.execute(
738 "UPDATE graph_state SET graph_version = ?1 WHERE id = 1",
739 params![next.get()],
740 )?;
741 transaction.execute("UPDATE index_status SET state = 'stale'", [])?;
742 transaction.commit()?;
743
744 Ok(CommitReceipt {
745 graph_version: next,
746 evidence_count,
747 entity_count,
748 relation_count,
749 claim_count,
750 event_count,
751 })
752}
753
754fn inspect_graph(connection: &mut Connection) -> Result<GraphInspection, StorageError> {
755 Ok(GraphInspection {
756 graph_version: current_graph_version(connection)?,
757 entity_count: count_rows(connection, "entities")?,
758 evidence_count: count_rows(connection, "evidence")?,
759 relation_count: count_rows(connection, "graph_relations")?,
760 claim_count: count_rows(connection, "graph_claims")?,
761 event_count: count_rows(connection, "graph_events")?,
762 mutation_count: count_rows(connection, "graph_mutations")?,
763 code_file_count: count_rows(connection, "code_files")?,
764 code_symbol_count: count_rows(connection, "code_symbols")?,
765 code_reference_count: count_rows(connection, "code_references")?,
766 code_chunk_count: count_rows(connection, "code_chunks")?,
767 code_parse_status_counts: code_graph::parse_status_counts(connection)?,
768 })
769}
770
771fn current_graph_version(connection: &mut Connection) -> Result<GraphVersion, StorageError> {
772 current_graph_version_in_transaction(connection)
773}
774
775fn current_graph_version_in_transaction(
776 connection: &Connection,
777) -> Result<GraphVersion, StorageError> {
778 let value = connection.query_row(
779 "SELECT graph_version FROM graph_state WHERE id = 1",
780 [],
781 |row| row.get::<_, u64>(0),
782 )?;
783
784 Ok(GraphVersion::new(value))
785}
786
787fn upsert_entity(
788 transaction: &rusqlite::Transaction<'_>,
789 label: &str,
790 graph_version: GraphVersion,
791) -> Result<String, StorageError> {
792 let entity_id = stable_id("entity", label);
793 transaction.execute(
794 "INSERT OR IGNORE INTO entities (id, label, created_graph_version)
795 VALUES (?1, ?2, ?3)",
796 params![entity_id, label, graph_version.get()],
797 )?;
798
799 Ok(entity_id)
800}
801
802fn add_scopes_for_evidence_ids(
803 connection: &Connection,
804 evidence_ids: &BTreeSet<String>,
805 affected_scopes: &mut BTreeSet<String>,
806) -> Result<(), StorageError> {
807 for evidence_id in evidence_ids {
808 if let Some(scope) = evidence_scope(connection, evidence_id)? {
809 affected_scopes.insert(scope);
810 }
811 }
812
813 Ok(())
814}
815
816fn evidence_scope(
817 connection: &Connection,
818 evidence_id: &str,
819) -> Result<Option<String>, StorageError> {
820 connection
821 .query_row(
822 "SELECT source_scope FROM evidence WHERE id = ?1",
823 params![evidence_id],
824 |row| row.get::<_, String>(0),
825 )
826 .optional()
827 .map_err(StorageError::from)
828}
829
830fn validate_parent_evidence(
831 connection: &Connection,
832 batch_evidence_scopes: &BTreeMap<String, String>,
833 evidence_id: &str,
834 source_scope: &str,
835 parent_evidence_id: &str,
836) -> Result<(), StorageError> {
837 if parent_evidence_id == evidence_id {
838 return Err(StorageError::InvalidInput(
839 "parent evidence id must reference a different evidence record".to_owned(),
840 ));
841 }
842 let parent_scope = if let Some(scope) = batch_evidence_scopes.get(parent_evidence_id) {
843 Some(scope.clone())
844 } else {
845 evidence_scope(connection, parent_evidence_id).map_err(|error| {
846 StorageError::InvalidInput(format!(
847 "parent evidence id '{parent_evidence_id}' could not be validated: {error}"
848 ))
849 })?
850 };
851
852 match parent_scope {
853 Some(parent_scope) if parent_scope == source_scope => Ok(()),
854 Some(parent_scope) => Err(StorageError::InvalidInput(format!(
855 "parent evidence id '{parent_evidence_id}' belongs to source scope \
856 '{parent_scope}' instead of '{source_scope}'"
857 ))),
858 None => Err(StorageError::InvalidInput(format!(
859 "parent evidence id '{parent_evidence_id}' does not exist in source scope \
860 '{source_scope}'"
861 ))),
862 }
863}
864
865fn evidence_ids_json(evidence_ids: &[String]) -> Result<String, StorageError> {
866 serde_json::to_string(evidence_ids)
867 .map_err(|error| StorageError::InvalidInput(error.to_string()))
868}
869
870fn validate_evidence_references(
871 transaction: &rusqlite::Transaction<'_>,
872 source_scope: &SourceScope,
873 evidence_ids: &[String],
874) -> Result<(), StorageError> {
875 for evidence_id in evidence_ids {
876 let actual_scope = transaction
877 .query_row(
878 "SELECT source_scope FROM evidence WHERE id = ?1",
879 params![evidence_id],
880 |row| row.get::<_, String>(0),
881 )
882 .optional()?;
883 let Some(actual_scope) = actual_scope else {
884 return Err(StorageError::InvalidInput(format!(
885 "structured fact references unknown evidence id '{evidence_id}'"
886 )));
887 };
888 if actual_scope != source_scope.as_str() {
889 return Err(StorageError::InvalidInput(format!(
890 "structured fact references evidence id '{evidence_id}' from source scope \
891 '{actual_scope}' instead of '{}'",
892 source_scope.as_str()
893 )));
894 }
895 }
896
897 Ok(())
898}
899
900fn replace_fact_evidence_links(
901 transaction: &rusqlite::Transaction<'_>,
902 fact_kind: &'static str,
903 fact_id: &str,
904 evidence_ids: &[String],
905) -> Result<(), StorageError> {
906 transaction.execute(
907 "DELETE FROM graph_fact_evidence WHERE fact_kind = ?1 AND fact_id = ?2",
908 params![fact_kind, fact_id],
909 )?;
910 for evidence_id in evidence_ids {
911 transaction.execute(
912 "INSERT OR IGNORE INTO graph_fact_evidence (fact_kind, fact_id, evidence_id)
913 VALUES (?1, ?2, ?3)",
914 params![fact_kind, fact_id, evidence_id],
915 )?;
916 }
917
918 Ok(())
919}
920
921#[cfg(test)]
922mod metadata_tests;
923
924#[cfg(test)]
925mod graph_tests;
926
927#[cfg(test)]
928mod index_refresh_queue_tests;
929
930#[cfg(test)]
931mod index_schema_migration_tests;
932
933#[cfg(test)]
934mod graphrag_phase4_tests;
935
936#[cfg(test)]
937mod index_refresh_tests;
938
939#[cfg(test)]
940#[path = "sqlite/health_tests.rs"]
941mod health_tests;
942
943#[cfg(test)]
944mod operations_tests;