Skip to main content

relay_knowledge/storage/
sqlite.rs

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