1use anyhow::{Context, Result};
2use chrono::Utc;
3use rusqlite::{params, Connection};
4use std::path::PathBuf;
5
6#[derive(Debug, Clone)]
7pub struct ExecutionRecord {
8 pub program: String,
9 pub args: String,
10 pub tokens_before: usize,
11 pub tokens_after: usize,
12 pub exit_code: i32,
13 pub duration_ms: u64,
14 pub timestamp: String,
15}
16
17pub struct ExecutionStore {
18 conn: Connection,
19}
20
21impl ExecutionStore {
22 pub fn open(path: &PathBuf) -> Result<Self> {
23 if let Some(parent) = path.parent() {
24 std::fs::create_dir_all(parent)?;
25 }
26 let conn = Connection::open(path).context("opening execution store")?;
27 conn.execute_batch(
28 "CREATE TABLE IF NOT EXISTS executions (
29 id INTEGER PRIMARY KEY AUTOINCREMENT,
30 program TEXT NOT NULL,
31 args TEXT NOT NULL,
32 tokens_before INTEGER NOT NULL,
33 tokens_after INTEGER NOT NULL,
34 exit_code INTEGER NOT NULL,
35 duration_ms INTEGER NOT NULL,
36 timestamp TEXT NOT NULL
37 );
38 CREATE INDEX IF NOT EXISTS idx_program ON executions(program);
39 CREATE INDEX IF NOT EXISTS idx_timestamp ON executions(timestamp);",
40 )?;
41 Ok(Self { conn })
42 }
43
44 pub fn in_memory() -> Result<Self> {
45 let conn = Connection::open_in_memory()?;
46 conn.execute_batch(
47 "CREATE TABLE IF NOT EXISTS executions (
48 id INTEGER PRIMARY KEY AUTOINCREMENT,
49 program TEXT NOT NULL,
50 args TEXT NOT NULL,
51 tokens_before INTEGER NOT NULL,
52 tokens_after INTEGER NOT NULL,
53 exit_code INTEGER NOT NULL,
54 duration_ms INTEGER NOT NULL,
55 timestamp TEXT NOT NULL
56 );",
57 )?;
58 Ok(Self { conn })
59 }
60
61 pub fn record(&self, rec: &ExecutionRecord) -> Result<()> {
62 self.conn.execute(
63 "INSERT INTO executions (program, args, tokens_before, tokens_after, exit_code, duration_ms, timestamp)
64 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
65 params![
66 rec.program, rec.args, rec.tokens_before, rec.tokens_after,
67 rec.exit_code, rec.duration_ms, rec.timestamp
68 ],
69 )?;
70 Ok(())
71 }
72
73 pub fn total_saved(&self) -> Result<usize> {
74 let val: i64 = self.conn.query_row(
75 "SELECT COALESCE(SUM(tokens_before - tokens_after), 0) FROM executions",
76 [],
77 |row| row.get(0),
78 )?;
79 Ok(val as usize)
80 }
81
82 pub fn summary(&self) -> Result<Vec<(String, usize, usize, f64)>> {
83 let mut stmt = self.conn.prepare(
84 "SELECT program,
85 SUM(tokens_before) as tb,
86 SUM(tokens_after) as ta,
87 AVG(CAST(tokens_before - tokens_after AS REAL) / NULLIF(tokens_before, 0)) * 100.0 as pct
88 FROM executions
89 GROUP BY program
90 ORDER BY tb DESC
91 LIMIT 20",
92 )?;
93 let rows = stmt.query_map([], |row| {
94 Ok((
95 row.get::<_, String>(0)?,
96 row.get::<_, i64>(1)? as usize,
97 row.get::<_, i64>(2)? as usize,
98 row.get::<_, f64>(3).unwrap_or(0.0),
99 ))
100 })?;
101 rows.collect::<Result<Vec<_>, _>>().map_err(Into::into)
102 }
103
104 pub fn recent(&self, limit: usize) -> Result<Vec<ExecutionRecord>> {
105 let mut stmt = self.conn.prepare(
106 "SELECT program, args, tokens_before, tokens_after, exit_code, duration_ms, timestamp
107 FROM executions ORDER BY id DESC LIMIT ?1",
108 )?;
109 let rows = stmt.query_map(params![limit as i64], |row| {
110 Ok(ExecutionRecord {
111 program: row.get(0)?,
112 args: row.get(1)?,
113 tokens_before: row.get::<_, i64>(2)? as usize,
114 tokens_after: row.get::<_, i64>(3)? as usize,
115 exit_code: row.get(4)?,
116 duration_ms: row.get::<_, i64>(5)? as u64,
117 timestamp: row.get(6)?,
118 })
119 })?;
120 rows.collect::<Result<Vec<_>, _>>().map_err(Into::into)
121 }
122
123 pub fn new_record(
124 program: &str,
125 args: &[String],
126 tokens_before: usize,
127 tokens_after: usize,
128 exit_code: i32,
129 duration_ms: u64,
130 ) -> ExecutionRecord {
131 ExecutionRecord {
132 program: program.to_string(),
133 args: args.join(" "),
134 tokens_before,
135 tokens_after,
136 exit_code,
137 duration_ms,
138 timestamp: Utc::now().to_rfc3339(),
139 }
140 }
141}