use std::{
path::Path,
sync::{Arc, Mutex},
};
use rusqlite::{params, Connection, OptionalExtension};
use crate::{Error, Result, Storage};
const TABLE_NAME: &str = "events";
pub struct Sqlite {
conn: Arc<Mutex<Connection>>,
in_transaction: bool,
}
impl Sqlite {
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();
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();
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();
storage.save("key1", vec![1, 2, 3]).unwrap();
storage.save("key2", vec![4, 5, 6]).unwrap();
storage.save("key3", vec![7, 8, 9]).unwrap();
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();
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();
{
let mut storage = Sqlite::open(&path).unwrap();
storage.save("key1", vec![1, 2, 3]).unwrap();
storage.save("key2", vec![4, 5, 6]).unwrap();
}
{
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();
storage.begin_transaction().unwrap();
storage.save("key1", vec![1, 2, 3]).unwrap();
storage.save("key2", vec![4, 5, 6]).unwrap();
storage.commit_transaction().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_transaction_rollback() {
let (_dir, path) = temp_db();
let mut storage = Sqlite::open(&path).unwrap();
storage.save("key1", vec![1, 2, 3]).unwrap();
storage.begin_transaction().unwrap();
storage.save("key1", vec![4, 5, 6]).unwrap();
storage.save("key2", vec![7, 8, 9]).unwrap();
storage.rollback_transaction().unwrap();
assert_eq!(storage.load("key1").unwrap(), Some(vec![1, 2, 3]));
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();
storage.begin_transaction().unwrap();
let result = storage.begin_transaction();
assert!(result.is_err());
assert!(result
.unwrap_err()
.to_string()
.contains("Transaction already in progress"));
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"));
}
}