use async_trait::async_trait; use rusqlite::Connection; use tokio::sync::Mutex; use std::path::Path; use chrono::{DateTime, Utc}; use super::{StorageRepository, UptimeMessage, UptimeStorageModel}; pub struct SQLiteRepository { conn: Mutex } impl SQLiteRepository { pub fn new>(path: P) -> Self { let conn = Connection::open(path).unwrap(); conn.pragma_update(None, "journal_mode", "WAL").unwrap(); conn.pragma_update(None, "foreign_keys", "ON").unwrap(); conn.pragma_update(None, "synchronous", "NORMAL").unwrap(); conn.execute_batch( r#" CREATE TABLE IF NOT EXISTS agents ( id TEXT PRIMARY KEY, first_seen TEXT NOT NULL, last_seen TEXT NOT NULL, message_count INTEGER NOT NULL DEFAULT 0 ) STRICT; CREATE INDEX IF NOT EXISTS idx_agents_id ON agents(id); CREATE INDEX IF NOT EXISTS idx_agents_times ON agents(first_seen, last_seen); "# ).unwrap(); Self { conn: Mutex::new(conn) } } } #[async_trait] impl StorageRepository for SQLiteRepository { async fn record_uptime(&self, agent: &str) -> anyhow::Result<()> { let conn = self.conn.lock().await; let now = Utc::now().to_rfc3339(); conn.execute(r#" INSERT INTO agents (id, first_seen, last_seen, message_count) VALUES (?1, ?2, ?2, 1) ON CONFLICT (id) DO UPDATE SET last_seen = excluded.last_seen, message_count = message_count + 1; "#, [agent, &now] )?; Ok(()) } async fn get_agents(&self) -> anyhow::Result> { let conn = self.conn.lock().await; let mut stmt = conn.prepare("SELECT id, first_seen, last_seen, message_count FROM agents")?; let result = stmt.query_map([], |row| { let first_seen: DateTime = row.get::<_, String>(1)?.parse().unwrap(); let last_seen: DateTime = row.get::<_, String>(2)?.parse().unwrap(); Ok(UptimeStorageModel { id: row.get(0)?, first_seen, last_seen, message_count: row.get(3)?, }) })?; let models: Result, _> = result.collect(); Ok(models?.into_iter().map(Into::into).collect()) } }