1use sqlx::{Row, SqlitePool};
4
5pub use nyx_agent_types::chain::ChainRecord;
6
7use super::attack_graph::AttackGraphStore;
8use crate::store::StoreError;
9
10pub struct ChainStore<'a> {
11 pool: &'a SqlitePool,
12}
13
14impl<'a> ChainStore<'a> {
15 pub fn new(pool: &'a SqlitePool) -> Self {
16 Self { pool }
17 }
18
19 pub async fn insert(&self, c: &ChainRecord) -> Result<(), StoreError> {
20 let cross_repo = if c.cross_repo { 1_i64 } else { 0_i64 };
21 sqlx::query(
22 r#"
23 INSERT INTO chains (
24 id, run_id, cross_repo, member_ids, rationale_blob,
25 attack_provenance, prompt_version, status, verification_attempt_id,
26 evidence_blob, severity
27 ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
28 "#,
29 )
30 .bind(&c.id)
31 .bind(&c.run_id)
32 .bind(cross_repo)
33 .bind(&c.member_ids)
34 .bind(&c.rationale_blob)
35 .bind(&c.attack_provenance)
36 .bind(&c.prompt_version)
37 .bind(&c.status)
38 .bind(&c.verification_attempt_id)
39 .bind(&c.evidence_blob)
40 .bind(&c.severity)
41 .execute(self.pool)
42 .await?;
43 AttackGraphStore::new(self.pool).record_chain(c).await?;
44 Ok(())
45 }
46
47 pub async fn get(&self, id: &str) -> Result<Option<ChainRecord>, StoreError> {
48 let row = sqlx::query(
49 r#"
50 SELECT id, run_id, cross_repo, member_ids,
51 rationale_blob, attack_provenance, prompt_version,
52 status, verification_attempt_id, evidence_blob, severity
53 FROM chains WHERE id = ?
54 "#,
55 )
56 .bind(id)
57 .fetch_optional(self.pool)
58 .await?;
59 row.map(row_to_chain_record).transpose()
60 }
61
62 pub async fn delete(&self, id: &str) -> Result<u64, StoreError> {
63 let res =
64 sqlx::query("DELETE FROM chains WHERE id = ?").bind(id).execute(self.pool).await?;
65 Ok(res.rows_affected())
66 }
67
68 pub async fn list_by_run(&self, run_id: &str) -> Result<Vec<ChainRecord>, StoreError> {
69 let rows = sqlx::query(
70 r#"
71 SELECT id, run_id, cross_repo, member_ids,
72 rationale_blob, attack_provenance, prompt_version,
73 status, verification_attempt_id, evidence_blob, severity
74 FROM chains WHERE run_id = ?
75 "#,
76 )
77 .bind(run_id)
78 .fetch_all(self.pool)
79 .await?;
80 rows.into_iter().map(row_to_chain_record).collect()
81 }
82
83 pub async fn list_by_run_and_status(
84 &self,
85 run_id: &str,
86 status: &str,
87 ) -> Result<Vec<ChainRecord>, StoreError> {
88 let rows = sqlx::query(
89 r#"
90 SELECT id, run_id, cross_repo, member_ids,
91 rationale_blob, attack_provenance, prompt_version,
92 status, verification_attempt_id, evidence_blob, severity
93 FROM chains WHERE run_id = ? AND status = ?
94 ORDER BY id
95 "#,
96 )
97 .bind(run_id)
98 .bind(status)
99 .fetch_all(self.pool)
100 .await?;
101 rows.into_iter().map(row_to_chain_record).collect()
102 }
103
104 pub async fn update_verification_state(
105 &self,
106 id: &str,
107 status: &str,
108 verification_attempt_id: Option<&str>,
109 evidence_blob: Option<&str>,
110 severity: Option<&str>,
111 ) -> Result<Option<ChainRecord>, StoreError> {
112 sqlx::query(
113 r#"
114 UPDATE chains SET
115 status = ?,
116 verification_attempt_id = COALESCE(?, verification_attempt_id),
117 evidence_blob = COALESCE(?, evidence_blob),
118 severity = COALESCE(?, severity)
119 WHERE id = ?
120 "#,
121 )
122 .bind(status)
123 .bind(verification_attempt_id)
124 .bind(evidence_blob)
125 .bind(severity)
126 .bind(id)
127 .execute(self.pool)
128 .await?;
129 self.get(id).await
130 }
131}
132
133fn row_to_chain_record(row: sqlx::sqlite::SqliteRow) -> Result<ChainRecord, StoreError> {
134 Ok(ChainRecord {
135 id: row.try_get("id")?,
136 run_id: row.try_get("run_id")?,
137 cross_repo: row.try_get::<i64, _>("cross_repo")? != 0,
138 member_ids: row.try_get("member_ids")?,
139 rationale_blob: row.try_get("rationale_blob")?,
140 attack_provenance: row.try_get("attack_provenance")?,
141 prompt_version: row.try_get("prompt_version")?,
142 status: row.try_get("status")?,
143 verification_attempt_id: row.try_get("verification_attempt_id")?,
144 evidence_blob: row.try_get("evidence_blob")?,
145 severity: row.try_get("severity")?,
146 })
147}
148
149#[cfg(test)]
150mod tests {
151 use crate::store::testutil::{fresh_store, sample_chain, sample_run};
152
153 #[tokio::test]
154 async fn insert_then_get_roundtrips() {
155 let (_tmp, s) = fresh_store().await;
156 s.runs().insert(&sample_run("r")).await.expect("run");
157 let c = sample_chain("c-1", "r", &["f-a", "f-b"]);
158 s.chains().insert(&c).await.expect("insert");
159 let got = s.chains().get("c-1").await.expect("get").expect("row");
160 assert_eq!(got, c);
161 assert!(!got.cross_repo);
162 }
163
164 #[tokio::test]
165 async fn list_by_run_returns_matching_only() {
166 let (_tmp, s) = fresh_store().await;
167 s.runs().insert(&sample_run("r1")).await.expect("r1");
168 s.runs().insert(&sample_run("r2")).await.expect("r2");
169 s.chains().insert(&sample_chain("c1", "r1", &["x"])).await.expect("c1");
170 s.chains().insert(&sample_chain("c2", "r2", &["y"])).await.expect("c2");
171 let got = s.chains().list_by_run("r1").await.expect("list");
172 assert_eq!(got.len(), 1);
173 assert_eq!(got[0].id, "c1");
174 }
175
176 #[tokio::test]
177 async fn cascade_from_run_delete() {
178 let (_tmp, s) = fresh_store().await;
179 s.runs().insert(&sample_run("doomed")).await.expect("run");
180 s.chains().insert(&sample_chain("c", "doomed", &["a"])).await.expect("chain");
181 s.runs().delete("doomed").await.expect("del");
182 assert!(
183 s.chains().get("c").await.expect("get").is_none(),
184 "FK cascade should have removed the chain"
185 );
186 }
187}