Skip to main content

chio_store_sqlite/
batch_approval_store.rs

1//! Phase 3.5 SQLite-backed `BatchApprovalStore`.
2//!
3//! Mirrors the single-request approval store: a WAL-journaled SQLite
4//! database with idempotent migrations. Batch records carry their own
5//! usage counters so the kernel can reconcile consumption without
6//! going back to the receipt log.
7
8use std::fs;
9use std::path::Path;
10
11use chio_core::capability::MonetaryAmount;
12use chio_kernel::{ApprovalStoreError, BatchApproval, BatchApprovalStore};
13use r2d2::Pool;
14use r2d2_sqlite::SqliteConnectionManager;
15use rusqlite::{params, OptionalExtension};
16
17pub struct SqliteBatchApprovalStore {
18    pool: Pool<SqliteConnectionManager>,
19}
20
21impl SqliteBatchApprovalStore {
22    pub fn open(path: impl AsRef<Path>) -> Result<Self, ApprovalStoreError> {
23        let path = path.as_ref();
24        if let Some(parent) = path.parent() {
25            if !parent.as_os_str().is_empty() {
26                fs::create_dir_all(parent)
27                    .map_err(|e| ApprovalStoreError::Backend(format!("create dir: {e}")))?;
28            }
29        }
30        let manager = SqliteConnectionManager::file(path);
31        let pool = Pool::builder()
32            .max_size(4)
33            .build(manager)
34            .map_err(|e| ApprovalStoreError::Backend(format!("pool build: {e}")))?;
35        let store = Self { pool };
36        store.run_migrations()?;
37        Ok(store)
38    }
39
40    pub fn open_in_memory() -> Result<Self, ApprovalStoreError> {
41        let manager = SqliteConnectionManager::memory();
42        let pool = Pool::builder()
43            .max_size(1)
44            .build(manager)
45            .map_err(|e| ApprovalStoreError::Backend(format!("pool build: {e}")))?;
46        let store = Self { pool };
47        store.run_migrations()?;
48        Ok(store)
49    }
50
51    fn run_migrations(&self) -> Result<(), ApprovalStoreError> {
52        let conn = self
53            .pool
54            .get()
55            .map_err(|e| ApprovalStoreError::Backend(format!("pool get: {e}")))?;
56        conn.execute_batch(
57            r#"
58            PRAGMA journal_mode = WAL;
59            PRAGMA synchronous = FULL;
60            PRAGMA busy_timeout = 5000;
61
62            CREATE TABLE IF NOT EXISTS chio_hitl_batches (
63                batch_id TEXT PRIMARY KEY,
64                approver_hex TEXT NOT NULL,
65                subject_id TEXT NOT NULL,
66                server_pattern TEXT NOT NULL,
67                tool_pattern TEXT NOT NULL,
68                per_call_currency TEXT,
69                per_call_units INTEGER,
70                total_currency TEXT,
71                total_units INTEGER,
72                max_calls INTEGER,
73                not_before INTEGER NOT NULL,
74                not_after INTEGER NOT NULL,
75                used_calls INTEGER NOT NULL DEFAULT 0,
76                used_total_units INTEGER NOT NULL DEFAULT 0,
77                revoked INTEGER NOT NULL DEFAULT 0
78            );
79            CREATE INDEX IF NOT EXISTS idx_chio_hitl_batches_subject
80                ON chio_hitl_batches(subject_id, revoked);
81            "#,
82        )
83        .map_err(|e| ApprovalStoreError::Backend(format!("migration: {e}")))?;
84        Ok(())
85    }
86}
87
88fn row_to_batch(row: &rusqlite::Row<'_>) -> rusqlite::Result<BatchApproval> {
89    let per_call_currency: Option<String> = row.get(5)?;
90    let per_call_units: Option<i64> = row.get(6)?;
91    let total_currency: Option<String> = row.get(7)?;
92    let total_units: Option<i64> = row.get(8)?;
93    let max_calls: Option<i64> = row.get(9)?;
94    let revoked: i64 = row.get(14)?;
95    Ok(BatchApproval {
96        batch_id: row.get(0)?,
97        approver_hex: row.get(1)?,
98        subject_id: row.get(2)?,
99        server_pattern: row.get(3)?,
100        tool_pattern: row.get(4)?,
101        max_amount_per_call: match (per_call_currency, per_call_units) {
102            (Some(currency), Some(units)) => Some(MonetaryAmount {
103                currency,
104                units: units.max(0) as u64,
105            }),
106            _ => None,
107        },
108        max_total_amount: match (total_currency, total_units) {
109            (Some(currency), Some(units)) => Some(MonetaryAmount {
110                currency,
111                units: units.max(0) as u64,
112            }),
113            _ => None,
114        },
115        max_calls: max_calls.map(|v| v.max(0) as u32),
116        not_before: row.get::<_, i64>(10)?.max(0) as u64,
117        not_after: row.get::<_, i64>(11)?.max(0) as u64,
118        used_calls: row.get::<_, i64>(12)?.max(0) as u32,
119        used_total_units: row.get::<_, i64>(13)?.max(0) as u64,
120        revoked: revoked != 0,
121    })
122}
123
124fn pattern_matches(pattern: &str, value: &str) -> bool {
125    if pattern == "*" {
126        return true;
127    }
128    if let Some(prefix) = pattern.strip_suffix('*') {
129        return value.starts_with(prefix);
130    }
131    pattern == value
132}
133
134fn amount_fits(batch: &BatchApproval, amount: Option<&MonetaryAmount>) -> bool {
135    let Some(amt) = amount else {
136        return batch.max_amount_per_call.is_none() && batch.max_total_amount.is_none();
137    };
138    if let Some(per_call) = &batch.max_amount_per_call {
139        if amt.currency != per_call.currency || amt.units > per_call.units {
140            return false;
141        }
142    }
143    if let Some(total) = &batch.max_total_amount {
144        if amt.currency != total.currency
145            || batch.used_total_units.saturating_add(amt.units) > total.units
146        {
147            return false;
148        }
149    }
150    true
151}
152
153impl BatchApprovalStore for SqliteBatchApprovalStore {
154    fn store(&self, batch: &BatchApproval) -> Result<(), ApprovalStoreError> {
155        let conn = self
156            .pool
157            .get()
158            .map_err(|e| ApprovalStoreError::Backend(format!("pool get: {e}")))?;
159        conn.execute(
160            r#"INSERT OR REPLACE INTO chio_hitl_batches (
161                batch_id, approver_hex, subject_id, server_pattern, tool_pattern,
162                per_call_currency, per_call_units, total_currency, total_units,
163                max_calls, not_before, not_after, used_calls, used_total_units, revoked
164            ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15)"#,
165            params![
166                batch.batch_id,
167                batch.approver_hex,
168                batch.subject_id,
169                batch.server_pattern,
170                batch.tool_pattern,
171                batch
172                    .max_amount_per_call
173                    .as_ref()
174                    .map(|a| a.currency.clone()),
175                batch.max_amount_per_call.as_ref().map(|a| a.units as i64),
176                batch.max_total_amount.as_ref().map(|a| a.currency.clone()),
177                batch.max_total_amount.as_ref().map(|a| a.units as i64),
178                batch.max_calls.map(|c| c as i64),
179                batch.not_before as i64,
180                batch.not_after as i64,
181                batch.used_calls as i64,
182                batch.used_total_units as i64,
183                batch.revoked as i64,
184            ],
185        )
186        .map_err(|e| ApprovalStoreError::Backend(format!("insert batch: {e}")))?;
187        Ok(())
188    }
189
190    fn find_matching(
191        &self,
192        subject_id: &str,
193        server_id: &str,
194        tool_name: &str,
195        amount: Option<&MonetaryAmount>,
196        now: u64,
197    ) -> Result<Option<BatchApproval>, ApprovalStoreError> {
198        let conn = self
199            .pool
200            .get()
201            .map_err(|e| ApprovalStoreError::Backend(format!("pool get: {e}")))?;
202        // Pull every non-revoked candidate for this subject within the
203        // time window, then apply pattern + amount filters in Rust.
204        let mut stmt = conn
205            .prepare(
206                r#"SELECT batch_id, approver_hex, subject_id, server_pattern, tool_pattern,
207                          per_call_currency, per_call_units, total_currency, total_units,
208                          max_calls, not_before, not_after, used_calls, used_total_units, revoked
209                   FROM chio_hitl_batches
210                   WHERE subject_id = ?1 AND revoked = 0
211                     AND not_before <= ?2 AND not_after > ?2"#,
212            )
213            .map_err(|e| ApprovalStoreError::Backend(format!("prepare: {e}")))?;
214        let rows = stmt
215            .query_map(params![subject_id, now as i64], row_to_batch)
216            .map_err(|e| ApprovalStoreError::Backend(format!("query: {e}")))?;
217        for row in rows {
218            let batch = row.map_err(|e| ApprovalStoreError::Backend(format!("row: {e}")))?;
219            if !pattern_matches(&batch.server_pattern, server_id) {
220                continue;
221            }
222            if !pattern_matches(&batch.tool_pattern, tool_name) {
223                continue;
224            }
225            if let Some(max) = batch.max_calls {
226                if batch.used_calls >= max {
227                    continue;
228                }
229            }
230            if !amount_fits(&batch, amount) {
231                continue;
232            }
233            return Ok(Some(batch));
234        }
235        Ok(None)
236    }
237
238    fn record_usage(
239        &self,
240        batch_id: &str,
241        amount: Option<&MonetaryAmount>,
242    ) -> Result<(), ApprovalStoreError> {
243        let conn = self
244            .pool
245            .get()
246            .map_err(|e| ApprovalStoreError::Backend(format!("pool get: {e}")))?;
247        let added_units = amount.map(|a| a.units as i64).unwrap_or(0);
248        let rows = conn.execute(
249            "UPDATE chio_hitl_batches SET used_calls = used_calls + 1, used_total_units = used_total_units + ?2 WHERE batch_id = ?1",
250            params![batch_id, added_units],
251        )
252        .map_err(|e| ApprovalStoreError::Backend(format!("update usage: {e}")))?;
253        if rows == 0 {
254            return Err(ApprovalStoreError::NotFound(batch_id.to_string()));
255        }
256        Ok(())
257    }
258
259    fn revoke(&self, batch_id: &str) -> Result<(), ApprovalStoreError> {
260        let conn = self
261            .pool
262            .get()
263            .map_err(|e| ApprovalStoreError::Backend(format!("pool get: {e}")))?;
264        let rows = conn
265            .execute(
266                "UPDATE chio_hitl_batches SET revoked = 1 WHERE batch_id = ?1",
267                params![batch_id],
268            )
269            .map_err(|e| ApprovalStoreError::Backend(format!("revoke: {e}")))?;
270        if rows == 0 {
271            return Err(ApprovalStoreError::NotFound(batch_id.to_string()));
272        }
273        Ok(())
274    }
275
276    fn get(&self, batch_id: &str) -> Result<Option<BatchApproval>, ApprovalStoreError> {
277        let conn = self
278            .pool
279            .get()
280            .map_err(|e| ApprovalStoreError::Backend(format!("pool get: {e}")))?;
281        let batch: Option<BatchApproval> = conn
282            .query_row(
283                r#"SELECT batch_id, approver_hex, subject_id, server_pattern, tool_pattern,
284                          per_call_currency, per_call_units, total_currency, total_units,
285                          max_calls, not_before, not_after, used_calls, used_total_units, revoked
286                   FROM chio_hitl_batches WHERE batch_id = ?1"#,
287                params![batch_id],
288                row_to_batch,
289            )
290            .optional()
291            .map_err(|e| ApprovalStoreError::Backend(format!("get: {e}")))?;
292        Ok(batch)
293    }
294}