Skip to main content

chio_store_sqlite/
batch_approval_store.rs

1//! 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::scope::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
21/// Batch-approval-store schema revision. Bump on every schema-affecting change.
22const BATCH_APPROVAL_STORE_SUPPORTED_SCHEMA_VERSION: i32 = 0;
23/// Stable key under which this store records its schema revision in the shared
24/// keyed metadata table, distinct from any co-located store's key.
25const BATCH_APPROVAL_STORE_SCHEMA_KEY: &str = "batch_approval";
26/// Tables shipped before schema stamping existed, used to adopt a pre-stamping
27/// batch-approval database rather than reject it as foreign.
28const BATCH_APPROVAL_STORE_LEGACY_ANCHOR_TABLES: &[&str] = &["chio_hitl_batches"];
29
30impl SqliteBatchApprovalStore {
31    pub fn open(path: impl AsRef<Path>) -> Result<Self, ApprovalStoreError> {
32        let path = path.as_ref();
33        // Resolve any `file:` URI to its on-disk parent before creating it, so a
34        // URI-configured store creates the real backing directory rather than a
35        // bogus scheme-prefixed one.
36        if let Some(parent) = crate::sqlite_parent_dir_to_create(path) {
37            fs::create_dir_all(&parent)
38                .map_err(|e| ApprovalStoreError::Backend(format!("create dir: {e}")))?;
39        }
40        let manager = SqliteConnectionManager::file(path);
41        let pool = Pool::builder()
42            .max_size(4)
43            .build(manager)
44            .map_err(|e| ApprovalStoreError::Backend(format!("pool build: {e}")))?;
45        let store = Self { pool };
46        store.run_migrations()?;
47        Ok(store)
48    }
49
50    pub fn open_in_memory() -> Result<Self, ApprovalStoreError> {
51        let manager = SqliteConnectionManager::memory();
52        let pool = Pool::builder()
53            .max_size(1)
54            .build(manager)
55            .map_err(|e| ApprovalStoreError::Backend(format!("pool build: {e}")))?;
56        let store = Self { pool };
57        store.run_migrations()?;
58        Ok(store)
59    }
60
61    fn run_migrations(&self) -> Result<(), ApprovalStoreError> {
62        let conn = self
63            .pool
64            .get()
65            .map_err(|e| ApprovalStoreError::Backend(format!("pool get: {e}")))?;
66        crate::check_schema_version(
67            &conn,
68            BATCH_APPROVAL_STORE_SCHEMA_KEY,
69            BATCH_APPROVAL_STORE_SUPPORTED_SCHEMA_VERSION,
70            BATCH_APPROVAL_STORE_LEGACY_ANCHOR_TABLES,
71        )
72        .map_err(|error| ApprovalStoreError::Backend(error.to_string()))?;
73        conn.execute_batch(
74            r#"
75            PRAGMA journal_mode = WAL;
76            PRAGMA synchronous = FULL;
77            PRAGMA busy_timeout = 5000;
78
79            CREATE TABLE IF NOT EXISTS chio_hitl_batches (
80                batch_id TEXT PRIMARY KEY,
81                approver_hex TEXT NOT NULL,
82                subject_id TEXT NOT NULL,
83                server_pattern TEXT NOT NULL,
84                tool_pattern TEXT NOT NULL,
85                per_call_currency TEXT,
86                per_call_units INTEGER,
87                total_currency TEXT,
88                total_units INTEGER,
89                max_calls INTEGER,
90                not_before INTEGER NOT NULL,
91                not_after INTEGER NOT NULL,
92                used_calls INTEGER NOT NULL DEFAULT 0,
93                used_total_units INTEGER NOT NULL DEFAULT 0,
94                revoked INTEGER NOT NULL DEFAULT 0
95            );
96            CREATE INDEX IF NOT EXISTS idx_chio_hitl_batches_subject
97                ON chio_hitl_batches(subject_id, revoked);
98            "#,
99        )
100        .map_err(|e| ApprovalStoreError::Backend(format!("migration: {e}")))?;
101        crate::stamp_schema_version(
102            &conn,
103            BATCH_APPROVAL_STORE_SCHEMA_KEY,
104            BATCH_APPROVAL_STORE_SUPPORTED_SCHEMA_VERSION,
105        )
106        .map_err(|error| ApprovalStoreError::Backend(error.to_string()))?;
107        Ok(())
108    }
109}
110
111fn row_to_batch(row: &rusqlite::Row<'_>) -> rusqlite::Result<BatchApproval> {
112    let per_call_currency: Option<String> = row.get(5)?;
113    let per_call_units: Option<i64> = row.get(6)?;
114    let total_currency: Option<String> = row.get(7)?;
115    let total_units: Option<i64> = row.get(8)?;
116    let max_calls: Option<i64> = row.get(9)?;
117    let revoked: i64 = row.get(14)?;
118    Ok(BatchApproval {
119        batch_id: row.get(0)?,
120        approver_hex: row.get(1)?,
121        subject_id: row.get(2)?,
122        server_pattern: row.get(3)?,
123        tool_pattern: row.get(4)?,
124        max_amount_per_call: match (per_call_currency, per_call_units) {
125            (Some(currency), Some(units)) => Some(MonetaryAmount {
126                currency,
127                units: units.max(0) as u64,
128            }),
129            _ => None,
130        },
131        max_total_amount: match (total_currency, total_units) {
132            (Some(currency), Some(units)) => Some(MonetaryAmount {
133                currency,
134                units: units.max(0) as u64,
135            }),
136            _ => None,
137        },
138        max_calls: max_calls.map(|v| v.max(0) as u32),
139        not_before: row.get::<_, i64>(10)?.max(0) as u64,
140        not_after: row.get::<_, i64>(11)?.max(0) as u64,
141        used_calls: row.get::<_, i64>(12)?.max(0) as u32,
142        used_total_units: row.get::<_, i64>(13)?.max(0) as u64,
143        revoked: revoked != 0,
144    })
145}
146
147fn pattern_matches(pattern: &str, value: &str) -> bool {
148    if pattern == "*" {
149        return true;
150    }
151    if let Some(prefix) = pattern.strip_suffix('*') {
152        return value.starts_with(prefix);
153    }
154    pattern == value
155}
156
157fn amount_fits(batch: &BatchApproval, amount: Option<&MonetaryAmount>) -> bool {
158    let Some(amt) = amount else {
159        return batch.max_amount_per_call.is_none() && batch.max_total_amount.is_none();
160    };
161    if let Some(per_call) = &batch.max_amount_per_call {
162        if amt.currency != per_call.currency || amt.units > per_call.units {
163            return false;
164        }
165    }
166    if let Some(total) = &batch.max_total_amount {
167        if amt.currency != total.currency
168            || batch.used_total_units.saturating_add(amt.units) > total.units
169        {
170            return false;
171        }
172    }
173    true
174}
175
176impl BatchApprovalStore for SqliteBatchApprovalStore {
177    fn store(&self, batch: &BatchApproval) -> Result<(), ApprovalStoreError> {
178        let conn = self
179            .pool
180            .get()
181            .map_err(|e| ApprovalStoreError::Backend(format!("pool get: {e}")))?;
182        conn.execute(
183            r#"INSERT OR REPLACE INTO chio_hitl_batches (
184                batch_id, approver_hex, subject_id, server_pattern, tool_pattern,
185                per_call_currency, per_call_units, total_currency, total_units,
186                max_calls, not_before, not_after, used_calls, used_total_units, revoked
187            ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15)"#,
188            params![
189                batch.batch_id,
190                batch.approver_hex,
191                batch.subject_id,
192                batch.server_pattern,
193                batch.tool_pattern,
194                batch
195                    .max_amount_per_call
196                    .as_ref()
197                    .map(|a| a.currency.clone()),
198                batch.max_amount_per_call.as_ref().map(|a| a.units as i64),
199                batch.max_total_amount.as_ref().map(|a| a.currency.clone()),
200                batch.max_total_amount.as_ref().map(|a| a.units as i64),
201                batch.max_calls.map(|c| c as i64),
202                batch.not_before as i64,
203                batch.not_after as i64,
204                batch.used_calls as i64,
205                batch.used_total_units as i64,
206                batch.revoked as i64,
207            ],
208        )
209        .map_err(|e| ApprovalStoreError::Backend(format!("insert batch: {e}")))?;
210        Ok(())
211    }
212
213    fn find_matching(
214        &self,
215        subject_id: &str,
216        server_id: &str,
217        tool_name: &str,
218        amount: Option<&MonetaryAmount>,
219        now: u64,
220    ) -> Result<Option<BatchApproval>, ApprovalStoreError> {
221        let conn = self
222            .pool
223            .get()
224            .map_err(|e| ApprovalStoreError::Backend(format!("pool get: {e}")))?;
225        // Pull every non-revoked candidate for this subject within the
226        // time window, then apply pattern + amount filters in Rust.
227        let mut stmt = conn
228            .prepare(
229                r#"SELECT batch_id, approver_hex, subject_id, server_pattern, tool_pattern,
230                          per_call_currency, per_call_units, total_currency, total_units,
231                          max_calls, not_before, not_after, used_calls, used_total_units, revoked
232                   FROM chio_hitl_batches
233                   WHERE subject_id = ?1 AND revoked = 0
234                     AND not_before <= ?2 AND not_after > ?2"#,
235            )
236            .map_err(|e| ApprovalStoreError::Backend(format!("prepare: {e}")))?;
237        let rows = stmt
238            .query_map(params![subject_id, now as i64], row_to_batch)
239            .map_err(|e| ApprovalStoreError::Backend(format!("query: {e}")))?;
240        for row in rows {
241            let batch = row.map_err(|e| ApprovalStoreError::Backend(format!("row: {e}")))?;
242            if !pattern_matches(&batch.server_pattern, server_id) {
243                continue;
244            }
245            if !pattern_matches(&batch.tool_pattern, tool_name) {
246                continue;
247            }
248            if let Some(max) = batch.max_calls {
249                if batch.used_calls >= max {
250                    continue;
251                }
252            }
253            if !amount_fits(&batch, amount) {
254                continue;
255            }
256            return Ok(Some(batch));
257        }
258        Ok(None)
259    }
260
261    fn record_usage(
262        &self,
263        batch_id: &str,
264        amount: Option<&MonetaryAmount>,
265    ) -> Result<(), ApprovalStoreError> {
266        let conn = self
267            .pool
268            .get()
269            .map_err(|e| ApprovalStoreError::Backend(format!("pool get: {e}")))?;
270        let added_units = amount.map(|a| a.units as i64).unwrap_or(0);
271        let rows = conn.execute(
272            "UPDATE chio_hitl_batches SET used_calls = used_calls + 1, used_total_units = used_total_units + ?2 WHERE batch_id = ?1",
273            params![batch_id, added_units],
274        )
275        .map_err(|e| ApprovalStoreError::Backend(format!("update usage: {e}")))?;
276        if rows == 0 {
277            return Err(ApprovalStoreError::NotFound(batch_id.to_string()));
278        }
279        Ok(())
280    }
281
282    fn revoke(&self, batch_id: &str) -> Result<(), ApprovalStoreError> {
283        let conn = self
284            .pool
285            .get()
286            .map_err(|e| ApprovalStoreError::Backend(format!("pool get: {e}")))?;
287        let rows = conn
288            .execute(
289                "UPDATE chio_hitl_batches SET revoked = 1 WHERE batch_id = ?1",
290                params![batch_id],
291            )
292            .map_err(|e| ApprovalStoreError::Backend(format!("revoke: {e}")))?;
293        if rows == 0 {
294            return Err(ApprovalStoreError::NotFound(batch_id.to_string()));
295        }
296        Ok(())
297    }
298
299    fn get(&self, batch_id: &str) -> Result<Option<BatchApproval>, ApprovalStoreError> {
300        let conn = self
301            .pool
302            .get()
303            .map_err(|e| ApprovalStoreError::Backend(format!("pool get: {e}")))?;
304        let batch: Option<BatchApproval> = conn
305            .query_row(
306                r#"SELECT batch_id, approver_hex, subject_id, server_pattern, tool_pattern,
307                          per_call_currency, per_call_units, total_currency, total_units,
308                          max_calls, not_before, not_after, used_calls, used_total_units, revoked
309                   FROM chio_hitl_batches WHERE batch_id = ?1"#,
310                params![batch_id],
311                row_to_batch,
312            )
313            .optional()
314            .map_err(|e| ApprovalStoreError::Backend(format!("get: {e}")))?;
315        Ok(batch)
316    }
317}