status-monitor/server/src/storage/sqlite.rs

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<Connection>
}

impl SQLiteRepository {
    pub fn new<P: AsRef<Path>>(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<Vec<UptimeMessage>> {
        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<Utc> = row.get::<_, String>(1)?.parse().unwrap();
            let last_seen: DateTime<Utc> = row.get::<_, String>(2)?.parse().unwrap();

            Ok(UptimeStorageModel {
                id: row.get(0)?,
                first_seen,
                last_seen,
                message_count: row.get(3)?,
            })
        })?;

        let models: Result<Vec<UptimeStorageModel>, _> = result.collect();

        Ok(models?.into_iter().map(Into::into).collect())
    }
}