1use std::str::FromStr;
4
5use async_trait::async_trait;
6use sqlx::sqlite::{SqliteConnectOptions, SqlitePoolOptions};
7use sqlx::SqlitePool;
8
9use crate::ledger::{LedgerError, LedgerStore, UsageEntry};
10
11pub struct SqliteLedger {
12 pool: SqlitePool,
13}
14
15impl SqliteLedger {
16 pub async fn connect(dsn: &str) -> Result<Self, LedgerError> {
24 let opts = SqliteConnectOptions::from_str(dsn)
25 .map_err(|e| LedgerError::Backend(e.to_string()))?
26 .create_if_missing(true);
27
28 let pool = SqlitePoolOptions::new()
29 .max_connections(1)
30 .connect_with(opts)
31 .await
32 .map_err(|e| LedgerError::Backend(e.to_string()))?;
33
34 sqlx::raw_sql(include_str!("../../migrations/0001_usage_events.sql"))
37 .execute(&pool)
38 .await
39 .map_err(|e| LedgerError::Backend(e.to_string()))?;
40
41 Ok(Self { pool })
42 }
43}
44
45#[async_trait]
46impl LedgerStore for SqliteLedger {
47 async fn record(&self, e: &UsageEntry) -> Result<(), LedgerError> {
48 sqlx::query(
49 "INSERT INTO usage_events \
50 (ts, tenant, workspace, route, provider, model, lane, \
51 input_tokens, output_tokens, cost_usd, request_id, status) \
52 VALUES (?,?,?,?,?,?,?,?,?,?,?,?)",
53 )
54 .bind(e.ts.to_rfc3339())
55 .bind(&e.tenant)
56 .bind(&e.workspace)
57 .bind(&e.route)
58 .bind(&e.provider)
59 .bind(&e.model)
60 .bind(&e.lane)
61 .bind(e.input_tokens as i64)
62 .bind(e.output_tokens as i64)
63 .bind(e.cost_usd)
64 .bind(&e.request_id)
65 .bind(&e.status)
66 .execute(&self.pool)
67 .await
68 .map_err(|e| LedgerError::Backend(e.to_string()))?;
69 Ok(())
70 }
71}