tiny-counter 0.1.0

Track event counts across time windows with fixed memory and fast queries
Documentation
use std::{
    path::Path,
    sync::{Arc, Mutex},
};

use rusqlite::{params, Connection, OptionalExtension};

use crate::{Error, Result, Storage};

const TABLE_NAME: &str = "events";

/// SQLite storage backend.
///
/// Stores data in a SQLite database with ACID guarantees. Suitable for
/// production use with full transaction support.
///
/// The connection is wrapped in Arc<Mutex<>> to satisfy Send + Sync requirements.
/// Supports atomic transactions via begin_transaction/commit_transaction.
pub struct Sqlite {
    conn: Arc<Mutex<Connection>>,
    in_transaction: bool,
}

impl Sqlite {
    /// Opens or creates a SQLite database at the given path.
    ///
    /// Creates the necessary tables if they don't exist.
    pub fn open(path: impl AsRef<Path>) -> Result<Self> {
        let conn = Connection::open(path)
            .map_err(|e| Error::Storage(format!("Failed to open database: {}", e)))?;

        let storage = Self {
            conn: Arc::new(Mutex::new(conn)),
            in_transaction: false,
        };
        storage.create_tables()?;
        Ok(storage)
    }

    fn create_tables(&self) -> Result<()> {
        let conn = self
            .conn
            .lock()
            .map_err(|e| Error::Storage(format!("Failed to acquire lock: {}", e)))?;

        conn.execute(
            &format!(
                "CREATE TABLE IF NOT EXISTS {} (
                    key TEXT PRIMARY KEY,
                    data BLOB NOT NULL
                )",
                TABLE_NAME
            ),
            [],
        )
        .map_err(|e| Error::Storage(format!("Failed to create table: {}", e)))?;
        Ok(())
    }
}

impl Storage for Sqlite {
    fn save(&mut self, key: &str, data: Vec<u8>) -> Result<()> {
        let conn = self
            .conn
            .lock()
            .map_err(|e| Error::Storage(format!("Failed to acquire lock: {}", e)))?;

        conn.execute(
            &format!(
                "INSERT OR REPLACE INTO {} (key, data) VALUES (?1, ?2)",
                TABLE_NAME
            ),
            params![key, &data],
        )
        .map_err(|e| Error::Storage(format!("Failed to save data: {}", e)))?;
        Ok(())
    }

    fn load(&self, key: &str) -> Result<Option<Vec<u8>>> {
        let conn = self
            .conn
            .lock()
            .map_err(|e| Error::Storage(format!("Failed to acquire lock: {}", e)))?;

        let mut stmt = conn
            .prepare(&format!("SELECT data FROM {} WHERE key = ?1", TABLE_NAME))
            .map_err(|e| Error::Storage(format!("Failed to prepare statement: {}", e)))?;

        let result = stmt
            .query_row(params![key], |row| row.get(0))
            .optional()
            .map_err(|e| Error::Storage(format!("Failed to load data: {}", e)))?;

        Ok(result)
    }

    fn delete(&mut self, key: &str) -> Result<()> {
        let conn = self
            .conn
            .lock()
            .map_err(|e| Error::Storage(format!("Failed to acquire lock: {}", e)))?;

        conn.execute(
            &format!("DELETE FROM {} WHERE key = ?1", TABLE_NAME),
            params![key],
        )
        .map_err(|e| Error::Storage(format!("Failed to delete data: {}", e)))?;
        Ok(())
    }

    fn list_keys(&self) -> Result<Vec<String>> {
        let conn = self
            .conn
            .lock()
            .map_err(|e| Error::Storage(format!("Failed to acquire lock: {}", e)))?;

        let mut stmt = conn
            .prepare(&format!("SELECT key FROM {}", TABLE_NAME))
            .map_err(|e| Error::Storage(format!("Failed to prepare statement: {}", e)))?;

        let keys = stmt
            .query_map([], |row| row.get(0))
            .map_err(|e| Error::Storage(format!("Failed to query keys: {}", e)))?
            .collect::<std::result::Result<Vec<String>, _>>()
            .map_err(|e| Error::Storage(format!("Failed to collect keys: {}", e)))?;

        Ok(keys)
    }

