Skip to main content

synapse/ledger/
sqlite.rs

1//! SQLite ledger backend (feature `ledger-sqlite`).
2
3use 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    /// Connect (DSN like `sqlite://synapse.db?mode=rwc` or `sqlite::memory:`)
17    /// and create the table if absent.
18    ///
19    /// Uses `max_connections(1)` so that both file-backed and in-memory databases
20    /// work correctly: with `sqlite::memory:` every connection gets its own
21    /// isolated database, so a single connection ensures the migration and all
22    /// subsequent writes share the same in-memory DB.
23    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        // Run the multi-statement migration via raw_sql which supports
35        // multiple `;`-separated statements in a single call.
36        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}