use anyhow::{Context, Result};
use chrono::Utc;
use rusqlite::{params, Connection};
use std::path::PathBuf;
#[derive(Debug, Clone)]
pub struct ExecutionRecord {
pub program: String,
pub args: String,
pub tokens_before: usize,
pub tokens_after: usize,
pub exit_code: i32,
pub duration_ms: u64,
pub timestamp: String,
}
pub struct ExecutionStore {
conn: Connection,
}
impl ExecutionStore {
pub fn open(path: &PathBuf) -> Result<Self> {
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)?;
}
let conn = Connection::open(path).context("opening execution store")?;
conn.execute_batch(
"CREATE TABLE IF NOT EXISTS executions (
id INTEGER PRIMARY KEY AUTOINCREMENT,
program TEXT NOT NULL,
args TEXT NOT NULL,
tokens_before INTEGER NOT NULL,
tokens_after INTEGER NOT NULL,
exit_code INTEGER NOT NULL,
duration_ms INTEGER NOT NULL,
timestamp TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_program ON executions(program);
CREATE INDEX IF NOT EXISTS idx_timestamp ON executions(timestamp);",
)?;
Ok(Self { conn })
}
pub fn in_memory() -> Result<Self> {
let conn = Connection::open_in_memory()?;
conn.execute_batch(
"CREATE TABLE IF NOT EXISTS executions (
id INTEGER PRIMARY KEY AUTOINCREMENT,
program TEXT NOT NULL,
args TEXT NOT NULL,
tokens_before INTEGER NOT NULL,
tokens_after INTEGER NOT NULL,
exit_code INTEGER NOT NULL,
duration_ms INTEGER NOT NULL,
timestamp TEXT NOT NULL
);",
)?;
Ok(Self { conn })
}
pub fn record(&self, rec: &ExecutionRecord) -> Result<()> {
self.conn.execute(
"INSERT INTO executions (program, args, tokens_before, tokens_after, exit_code, duration_ms, timestamp)
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
params![
rec.program, rec.args, rec.tokens_before, rec.tokens_after,
rec.exit_code, rec.duration_ms, rec.timestamp
],
)?;
Ok(())
}
pub fn total_saved(&self) -> Result<usize> {
let val: i64 = self.conn.query_row(
"SELECT COALESCE(SUM(tokens_before - tokens_after), 0) FROM executions",
[],
|row| row.get(0),
)?;
Ok(val as usize)
}
pub fn summary(&self) -> Result<Vec<(String, usize, usize, f64)>> {
let mut stmt = self.conn.prepare(
"SELECT program,
SUM(tokens_before) as tb,
SUM(tokens_after) as ta,
AVG(CAST(tokens_before - tokens_after AS REAL) / NULLIF(tokens_before, 0)) * 100.0 as pct
FROM executions
GROUP BY program
ORDER BY tb DESC
LIMIT 20",
)?;
let rows = stmt.query_map([], |row| {
Ok((
row.get::<_, String>(0)?,
row.get::<_, i64>(1)? as usize,
row.get::<_, i64>(2)? as usize,
row.get::<_, f64>(3).unwrap_or(0.0),
))
})?;
rows.collect::<Result<Vec<_>, _>>().map_err(Into::into)
}
pub fn recent(&self, limit: usize) -> Result<Vec<ExecutionRecord>> {
let mut stmt = self.conn.prepare(
"SELECT program, args, tokens_before, tokens_after, exit_code, duration_ms, timestamp
FROM executions ORDER BY id DESC LIMIT ?1",
)?;
let rows = stmt.query_map(params![limit as i64], |row| {
Ok(ExecutionRecord {
program: row.get(0)?,
args: row.get(1)?,
tokens_before: row.get::<_, i64>(2)? as usize,
tokens_after: row.get::<_, i64>(3)? as usize,
exit_code: row.get(4)?,
duration_ms: row.get::<_, i64>(5)? as u64,
timestamp: row.get(6)?,
})
})?;
rows.collect::<Result<Vec<_>, _>>().map_err(Into::into)
}
pub fn new_record(
program: &str,
args: &[String],
tokens_before: usize,
tokens_after: usize,
exit_code: i32,
duration_ms: u64,
) -> ExecutionRecord {
ExecutionRecord {
program: program.to_string(),
args: args.join(" "),
tokens_before,
tokens_after,
exit_code,
duration_ms,
timestamp: Utc::now().to_rfc3339(),
}
}
}