1use std::path::{Path, PathBuf};
10use std::time::Duration;
11
12use rusqlite::OptionalExtension;
13
14const CURRENT_SCHEMA_VERSION: i64 = 2;
15
16const SCHEMA_V2_SQL: &str = r#"
17CREATE TABLE IF NOT EXISTS cea_nodes (
18 node_id TEXT PRIMARY KEY,
19 node_kind TEXT NOT NULL,
20 sig_json TEXT NOT NULL,
21 first_seen TEXT NOT NULL,
22 last_seen TEXT NOT NULL
23);
24CREATE TABLE IF NOT EXISTS cea_edges (
25 edge_id TEXT PRIMARY KEY,
26 cause_node_id TEXT NOT NULL,
27 effect_node_id TEXT NOT NULL,
28 weight REAL NOT NULL,
29 count INTEGER NOT NULL,
30 confidence REAL NOT NULL,
31 alpha REAL NOT NULL,
32 beta REAL NOT NULL,
33 version_id TEXT NOT NULL,
34 last_seen TEXT NOT NULL,
35 UNIQUE(cause_node_id, effect_node_id, version_id)
36);
37CREATE INDEX IF NOT EXISTS idx_cea_edges_version_id ON cea_edges(version_id);
38CREATE TABLE IF NOT EXISTS cea_run_log (
39 run_hash TEXT PRIMARY KEY,
40 eval_id TEXT NOT NULL,
41 edges_added INTEGER NOT NULL,
42 edges_updated INTEGER NOT NULL,
43 processed_at TEXT NOT NULL
44);
45"#;
46
47pub struct SqliteCeaStore {
48 path: PathBuf,
49}
50
51impl SqliteCeaStore {
52 pub fn open(path: &Path) -> Result<Self, cea_store::CeaStoreError> {
53 let store = Self {
54 path: path.to_path_buf(),
55 };
56 let _ = store.connect()?;
57 Ok(store)
58 }
59
60 fn connect(&self) -> Result<rusqlite::Connection, cea_store::CeaStoreError> {
61 let connection = rusqlite::Connection::open(&self.path)
62 .map_err(|error| cea_store::CeaStoreError::Backend(error.to_string()))?;
63 configure_connection(&connection)?;
64 ensure_schema(&connection)?;
65 Ok(connection)
66 }
67}
68
69pub struct SqliteCeaStoreConn<'conn> {
70 conn: &'conn rusqlite::Connection,
71}
72
73impl<'conn> SqliteCeaStoreConn<'conn> {
74 pub fn new(conn: &'conn rusqlite::Connection) -> Self {
75 Self { conn }
76 }
77}
78
79pub struct SqliteCeaStoreTx<'tx, 'conn> {
80 tx: &'tx rusqlite::Transaction<'conn>,
81}
82
83impl<'tx, 'conn> SqliteCeaStoreTx<'tx, 'conn> {
84 pub fn new(tx: &'tx rusqlite::Transaction<'conn>) -> Self {
85 Self { tx }
86 }
87}
88
89struct SqliteCeaWriteTx<'tx, 'conn> {
90 tx: &'tx rusqlite::Transaction<'conn>,
91}
92
93fn configure_connection(conn: &rusqlite::Connection) -> Result<(), cea_store::CeaStoreError> {
94 conn.pragma_update(None, "journal_mode", "WAL")
95 .map_err(|error| cea_store::CeaStoreError::Backend(error.to_string()))?;
96 conn.pragma_update(None, "foreign_keys", "ON")
97 .map_err(|error| cea_store::CeaStoreError::Backend(error.to_string()))?;
98 conn.busy_timeout(Duration::from_secs(5))
99 .map_err(|error| cea_store::CeaStoreError::Backend(error.to_string()))?;
100 Ok(())
101}
102
103fn ensure_schema(conn: &rusqlite::Connection) -> Result<(), cea_store::CeaStoreError> {
104 let mut version = read_schema_version(conn)?;
105
106 if version > CURRENT_SCHEMA_VERSION {
107 return Err(cea_store::CeaStoreError::Backend(format!(
108 "cea-sqlite schema version {version} is newer than supported {CURRENT_SCHEMA_VERSION}"
109 )));
110 }
111
112 if version == 0 {
113 version = bootstrap_or_adopt_schema(conn)?;
114 }
115
116 while version < CURRENT_SCHEMA_VERSION {
117 version = migrate_schema(conn, version)?;
118 }
119
120 Ok(())
121}
122
123fn read_schema_version(conn: &rusqlite::Connection) -> Result<i64, cea_store::CeaStoreError> {
124 conn.pragma_query_value(None, "user_version", |row| row.get(0))
125 .map_err(|error| cea_store::CeaStoreError::Backend(error.to_string()))
126}
127
128fn write_schema_version(
129 conn: &rusqlite::Connection,
130 version: i64,
131) -> Result<(), cea_store::CeaStoreError> {
132 conn.pragma_update(None, "user_version", version)
133 .map_err(|error| cea_store::CeaStoreError::Backend(error.to_string()))
134}
135
136fn bootstrap_or_adopt_schema(conn: &rusqlite::Connection) -> Result<i64, cea_store::CeaStoreError> {
137 if !table_exists(conn, "cea_nodes")?
138 && !table_exists(conn, "cea_edges")?
139 && !table_exists(conn, "cea_run_log")?
140 {
141 conn.execute_batch(SCHEMA_V2_SQL)
142 .map_err(|error| cea_store::CeaStoreError::Backend(error.to_string()))?;
143 write_schema_version(conn, CURRENT_SCHEMA_VERSION)?;
144 return Ok(CURRENT_SCHEMA_VERSION);
145 }
146
147 let version = detect_existing_schema_version(conn)?;
148 if version == CURRENT_SCHEMA_VERSION {
149 conn.execute_batch(SCHEMA_V2_SQL)
150 .map_err(|error| cea_store::CeaStoreError::Backend(error.to_string()))?;
151 }
152 write_schema_version(conn, version)?;
153 Ok(version)
154}
155
156fn detect_existing_schema_version(
157 conn: &rusqlite::Connection,
158) -> Result<i64, cea_store::CeaStoreError> {
159 let has_nodes = table_exists(conn, "cea_nodes")?;
160 let has_edges = table_exists(conn, "cea_edges")?;
161 let has_run_log = table_exists(conn, "cea_run_log")?;
162 if !(has_nodes && has_edges && has_run_log) {
163 return Err(cea_store::CeaStoreError::Backend(
164 "incomplete cea-sqlite schema; expected cea_nodes, cea_edges, and cea_run_log"
165 .to_string(),
166 ));
167 }
168
169 if has_legacy_edges_schema(conn)? || has_blocked_run_log_schema(conn)? {
170 return Ok(1);
171 }
172
173 Ok(CURRENT_SCHEMA_VERSION)
174}
175
176fn migrate_schema(
177 conn: &rusqlite::Connection,
178 version: i64,
179) -> Result<i64, cea_store::CeaStoreError> {
180 match version {
181 1 => {
182 migrate_v1_to_v2(conn)?;
183 conn.execute_batch(SCHEMA_V2_SQL)
184 .map_err(|error| cea_store::CeaStoreError::Backend(error.to_string()))?;
185 write_schema_version(conn, CURRENT_SCHEMA_VERSION)?;
186 Ok(CURRENT_SCHEMA_VERSION)
187 }
188 _ => Err(cea_store::CeaStoreError::Backend(format!(
189 "unsupported cea-sqlite schema version migration from {version}"
190 ))),
191 }
192}
193
194fn has_legacy_edges_schema(conn: &rusqlite::Connection) -> Result<bool, cea_store::CeaStoreError> {
195 if !table_exists(conn, "cea_edges")? {
196 return Ok(false);
197 }
198 Ok(!column_exists(conn, "cea_edges", "alpha")? || !column_exists(conn, "cea_edges", "beta")?)
199}
200
201fn has_blocked_run_log_schema(
202 conn: &rusqlite::Connection,
203) -> Result<bool, cea_store::CeaStoreError> {
204 let sql: Option<String> = conn
205 .query_row(
206 "SELECT sql FROM sqlite_master WHERE type = 'table' AND name = 'cea_run_log'",
207 [],
208 |row| row.get(0),
209 )
210 .optional()
211 .map_err(|error| cea_store::CeaStoreError::Backend(error.to_string()))?;
212 Ok(sql
213 .as_deref()
214 .is_some_and(|sql| sql.contains("CHECK(eval_id != 'blocked')")))
215}
216
217fn table_exists(
218 conn: &rusqlite::Connection,
219 table: &str,
220) -> Result<bool, cea_store::CeaStoreError> {
221 let exists: Option<i64> = conn
222 .query_row(
223 "SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = ?1",
224 [table],
225 |row| row.get(0),
226 )
227 .optional()
228 .map_err(|error| cea_store::CeaStoreError::Backend(error.to_string()))?;
229 Ok(exists.is_some())
230}
231
232fn column_exists(
233 conn: &rusqlite::Connection,
234 table: &str,
235 column: &str,
236) -> Result<bool, cea_store::CeaStoreError> {
237 let mut stmt = conn
238 .prepare(&format!("PRAGMA table_info({table})"))
239 .map_err(|error| cea_store::CeaStoreError::Backend(error.to_string()))?;
240 let columns = stmt
241 .query_map([], |row| row.get::<_, String>(1))
242 .map_err(|error| cea_store::CeaStoreError::Backend(error.to_string()))?;
243 for name in columns {
244 if name.map_err(|error| cea_store::CeaStoreError::Backend(error.to_string()))? == column {
245 return Ok(true);
246 }
247 }
248 Ok(false)
249}
250
251fn migrate_v1_to_v2(conn: &rusqlite::Connection) -> Result<(), cea_store::CeaStoreError> {
252 if has_legacy_edges_schema(conn)? {
253 conn.execute_batch(
254 r#"
255CREATE TABLE cea_edges_v2 (
256 edge_id TEXT PRIMARY KEY,
257 cause_node_id TEXT NOT NULL,
258 effect_node_id TEXT NOT NULL,
259 weight REAL NOT NULL,
260 count INTEGER NOT NULL,
261 confidence REAL NOT NULL,
262 alpha REAL NOT NULL,
263 beta REAL NOT NULL,
264 version_id TEXT NOT NULL,
265 last_seen TEXT NOT NULL,
266 UNIQUE(cause_node_id, effect_node_id, version_id)
267);
268INSERT INTO cea_edges_v2 (
269 edge_id, cause_node_id, effect_node_id, weight, count, confidence, alpha, beta, version_id, last_seen
270)
271SELECT
272 edge_id,
273 cause_node_id,
274 effect_node_id,
275 weight,
276 count,
277 confidence,
278 CASE
279 WHEN count > 0 THEN CAST(count AS REAL)
280 ELSE 1.0
281 END,
282 1.0,
283 version_id,
284 last_seen
285FROM cea_edges;
286DROP TABLE cea_edges;
287ALTER TABLE cea_edges_v2 RENAME TO cea_edges;
288"#,
289 )
290 .map_err(|error| cea_store::CeaStoreError::Backend(error.to_string()))?;
291 }
292
293 if has_blocked_run_log_schema(conn)? {
294 conn.execute_batch(
295 r#"
296CREATE TABLE cea_run_log_v2 (
297 run_hash TEXT PRIMARY KEY,
298 eval_id TEXT NOT NULL,
299 edges_added INTEGER NOT NULL,
300 edges_updated INTEGER NOT NULL,
301 processed_at TEXT NOT NULL
302);
303INSERT INTO cea_run_log_v2 (run_hash, eval_id, edges_added, edges_updated, processed_at)
304SELECT run_hash, eval_id, edges_added, edges_updated, processed_at
305FROM cea_run_log;
306DROP TABLE cea_run_log;
307ALTER TABLE cea_run_log_v2 RENAME TO cea_run_log;
308"#,
309 )
310 .map_err(|error| cea_store::CeaStoreError::Backend(error.to_string()))?;
311 }
312
313 Ok(())
314}
315
316fn positive_stats(weight_delta: f64) -> cea_core::EdgeStats {
317 let mut stats = cea_core::EdgeStats::default();
318 stats.observe_positive(weight_delta);
319 stats
320}
321
322fn stats_from_persisted(alpha: f64, beta: f64, count: i64) -> cea_core::EdgeStats {
323 cea_core::EdgeStats {
324 alpha,
325 beta,
326 observations: count.max(0) as u64,
327 }
328}
329
330fn map_edge_row(row: &rusqlite::Row<'_>) -> rusqlite::Result<cea_store::CeaEdgeRow> {
331 let count: i64 = row.get(4)?;
332 let alpha: f64 = row.get(6)?;
333 let beta: f64 = row.get(7)?;
334 let stats = stats_from_persisted(alpha, beta, count);
335 Ok(cea_store::CeaEdgeRow {
336 edge_id: row.get(0)?,
337 cause_node_id: row.get(1)?,
338 effect_node_id: row.get(2)?,
339 weight: row.get(3)?,
340 alpha,
341 beta,
342 count,
343 confidence: stats.confidence(),
344 version_id: row.get(8)?,
345 last_seen: row.get(9)?,
346 })
347}
348
349impl cea_store::CeaStoreWriteTx for SqliteCeaWriteTx<'_, '_> {
350 fn has_run(&self, run_hash: &str) -> Result<bool, cea_store::CeaStoreError> {
351 let count: i64 = self
352 .tx
353 .query_row(
354 "SELECT COUNT(*) FROM cea_run_log WHERE run_hash = ?1",
355 [run_hash],
356 |row| row.get(0),
357 )
358 .map_err(|error| cea_store::CeaStoreError::Backend(error.to_string()))?;
359 Ok(count > 0)
360 }
361
362 fn upsert_node(
363 &self,
364 node_id: &str,
365 node_kind: &str,
366 sig_json: &str,
367 ) -> Result<(), cea_store::CeaStoreError> {
368 let now = chrono::Utc::now().to_rfc3339();
369 self.tx
370 .execute(
371 "INSERT INTO cea_nodes (node_id, node_kind, sig_json, first_seen, last_seen) VALUES (?1, ?2, ?3, ?4, ?4) ON CONFLICT(node_id) DO UPDATE SET last_seen = ?4",
372 rusqlite::params![node_id, node_kind, sig_json, now],
373 )
374 .map_err(|error| cea_store::CeaStoreError::Backend(error.to_string()))?;
375 Ok(())
376 }
377
378 fn upsert_edge(
379 &self,
380 edge_id: &str,
381 cause_node_id: &str,
382 effect_node_id: &str,
383 weight_delta: f64,
384 version_id: &str,
385 ) -> Result<bool, cea_store::CeaStoreError> {
386 let now = chrono::Utc::now().to_rfc3339();
387 let weight_delta = weight_delta.max(0.0);
388 let existing = self
389 .tx
390 .query_row(
391 "SELECT weight, count, alpha, beta FROM cea_edges WHERE cause_node_id = ?1 AND effect_node_id = ?2 AND version_id = ?3",
392 rusqlite::params![cause_node_id, effect_node_id, version_id],
393 |row| {
394 Ok((
395 row.get::<_, f64>(0)?,
396 row.get::<_, i64>(1)?,
397 row.get::<_, f64>(2)?,
398 row.get::<_, f64>(3)?,
399 ))
400 },
401 )
402 .optional()
403 .map_err(|error| cea_store::CeaStoreError::Backend(error.to_string()))?;
404
405 if let Some((weight, count, alpha, beta)) = existing {
406 let mut stats = stats_from_persisted(alpha, beta, count);
407 stats.observe_positive(weight_delta);
408 self.tx
409 .execute(
410 "UPDATE cea_edges SET weight = ?1, count = ?2, confidence = ?3, alpha = ?4, beta = ?5, last_seen = ?6 WHERE cause_node_id = ?7 AND effect_node_id = ?8 AND version_id = ?9",
411 rusqlite::params![
412 weight + weight_delta,
413 stats.observations as i64,
414 stats.confidence(),
415 stats.alpha,
416 stats.beta,
417 now,
418 cause_node_id,
419 effect_node_id,
420 version_id
421 ],
422 )
423 .map_err(|error| cea_store::CeaStoreError::Backend(error.to_string()))?;
424 return Ok(false);
425 }
426
427 let stats = positive_stats(weight_delta);
428 self.tx
429 .execute(
430 "INSERT INTO cea_edges (edge_id, cause_node_id, effect_node_id, weight, count, confidence, alpha, beta, version_id, last_seen) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)",
431 rusqlite::params![
432 edge_id,
433 cause_node_id,
434 effect_node_id,
435 weight_delta,
436 stats.observations as i64,
437 stats.confidence(),
438 stats.alpha,
439 stats.beta,
440 version_id,
441 now
442 ],
443 )
444 .map_err(|error| cea_store::CeaStoreError::Backend(error.to_string()))?;
445
446 Ok(true)
447 }
448
449 fn load_effect_ids_for_cause(
450 &self,
451 cause_node_id: &str,
452 version_id: &str,
453 ) -> Result<Vec<String>, cea_store::CeaStoreError> {
454 let mut stmt = self
455 .tx
456 .prepare(
457 "SELECT effect_node_id FROM cea_edges WHERE cause_node_id = ?1 AND version_id = ?2 ORDER BY effect_node_id",
458 )
459 .map_err(|error| cea_store::CeaStoreError::Backend(error.to_string()))?;
460 let rows = stmt
461 .query_map(rusqlite::params![cause_node_id, version_id], |row| {
462 row.get(0)
463 })
464 .map_err(|error| cea_store::CeaStoreError::Backend(error.to_string()))?;
465
466 let mut effect_ids = Vec::new();
467 for row in rows {
468 effect_ids
469 .push(row.map_err(|error| cea_store::CeaStoreError::Backend(error.to_string()))?);
470 }
471 Ok(effect_ids)
472 }
473
474 fn reinforce_negative_edge(
475 &self,
476 cause_node_id: &str,
477 effect_node_id: &str,
478 amount: f64,
479 version_id: &str,
480 ) -> Result<(), cea_store::CeaStoreError> {
481 let now = chrono::Utc::now().to_rfc3339();
482 let amount = amount.max(0.0);
483 let existing = self
484 .tx
485 .query_row(
486 "SELECT count, alpha, beta FROM cea_edges WHERE cause_node_id = ?1 AND effect_node_id = ?2 AND version_id = ?3",
487 rusqlite::params![cause_node_id, effect_node_id, version_id],
488 |row| {
489 Ok((
490 row.get::<_, i64>(0)?,
491 row.get::<_, f64>(1)?,
492 row.get::<_, f64>(2)?,
493 ))
494 },
495 )
496 .optional()
497 .map_err(|error| cea_store::CeaStoreError::Backend(error.to_string()))?
498 .ok_or_else(|| {
499 cea_store::CeaStoreError::Backend(format!(
500 "missing edge for negative observation: {cause_node_id} -> {effect_node_id} ({version_id})"
501 ))
502 })?;
503
504 let (count, alpha, beta) = existing;
505 let mut stats = stats_from_persisted(alpha, beta, count);
506 stats.observe_negative(amount);
507 self.tx
508 .execute(
509 "UPDATE cea_edges SET count = ?1, confidence = ?2, alpha = ?3, beta = ?4, last_seen = ?5 WHERE cause_node_id = ?6 AND effect_node_id = ?7 AND version_id = ?8",
510 rusqlite::params![
511 stats.observations as i64,
512 stats.confidence(),
513 stats.alpha,
514 stats.beta,
515 now,
516 cause_node_id,
517 effect_node_id,
518 version_id
519 ],
520 )
521 .map_err(|error| cea_store::CeaStoreError::Backend(error.to_string()))?;
522 Ok(())
523 }
524
525 fn insert_run_log(
526 &self,
527 run_hash: &str,
528 eval_id: &str,
529 edges_added: i64,
530 edges_updated: i64,
531 ) -> Result<(), cea_store::CeaStoreError> {
532 let now = chrono::Utc::now().to_rfc3339();
533 self.tx
534 .execute(
535 "INSERT INTO cea_run_log (run_hash, eval_id, edges_added, edges_updated, processed_at) VALUES (?1, ?2, ?3, ?4, ?5)",
536 rusqlite::params![run_hash, eval_id, edges_added, edges_updated, now],
537 )
538 .map_err(|error| cea_store::CeaStoreError::Backend(error.to_string()))?;
539 Ok(())
540 }
541}
542
543impl<'conn> cea_store::CeaStore for SqliteCeaStoreConn<'conn> {
544 fn with_write_tx<T, F>(&self, f: F) -> Result<T, cea_store::CeaStoreError>
545 where
546 F: FnOnce(&dyn cea_store::CeaStoreWriteTx) -> Result<T, cea_store::CeaStoreError>,
547 {
548 let _ = f;
549 Err(cea_store::CeaStoreError::Backend(
550 "SqliteCeaStoreConn does not open write transactions; use SqliteCeaStore or SqliteCeaStoreTx"
551 .to_string(),
552 ))
553 }
554
555 fn load_nodes(&self) -> Result<Vec<cea_store::CeaNodeRow>, cea_store::CeaStoreError> {
556 let mut statement = self
557 .conn
558 .prepare("SELECT node_id, node_kind, sig_json FROM cea_nodes")
559 .map_err(|error| cea_store::CeaStoreError::Backend(error.to_string()))?;
560 let rows = statement
561 .query_map([], |row| {
562 Ok(cea_store::CeaNodeRow {
563 node_id: row.get(0)?,
564 node_kind: row.get(1)?,
565 sig_json: row.get(2)?,
566 })
567 })
568 .map_err(|error| cea_store::CeaStoreError::Backend(error.to_string()))?;
569
570 let mut nodes = Vec::new();
571 for row in rows {
572 nodes.push(row.map_err(|error| cea_store::CeaStoreError::Backend(error.to_string()))?);
573 }
574 Ok(nodes)
575 }
576
577 fn load_edges(
578 &self,
579 version_id: Option<&str>,
580 ) -> Result<Vec<cea_store::CeaEdgeRow>, cea_store::CeaStoreError> {
581 let query = if version_id.is_some() {
582 "SELECT edge_id, cause_node_id, effect_node_id, weight, count, confidence, alpha, beta, version_id, last_seen FROM cea_edges WHERE version_id = ?1"
583 } else {
584 "SELECT edge_id, cause_node_id, effect_node_id, weight, count, confidence, alpha, beta, version_id, last_seen FROM cea_edges"
585 };
586
587 let mut statement = self
588 .conn
589 .prepare(query)
590 .map_err(|error| cea_store::CeaStoreError::Backend(error.to_string()))?;
591
592 let rows = if let Some(version_id) = version_id {
593 statement
594 .query_map([version_id], map_edge_row)
595 .map_err(|error| cea_store::CeaStoreError::Backend(error.to_string()))?
596 } else {
597 statement
598 .query_map([], map_edge_row)
599 .map_err(|error| cea_store::CeaStoreError::Backend(error.to_string()))?
600 };
601
602 let mut edges = Vec::new();
603 for row in rows {
604 edges.push(row.map_err(|error| cea_store::CeaStoreError::Backend(error.to_string()))?);
605 }
606 Ok(edges)
607 }
608}
609
610impl<'tx, 'conn> cea_store::CeaStore for SqliteCeaStoreTx<'tx, 'conn> {
611 fn with_write_tx<T, F>(&self, f: F) -> Result<T, cea_store::CeaStoreError>
612 where
613 F: FnOnce(&dyn cea_store::CeaStoreWriteTx) -> Result<T, cea_store::CeaStoreError>,
614 {
615 let write_tx = SqliteCeaWriteTx { tx: self.tx };
616 f(&write_tx)
617 }
618
619 fn load_nodes(&self) -> Result<Vec<cea_store::CeaNodeRow>, cea_store::CeaStoreError> {
620 let mut statement = self
621 .tx
622 .prepare("SELECT node_id, node_kind, sig_json FROM cea_nodes")
623 .map_err(|error| cea_store::CeaStoreError::Backend(error.to_string()))?;
624 let rows = statement
625 .query_map([], |row| {
626 Ok(cea_store::CeaNodeRow {
627 node_id: row.get(0)?,
628 node_kind: row.get(1)?,
629 sig_json: row.get(2)?,
630 })
631 })
632 .map_err(|error| cea_store::CeaStoreError::Backend(error.to_string()))?;
633
634 let mut nodes = Vec::new();
635 for row in rows {
636 nodes.push(row.map_err(|error| cea_store::CeaStoreError::Backend(error.to_string()))?);
637 }
638 Ok(nodes)
639 }
640
641 fn load_edges(
642 &self,
643 version_id: Option<&str>,
644 ) -> Result<Vec<cea_store::CeaEdgeRow>, cea_store::CeaStoreError> {
645 let query = if version_id.is_some() {
646 "SELECT edge_id, cause_node_id, effect_node_id, weight, count, confidence, alpha, beta, version_id, last_seen FROM cea_edges WHERE version_id = ?1"
647 } else {
648 "SELECT edge_id, cause_node_id, effect_node_id, weight, count, confidence, alpha, beta, version_id, last_seen FROM cea_edges"
649 };
650
651 let mut statement = self
652 .tx
653 .prepare(query)
654 .map_err(|error| cea_store::CeaStoreError::Backend(error.to_string()))?;
655
656 let rows = if let Some(version_id) = version_id {
657 statement
658 .query_map([version_id], map_edge_row)
659 .map_err(|error| cea_store::CeaStoreError::Backend(error.to_string()))?
660 } else {
661 statement
662 .query_map([], map_edge_row)
663 .map_err(|error| cea_store::CeaStoreError::Backend(error.to_string()))?
664 };
665
666 let mut edges = Vec::new();
667 for row in rows {
668 edges.push(row.map_err(|error| cea_store::CeaStoreError::Backend(error.to_string()))?);
669 }
670 Ok(edges)
671 }
672}
673
674impl cea_store::CeaStore for SqliteCeaStore {
675 fn with_write_tx<T, F>(&self, f: F) -> Result<T, cea_store::CeaStoreError>
676 where
677 F: FnOnce(&dyn cea_store::CeaStoreWriteTx) -> Result<T, cea_store::CeaStoreError>,
678 {
679 let mut connection = self.connect()?;
680 let tx = connection
681 .transaction()
682 .map_err(|error| cea_store::CeaStoreError::Backend(error.to_string()))?;
683 let write_tx = SqliteCeaWriteTx { tx: &tx };
684 let result = f(&write_tx)?;
685 tx.commit()
686 .map_err(|error| cea_store::CeaStoreError::Backend(error.to_string()))?;
687 Ok(result)
688 }
689
690 fn load_nodes(&self) -> Result<Vec<cea_store::CeaNodeRow>, cea_store::CeaStoreError> {
691 let connection = self.connect()?;
692 let mut statement = connection
693 .prepare("SELECT node_id, node_kind, sig_json FROM cea_nodes")
694 .map_err(|error| cea_store::CeaStoreError::Backend(error.to_string()))?;
695 let rows = statement
696 .query_map([], |row| {
697 Ok(cea_store::CeaNodeRow {
698 node_id: row.get(0)?,
699 node_kind: row.get(1)?,
700 sig_json: row.get(2)?,
701 })
702 })
703 .map_err(|error| cea_store::CeaStoreError::Backend(error.to_string()))?;
704
705 let mut nodes = Vec::new();
706 for row in rows {
707 nodes.push(row.map_err(|error| cea_store::CeaStoreError::Backend(error.to_string()))?);
708 }
709 Ok(nodes)
710 }
711
712 fn load_edges(
713 &self,
714 version_id: Option<&str>,
715 ) -> Result<Vec<cea_store::CeaEdgeRow>, cea_store::CeaStoreError> {
716 let connection = self.connect()?;
717
718 let query = if version_id.is_some() {
719 "SELECT edge_id, cause_node_id, effect_node_id, weight, count, confidence, alpha, beta, version_id, last_seen FROM cea_edges WHERE version_id = ?1"
720 } else {
721 "SELECT edge_id, cause_node_id, effect_node_id, weight, count, confidence, alpha, beta, version_id, last_seen FROM cea_edges"
722 };
723
724 let mut statement = connection
725 .prepare(query)
726 .map_err(|error| cea_store::CeaStoreError::Backend(error.to_string()))?;
727
728 let rows = if let Some(version_id) = version_id {
729 statement
730 .query_map([version_id], map_edge_row)
731 .map_err(|error| cea_store::CeaStoreError::Backend(error.to_string()))?
732 } else {
733 statement
734 .query_map([], map_edge_row)
735 .map_err(|error| cea_store::CeaStoreError::Backend(error.to_string()))?
736 };
737
738 let mut edges = Vec::new();
739 for row in rows {
740 edges.push(row.map_err(|error| cea_store::CeaStoreError::Backend(error.to_string()))?);
741 }
742 Ok(edges)
743 }
744}
745
746#[cfg(test)]
747mod tests {
748 use std::path::Path;
749
750 use cea_core::{
751 AnchorKind, AttributedRunResult, AttributionTriple, EditOpKind, EditOpSignature, FileIndex,
752 OpIndex, ScopeTag,
753 };
754 use check_runner::{CheckResult, EffectSignature, ParsedCheckOutput};
755 use tempfile::TempDir;
756
757 use super::{SqliteCeaStore, CURRENT_SCHEMA_VERSION, SCHEMA_V2_SQL};
758
759 fn create_current_schema(path: &Path) {
760 let conn = rusqlite::Connection::open(path).unwrap();
761 conn.execute_batch(SCHEMA_V2_SQL).unwrap();
762 }
763
764 fn create_legacy_schema(path: &Path) {
765 let conn = rusqlite::Connection::open(path).unwrap();
766 conn.execute_batch(
767 r#"
768CREATE TABLE cea_nodes (
769 node_id TEXT PRIMARY KEY,
770 node_kind TEXT NOT NULL,
771 sig_json TEXT NOT NULL,
772 first_seen TEXT NOT NULL,
773 last_seen TEXT NOT NULL
774);
775CREATE TABLE cea_edges (
776 edge_id TEXT PRIMARY KEY,
777 cause_node_id TEXT NOT NULL,
778 effect_node_id TEXT NOT NULL,
779 weight REAL NOT NULL,
780 count INTEGER NOT NULL,
781 confidence REAL NOT NULL,
782 version_id TEXT NOT NULL,
783 last_seen TEXT NOT NULL,
784 UNIQUE(cause_node_id, effect_node_id, version_id)
785);
786CREATE TABLE cea_run_log (
787 run_hash TEXT PRIMARY KEY,
788 eval_id TEXT NOT NULL CHECK(eval_id != 'blocked'),
789 edges_added INTEGER NOT NULL,
790 edges_updated INTEGER NOT NULL,
791 processed_at TEXT NOT NULL
792);
793"#,
794 )
795 .unwrap();
796 }
797
798 fn install_run_log_failure_trigger(path: &Path) {
799 let conn = rusqlite::Connection::open(path).unwrap();
800 conn.execute_batch(
801 r#"
802CREATE TRIGGER fail_blocked_run_log_insert
803BEFORE INSERT ON cea_run_log
804WHEN NEW.eval_id = 'blocked'
805BEGIN
806 SELECT RAISE(ABORT, 'blocked run log insert');
807END;
808"#,
809 )
810 .unwrap();
811 }
812
813 #[test]
814 fn open_creates_schema_for_fresh_database() {
815 let dir = TempDir::new().unwrap();
816 let path = dir.path().join("cea.db");
817
818 let store = SqliteCeaStore::open(&path).unwrap();
819 let conn = store.connect().unwrap();
820 let user_version: i64 = conn
821 .pragma_query_value(None, "user_version", |row| row.get(0))
822 .unwrap();
823 let nodes: i64 = conn
824 .query_row("SELECT COUNT(*) FROM cea_nodes", [], |row| row.get(0))
825 .unwrap();
826 let edges: i64 = conn
827 .query_row("SELECT COUNT(*) FROM cea_edges", [], |row| row.get(0))
828 .unwrap();
829 let runs: i64 = conn
830 .query_row("SELECT COUNT(*) FROM cea_run_log", [], |row| row.get(0))
831 .unwrap();
832
833 assert_eq!(user_version, CURRENT_SCHEMA_VERSION);
834 assert_eq!((nodes, edges, runs), (0, 0, 0));
835 assert!(super::column_exists(&conn, "cea_edges", "alpha").unwrap());
836 assert!(super::column_exists(&conn, "cea_edges", "beta").unwrap());
837 }
838
839 #[test]
840 fn open_adopts_existing_current_schema_without_manual_versioning() {
841 let dir = TempDir::new().unwrap();
842 let path = dir.path().join("cea.db");
843 create_current_schema(&path);
844
845 let conn = rusqlite::Connection::open(&path).unwrap();
846 let user_version: i64 = conn
847 .pragma_query_value(None, "user_version", |row| row.get(0))
848 .unwrap();
849 assert_eq!(user_version, 0);
850 drop(conn);
851
852 let store = SqliteCeaStore::open(&path).unwrap();
853 let conn = store.connect().unwrap();
854 let user_version: i64 = conn
855 .pragma_query_value(None, "user_version", |row| row.get(0))
856 .unwrap();
857 assert_eq!(user_version, CURRENT_SCHEMA_VERSION);
858 }
859
860 #[test]
861 fn open_migrates_legacy_schema_and_preserves_rows() {
862 let dir = TempDir::new().unwrap();
863 let path = dir.path().join("cea.db");
864 create_legacy_schema(&path);
865
866 let legacy_conn = rusqlite::Connection::open(&path).unwrap();
867 legacy_conn
868 .execute(
869 "INSERT INTO cea_edges (edge_id, cause_node_id, effect_node_id, weight, count, confidence, version_id, last_seen) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)",
870 rusqlite::params![
871 "legacy-edge",
872 "cause",
873 "effect",
874 3.5_f64,
875 4_i64,
876 0.6_f64,
877 "v1",
878 "now"
879 ],
880 )
881 .unwrap();
882 drop(legacy_conn);
883
884 let store = SqliteCeaStore::open(&path).unwrap();
885 let conn = store.connect().unwrap();
886 let user_version: i64 = conn
887 .pragma_query_value(None, "user_version", |row| row.get(0))
888 .unwrap();
889 let (weight, count, alpha, beta): (f64, i64, f64, f64) = conn
890 .query_row(
891 "SELECT weight, count, alpha, beta FROM cea_edges WHERE edge_id = 'legacy-edge'",
892 [],
893 |row| Ok((row.get(0)?, row.get(1)?, row.get(2)?, row.get(3)?)),
894 )
895 .unwrap();
896 let run_log_sql: String = conn
897 .query_row(
898 "SELECT sql FROM sqlite_master WHERE type = 'table' AND name = 'cea_run_log'",
899 [],
900 |row| row.get(0),
901 )
902 .unwrap();
903
904 assert_eq!(user_version, CURRENT_SCHEMA_VERSION);
905 assert!(super::column_exists(&conn, "cea_edges", "alpha").unwrap());
906 assert!(super::column_exists(&conn, "cea_edges", "beta").unwrap());
907 assert_eq!((weight, count, alpha, beta), (3.5, 4, 4.0, 1.0));
908 assert!(
909 !run_log_sql.contains("CHECK(eval_id != 'blocked')"),
910 "legacy run-log constraint should be removed during migration"
911 );
912 drop(conn);
913
914 let outcome =
915 cea_store::update_graph(&store, &sample_result("blocked"), "blocked", "v1", 0.95)
916 .unwrap();
917 assert_eq!(
918 outcome,
919 cea_store::UpdateResult::Applied {
920 edges_added: 1,
921 edges_updated: 0,
922 }
923 );
924 }
925
926 fn sample_result(seed: &str) -> AttributedRunResult {
927 sample_result_with(seed, seed)
928 }
929
930 fn sample_result_with(cause_seed: &str, effect_seed: &str) -> AttributedRunResult {
931 AttributedRunResult::new(
932 vec![AttributionTriple {
933 cause: EditOpSignature {
934 op_kind: EditOpKind::Insert,
935 anchor_kind: AnchorKind::AfterLine,
936 lines_added: 1,
937 lines_removed: 0,
938 context_hash: format!("{cause_seed:0>8}"),
939 file_extension: "rs".to_string(),
940 scope_tag: ScopeTag::Function,
941 op_index: OpIndex(0),
942 file_index: FileIndex(0),
943 },
944 effect: EffectSignature {
945 check_kind: "clippy".to_string(),
946 outcome: "warning".to_string(),
947 severity: "warning".to_string(),
948 message_class: format!("unused_{effect_seed}"),
949 line_offset_from_edit: Some(1),
950 },
951 distance: 1,
952 weight: 1.0,
953 }],
954 CheckResult {
955 fmt_pass: true,
956 clippy_pass: false,
957 test_pass: true,
958 fmt_output: ParsedCheckOutput::default(),
959 clippy_output: ParsedCheckOutput::default(),
960 test_output: ParsedCheckOutput::default(),
961 total_duration_ms: 1,
962 },
963 )
964 }
965
966 #[test]
967 fn update_graph_persists_nodes_edges_and_run_log() {
968 let dir = TempDir::new().unwrap();
969 let path = dir.path().join("cea.db");
970 create_current_schema(&path);
971 let store = SqliteCeaStore::open(&path).unwrap();
972 let result = sample_result("ok");
973
974 let outcome = cea_store::update_graph(&store, &result, "eval-ok", "v1", 0.95).unwrap();
975 assert_eq!(
976 outcome,
977 cea_store::UpdateResult::Applied {
978 edges_added: 1,
979 edges_updated: 0,
980 }
981 );
982
983 let conn = rusqlite::Connection::open(&path).unwrap();
984 let nodes: i64 = conn
985 .query_row("SELECT COUNT(*) FROM cea_nodes", [], |row| row.get(0))
986 .unwrap();
987 let edges: i64 = conn
988 .query_row("SELECT COUNT(*) FROM cea_edges", [], |row| row.get(0))
989 .unwrap();
990 let runs: i64 = conn
991 .query_row("SELECT COUNT(*) FROM cea_run_log", [], |row| row.get(0))
992 .unwrap();
993 assert_eq!((nodes, edges, runs), (2, 1, 1));
994 }
995
996 #[test]
997 fn update_graph_rolls_back_when_run_log_insert_fails() {
998 let dir = TempDir::new().unwrap();
999 let path = dir.path().join("cea.db");
1000 create_current_schema(&path);
1001 let store = SqliteCeaStore::open(&path).unwrap();
1002 install_run_log_failure_trigger(&path);
1003 let result = sample_result("rollback");
1004
1005 let err = cea_store::update_graph(&store, &result, "blocked", "v1", 0.95).unwrap_err();
1006 assert!(matches!(err, cea_store::CeaStoreError::Backend(_)));
1007
1008 let conn = rusqlite::Connection::open(&path).unwrap();
1009 let nodes: i64 = conn
1010 .query_row("SELECT COUNT(*) FROM cea_nodes", [], |row| row.get(0))
1011 .unwrap();
1012 let edges: i64 = conn
1013 .query_row("SELECT COUNT(*) FROM cea_edges", [], |row| row.get(0))
1014 .unwrap();
1015 let runs: i64 = conn
1016 .query_row("SELECT COUNT(*) FROM cea_run_log", [], |row| row.get(0))
1017 .unwrap();
1018 assert_eq!((nodes, edges, runs), (0, 0, 0));
1019 }
1020
1021 #[test]
1022 fn load_edges_filters_by_version() {
1023 let dir = TempDir::new().unwrap();
1024 let path = dir.path().join("cea.db");
1025 create_current_schema(&path);
1026 let store = SqliteCeaStore::open(&path).unwrap();
1027
1028 cea_store::update_graph(&store, &sample_result("one"), "eval-1", "v1", 0.95).unwrap();
1029 cea_store::update_graph(&store, &sample_result("two"), "eval-2", "v2", 0.95).unwrap();
1030
1031 let all = cea_store::CeaStore::load_edges(&store, None).unwrap();
1032 let filtered = cea_store::CeaStore::load_edges(&store, Some("v1")).unwrap();
1033 assert_eq!(all.len(), 2);
1034 assert_eq!(filtered.len(), 1);
1035 assert_eq!(filtered[0].version_id, "v1");
1036 }
1037
1038 #[test]
1039 fn update_graph_persists_negative_evidence_for_absent_effects() {
1040 let dir = TempDir::new().unwrap();
1041 let path = dir.path().join("cea.db");
1042 create_current_schema(&path);
1043 let store = SqliteCeaStore::open(&path).unwrap();
1044
1045 cea_store::update_graph(
1046 &store,
1047 &sample_result_with("shared", "one"),
1048 "eval-1",
1049 "v1",
1050 0.95,
1051 )
1052 .unwrap();
1053 let second = cea_store::update_graph(
1054 &store,
1055 &sample_result_with("shared", "two"),
1056 "eval-2",
1057 "v1",
1058 0.95,
1059 )
1060 .unwrap();
1061 assert_eq!(
1062 second,
1063 cea_store::UpdateResult::Applied {
1064 edges_added: 1,
1065 edges_updated: 1,
1066 }
1067 );
1068
1069 let edges = cea_store::CeaStore::load_edges(&store, Some("v1")).unwrap();
1070 let cause_id = cea_core::edit_op_node_id(&EditOpSignature {
1071 op_kind: EditOpKind::Insert,
1072 anchor_kind: AnchorKind::AfterLine,
1073 lines_added: 1,
1074 lines_removed: 0,
1075 context_hash: format!("{:0>8}", "shared"),
1076 file_extension: "rs".to_string(),
1077 scope_tag: ScopeTag::Function,
1078 op_index: OpIndex(0),
1079 file_index: FileIndex(0),
1080 });
1081 let first_effect_id = cea_core::effect_node_id(&EffectSignature {
1082 check_kind: "clippy".to_string(),
1083 outcome: "warning".to_string(),
1084 severity: "warning".to_string(),
1085 message_class: "unused_one".to_string(),
1086 line_offset_from_edit: Some(1),
1087 });
1088 let second_effect_id = cea_core::effect_node_id(&EffectSignature {
1089 check_kind: "clippy".to_string(),
1090 outcome: "warning".to_string(),
1091 severity: "warning".to_string(),
1092 message_class: "unused_two".to_string(),
1093 line_offset_from_edit: Some(1),
1094 });
1095 let first_edge = edges
1096 .iter()
1097 .find(|edge| edge.cause_node_id == cause_id && edge.effect_node_id == first_effect_id)
1098 .unwrap();
1099 let second_edge = edges
1100 .iter()
1101 .find(|edge| edge.cause_node_id == cause_id && edge.effect_node_id == second_effect_id)
1102 .unwrap();
1103
1104 assert_eq!(first_edge.beta, 2.0);
1105 assert_eq!(first_edge.count, 2);
1106 assert_eq!(second_edge.beta, 1.0);
1107 assert_eq!(second_edge.count, 1);
1108 }
1109
1110 #[test]
1111 fn load_graph_round_trips_persisted_edge_stats() {
1112 let dir = TempDir::new().unwrap();
1113 let path = dir.path().join("cea.db");
1114 create_current_schema(&path);
1115
1116 let cause = cea_core::EditOpSignature {
1117 op_kind: EditOpKind::Insert,
1118 anchor_kind: AnchorKind::AfterLine,
1119 lines_added: 1,
1120 lines_removed: 0,
1121 context_hash: "0000stats".to_string(),
1122 file_extension: "rs".to_string(),
1123 scope_tag: ScopeTag::Function,
1124 op_index: OpIndex(0),
1125 file_index: FileIndex(0),
1126 };
1127 let effect = EffectSignature {
1128 check_kind: "clippy".to_string(),
1129 outcome: "warning".to_string(),
1130 severity: "warning".to_string(),
1131 message_class: "unused_stats".to_string(),
1132 line_offset_from_edit: Some(1),
1133 };
1134 let cause_id = cea_core::edit_op_node_id(&cause);
1135 let effect_id = cea_core::effect_node_id(&effect);
1136 let stats = cea_core::EdgeStats {
1137 alpha: 4.0,
1138 beta: 2.0,
1139 observations: 3,
1140 };
1141
1142 let conn = rusqlite::Connection::open(&path).unwrap();
1143 conn.execute(
1144 "INSERT INTO cea_nodes (node_id, node_kind, sig_json, first_seen, last_seen) VALUES (?1, 'cause', ?2, 'now', 'now')",
1145 rusqlite::params![cause_id, serde_json::to_string(&cause).unwrap()],
1146 )
1147 .unwrap();
1148 conn.execute(
1149 "INSERT INTO cea_nodes (node_id, node_kind, sig_json, first_seen, last_seen) VALUES (?1, 'effect', ?2, 'now', 'now')",
1150 rusqlite::params![effect_id, serde_json::to_string(&effect).unwrap()],
1151 )
1152 .unwrap();
1153 conn.execute(
1154 "INSERT INTO cea_edges (edge_id, cause_node_id, effect_node_id, weight, count, confidence, alpha, beta, version_id, last_seen) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, 'v1', 'now')",
1155 rusqlite::params![
1156 format!("{cause_id}_{effect_id}_v1"),
1157 cause_id,
1158 effect_id,
1159 3.0_f64,
1160 stats.observations as i64,
1161 0.0_f64,
1162 stats.alpha,
1163 stats.beta,
1164 ],
1165 )
1166 .unwrap();
1167
1168 let store = SqliteCeaStore::open(&path).unwrap();
1169 let graph = cea_store::load_graph(&store, Some("v1")).unwrap();
1170 let cause_index = graph.node_index_map[&cea_core::edit_op_node_id(&cause)];
1171 let effect_index = graph.node_index_map[&cea_core::effect_node_id(&effect)];
1172 let edge = graph
1173 .graph
1174 .edge_weight(graph.graph.find_edge(cause_index, effect_index).unwrap())
1175 .unwrap();
1176 assert_eq!(edge.stats, stats);
1177 assert!((edge.confidence - stats.confidence()).abs() < 1e-12);
1178 }
1179
1180 #[test]
1181 fn conn_write_transactions_are_explicitly_unsupported() {
1182 let dir = TempDir::new().unwrap();
1183 let path = dir.path().join("cea.db");
1184 create_current_schema(&path);
1185
1186 let conn = rusqlite::Connection::open(&path).unwrap();
1187 let store = super::SqliteCeaStoreConn::new(&conn);
1188 let err = cea_store::CeaStore::with_write_tx(&store, |_| Ok(())).unwrap_err();
1189 assert!(
1190 err.to_string().contains("does not open write transactions"),
1191 "unexpected error: {err}"
1192 );
1193 }
1194
1195 #[test]
1196 fn direct_store_connections_apply_safety_pragmas() {
1197 let dir = TempDir::new().unwrap();
1198 let path = dir.path().join("cea.db");
1199 create_current_schema(&path);
1200
1201 let store = SqliteCeaStore::open(&path).unwrap();
1202 let conn = store.connect().unwrap();
1203
1204 let foreign_keys: i64 = conn
1205 .pragma_query_value(None, "foreign_keys", |row| row.get(0))
1206 .unwrap();
1207 let journal_mode: String = conn
1208 .pragma_query_value(None, "journal_mode", |row| row.get(0))
1209 .unwrap();
1210 let busy_timeout: i64 = conn
1211 .pragma_query_value(None, "busy_timeout", |row| row.get(0))
1212 .unwrap();
1213
1214 assert_eq!(foreign_keys, 1, "foreign keys should be enabled");
1215 assert_eq!(journal_mode.to_lowercase(), "wal");
1216 assert_eq!(busy_timeout, 5_000);
1217 }
1218}