    fn begin_transaction(&mut self) -> Result<()> {
        if self.in_transaction {
            return Err(Error::Storage(
                "Transaction already in progress".to_string(),
            ));
        }

        let conn = self
            .conn
            .lock()
            .map_err(|e| Error::Storage(format!("Failed to acquire lock: {}", e)))?;

        conn.execute("BEGIN TRANSACTION", [])
            .map_err(|e| Error::Storage(format!("Failed to begin transaction: {}", e)))?;

        self.in_transaction = true;
        Ok(())
    }

    fn commit_transaction(&mut self) -> Result<()> {
        if !self.in_transaction {
            return Err(Error::Storage("No transaction in progress".to_string()));
        }

        let conn = self
            .conn
            .lock()
            .map_err(|e| Error::Storage(format!("Failed to acquire lock: {}", e)))?;

        conn.execute("COMMIT", [])
            .map_err(|e| Error::Storage(format!("Failed to commit transaction: {}", e)))?;

        self.in_transaction = false;
        Ok(())
    }

    fn rollback_transaction(&mut self) -> Result<()> {
        if !self.in_transaction {
            return Err(Error::Storage("No transaction in progress".to_string()));
        }

        let conn = self
            .conn
            .lock()
            .map_err(|e| Error::Storage(format!("Failed to acquire lock: {}", e)))?;

        conn.execute("ROLLBACK", [])
            .map_err(|e| Error::Storage(format!("Failed to rollback transaction: {}", e)))?;

        self.in_transaction = false;
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use tempfile::TempDir;

    use super::*;

    fn temp_db() -> (TempDir, String) {
        let dir = TempDir::new().unwrap();
        let path = dir.path().join("test.db");
        (dir, path.to_string_lossy().to_string())
    }

    #[test]
    fn test_save_load_roundtrip() {
        let (_dir, path) = temp_db();
        let mut storage = Sqlite::open(&path).unwrap();

        let data = vec![1, 2, 3, 4, 5];
        storage.save("test_key", data.clone()).unwrap();

        let loaded = storage.load("test_key").unwrap();
        assert_eq!(loaded, Some(data));
    }

    #[test]
    fn test_database_created_if_not_exists() {
        let (_dir, path) = temp_db();
        let storage = Sqlite::open(&path).unwrap();

        // Verify table exists by querying it
        let result = storage.load("any_key").unwrap();
        assert_eq!(result, None);
    }

    #[test]
    fn test_tables_created_automatically() {
        let (_dir, path) = temp_db();
        let storage = Sqlite::open(&path).unwrap();

        // Should be able to perform operations immediately
        let mut storage = storage;
        storage.save("test", vec![1, 2, 3]).unwrap();
        let loaded = storage.load("test").unwrap();
        assert_eq!(loaded, Some(vec![1, 2, 3]));
    }

    #[test]
    fn test_multiple_saves_are_atomic() {
        let (_dir, path) = temp_db();
        let mut storage = Sqlite::open(&path).unwrap();

        // Save multiple entries
        storage.save("key1", vec![1, 2, 3]).unwrap();
        storage.save("key2", vec![4, 5, 6]).unwrap();
        storage.save("key3", vec![7, 8, 9]).unwrap();

        // All should be retrievable
        assert_eq!(storage.load("key1").unwrap(), Some(vec![1, 2, 3]));
        assert_eq!(storage.load("key2").unwrap(), Some(vec![4, 5, 6]));
        assert_eq!(storage.load("key3").unwrap(), Some(vec![7, 8, 9]));
    }

    #[test]
    fn test_save_large_blob() {
        let (_dir, path) = temp_db();
        let mut storage = Sqlite::open(&path).unwrap();

        // Create a 2MB blob
        let large_data: Vec<u8> = (0..2_000_000).map(|i| (i % 256) as u8).collect();
        storage.save("large", large_data.clone()).unwrap();

        let loaded = storage.load("large").unwrap();
        assert_eq!(loaded, Some(large_data));
    }

    #[test]
    fn test_delete_removes_data() {
        let (_dir, path) = temp_db();
        let mut storage = Sqlite::open(&path).unwrap();

        storage.save("test", vec![1, 2, 3]).unwrap();
        storage.delete("test").unwrap();

        let loaded = storage.load("test").unwrap();
        assert_eq!(loaded, None);
    }

    #[test]
    fn test_overwrite_existing_key() {
        let (_dir, path) = temp_db();
        let mut storage = Sqlite::open(&path).unwrap();

        storage.save("test", vec![1, 2, 3]).unwrap();
        storage.save("test", vec![4, 5, 6]).unwrap();

        let loaded = storage.load("test").unwrap();
        assert_eq!(loaded, Some(vec![4, 5, 6]));
    }

    #[test]
    fn test_persistence_across_connections() {
        let (_dir, path) = temp_db();

        // First connection: save data
        {
            let mut storage = Sqlite::open(&path).unwrap();
            storage.save("key1", vec![1, 2, 3]).unwrap();
            storage.save("key2", vec![4, 5, 6]).unwrap();
        }

        // Second connection: load data
        {
            let storage = Sqlite::open(&path).unwrap();
            assert_eq!(storage.load("key1").unwrap(), Some(vec![1, 2, 3]));
            assert_eq!(storage.load("key2").unwrap(), Some(vec![4, 5, 6]));
        }
    }

    #[test]
    fn test_load_nonexistent_returns_none() {
        let (_dir, path) = temp_db();
        let storage = Sqlite::open(&path).unwrap();

        let loaded = storage.load("nonexistent").unwrap();
        assert_eq!(loaded, None);
    }

    #[test]
    fn test_transaction_commit() {
        let (_dir, path) = temp_db();
        let mut storage = Sqlite::open(&path).unwrap();

        // Begin transaction
        storage.begin_transaction().unwrap();

        // Save data within transaction
        storage.save("key1", vec![1, 2, 3]).unwrap();
        storage.save("key2", vec![4, 5, 6]).unwrap();

        // Data should not be visible yet (in transaction)
        // Commit transaction
        storage.commit_transaction().unwrap();

        // Data should now be persisted
        assert_eq!(storage.load("key1").unwrap(), Some(vec![1, 2, 3]));
        assert_eq!(storage.load("key2").unwrap(), Some(vec![4, 5, 6]));
    }

    #[test]
    fn test_transaction_rollback() {
        let (_dir, path) = temp_db();
        let mut storage = Sqlite::open(&path).unwrap();

        // Save initial data
        storage.save("key1", vec![1, 2, 3]).unwrap();

        // Begin transaction
        storage.begin_transaction().unwrap();

        // Modify data within transaction
        storage.save("key1", vec![4, 5, 6]).unwrap();
        storage.save("key2", vec![7, 8, 9]).unwrap();

        // Rollback transaction
        storage.rollback_transaction().unwrap();

        // Original data should be unchanged
        assert_eq!(storage.load("key1").unwrap(), Some(vec![1, 2, 3]));
        // New key should not exist
        assert_eq!(storage.load("key2").unwrap(), None);
    }

    #[test]
    fn test_transaction_nested_begin_fails() {
        let (_dir, path) = temp_db();
        let mut storage = Sqlite::open(&path).unwrap();

        // Begin first transaction
        storage.begin_transaction().unwrap();

        // Attempting to begin another should fail
        let result = storage.begin_transaction();
        assert!(result.is_err());
        assert!(result
            .unwrap_err()
            .to_string()
            .contains("Transaction already in progress"));

        // Clean up
        storage.rollback_transaction().unwrap();
    }

    #[test]
    fn test_commit_without_begin_fails() {
        let (_dir, path) = temp_db();
        let mut storage = Sqlite::open(&path).unwrap();

        let result = storage.commit_transaction();
        assert!(result.is_err());
        assert!(result
            .unwrap_err()
            .to_string()
            .contains("No transaction in progress"));
    }

    #[test]
    fn test_rollback_without_begin_fails() {
        let (_dir, path) = temp_db();
        let mut storage = Sqlite::open(&path).unwrap();

        let result = storage.rollback_transaction();
        assert!(result.is_err());
        assert!(result
            .unwrap_err()
            .to_string()
            .contains("No transaction in progress"));
    }
}