//! Usage rollup DAO //! //! Aggregates proxy_request_logs into daily rollups and prunes old detail rows. use crate::database::{lock_conn, Database}; use crate::error::AppError; impl Database { /// Aggregate proxy_request_logs older than `retain_days` into usage_daily_rollups, /// then delete the aggregated detail rows. /// Returns the number of deleted detail rows. pub fn rollup_and_prune(&self, retain_days: i64) -> Result { let cutoff = chrono::Utc::now().timestamp() - retain_days * 86400; let conn = lock_conn!(self.conn); // Check if there are any rows to process let count: i64 = conn .query_row( "SELECT COUNT(*) FROM proxy_request_logs WHERE created_at < ?1", [cutoff], |row| row.get(0), ) .map_err(|e| AppError::Database(e.to_string()))?; if count == 0 { return Ok(0); } // Use a savepoint for atomicity conn.execute("SAVEPOINT rollup_prune;", []) .map_err(|e| AppError::Database(e.to_string()))?; let result = Self::do_rollup_and_prune(&conn, cutoff); match result { Ok(deleted) => { conn.execute("RELEASE rollup_prune;", []) .map_err(|e| AppError::Database(e.to_string()))?; if deleted > 0 { log::info!( "Rolled up and pruned {deleted} proxy_request_logs (retain={retain_days}d)" ); } Ok(deleted) } Err(e) => { conn.execute("ROLLBACK TO rollup_prune;", []).ok(); conn.execute("RELEASE rollup_prune;", []).ok(); Err(e) } } } fn do_rollup_and_prune(conn: &rusqlite::Connection, cutoff: i64) -> Result { // Aggregate old logs, merging with any pre-existing rollup rows via LEFT JOIN. conn.execute( "INSERT OR REPLACE INTO usage_daily_rollups (date, app_type, provider_id, model, request_count, success_count, input_tokens, output_tokens, cache_read_tokens, cache_creation_tokens, total_cost_usd, avg_latency_ms) SELECT d, a, p, m, COALESCE(old.request_count, 0) + new_req, COALESCE(old.success_count, 0) + new_succ, COALESCE(old.input_tokens, 0) + new_in, COALESCE(old.output_tokens, 0) + new_out, COALESCE(old.cache_read_tokens, 0) + new_cr, COALESCE(old.cache_creation_tokens, 0) + new_cc, CAST(COALESCE(CAST(old.total_cost_usd AS REAL), 0) + new_cost AS TEXT), CASE WHEN COALESCE(old.request_count, 0) + new_req > 0 THEN (COALESCE(old.avg_latency_ms, 0) * COALESCE(old.request_count, 0) + new_lat * new_req) / (COALESCE(old.request_count, 0) + new_req) ELSE 0 END FROM ( SELECT date(created_at, 'unixepoch', 'localtime') as d, app_type as a, provider_id as p, model as m, COUNT(*) as new_req, SUM(CASE WHEN status_code >= 200 AND status_code < 300 THEN 1 ELSE 0 END) as new_succ, COALESCE(SUM(input_tokens), 0) as new_in, COALESCE(SUM(output_tokens), 0) as new_out, COALESCE(SUM(cache_read_tokens), 0) as new_cr, COALESCE(SUM(cache_creation_tokens), 0) as new_cc, COALESCE(SUM(CAST(total_cost_usd AS REAL)), 0) as new_cost, COALESCE(AVG(latency_ms), 0) as new_lat FROM proxy_request_logs WHERE created_at < ?1 GROUP BY d, a, p, m ) agg LEFT JOIN usage_daily_rollups old ON old.date = agg.d AND old.app_type = agg.a AND old.provider_id = agg.p AND old.model = agg.m", [cutoff], ) .map_err(|e| AppError::Database(format!("Rollup aggregation failed: {e}")))?; // Delete the aggregated detail rows let deleted = conn .execute( "DELETE FROM proxy_request_logs WHERE created_at < ?1", [cutoff], ) .map_err(|e| AppError::Database(format!("Pruning old logs failed: {e}")))?; Ok(deleted as u64) } } #[cfg(test)] mod tests { use crate::database::Database; use crate::error::AppError; #[test] fn test_rollup_and_prune() -> Result<(), AppError> { let db = Database::memory()?; let now = chrono::Utc::now().timestamp(); let old_ts = now - 40 * 86400; // 40 days ago let recent_ts = now - 5 * 86400; // 5 days ago { let conn = crate::database::lock_conn!(db.conn); for i in 0..5 { conn.execute( "INSERT INTO proxy_request_logs ( request_id, provider_id, app_type, model, input_tokens, output_tokens, total_cost_usd, latency_ms, status_code, created_at ) VALUES (?1, 'p1', 'claude', 'claude-3', 100, 50, '0.01', 100, 200, ?2)", rusqlite::params![format!("old-{i}"), old_ts + i as i64], )?; } for i in 0..3 { conn.execute( "INSERT INTO proxy_request_logs ( request_id, provider_id, app_type, model, input_tokens, output_tokens, total_cost_usd, latency_ms, status_code, created_at ) VALUES (?1, 'p1', 'claude', 'claude-3', 200, 100, '0.02', 150, 200, ?2)", rusqlite::params![format!("recent-{i}"), recent_ts + i as i64], )?; } } let deleted = db.rollup_and_prune(30)?; assert_eq!(deleted, 5); // Verify rollup data let conn = crate::database::lock_conn!(db.conn); let count: i64 = conn.query_row( "SELECT request_count FROM usage_daily_rollups WHERE app_type = 'claude'", [], |row| row.get(0), )?; assert_eq!(count, 5); // Verify recent logs untouched let remaining: i64 = conn.query_row("SELECT COUNT(*) FROM proxy_request_logs", [], |row| { row.get(0) })?; assert_eq!(remaining, 3); Ok(()) } #[test] fn test_rollup_noop_when_no_old_data() -> Result<(), AppError> { let db = Database::memory()?; assert_eq!(db.rollup_and_prune(30)?, 0); Ok(()) } #[test] fn test_rollup_merges_with_existing() -> Result<(), AppError> { let db = Database::memory()?; let now = chrono::Utc::now().timestamp(); let old_ts = now - 40 * 86400; { let conn = crate::database::lock_conn!(db.conn); let date_str = chrono::DateTime::from_timestamp(old_ts, 0) .unwrap() .format("%Y-%m-%d") .to_string(); conn.execute( "INSERT INTO usage_daily_rollups (date, app_type, provider_id, model, request_count, success_count, input_tokens, output_tokens, total_cost_usd, avg_latency_ms) VALUES (?1, 'claude', 'p1', 'claude-3', 10, 10, 1000, 500, '0.10', 100)", [&date_str], )?; for i in 0..3 { conn.execute( "INSERT INTO proxy_request_logs ( request_id, provider_id, app_type, model, input_tokens, output_tokens, total_cost_usd, latency_ms, status_code, created_at ) VALUES (?1, 'p1', 'claude', 'claude-3', 100, 50, '0.01', 200, 200, ?2)", rusqlite::params![format!("merge-{i}"), old_ts + i as i64], )?; } } let deleted = db.rollup_and_prune(30)?; assert_eq!(deleted, 3); let conn = crate::database::lock_conn!(db.conn); let (count, input): (i64, i64) = conn.query_row( "SELECT request_count, input_tokens FROM usage_daily_rollups WHERE app_type = 'claude' AND provider_id = 'p1'", [], |row| Ok((row.get(0)?, row.get(1)?)), )?; assert_eq!(count, 13, "10 existing + 3 new"); assert_eq!(input, 1300, "1000 existing + 300 new"); Ok(()) } }