stemedb/crates/stemedb-storage/src/sled_backend.rs

use crate::error::{Result, StorageError};
use crate::traits::KVStore;
use async_trait::async_trait;
use sled::Db;
use std::path::Path;

/// Sled-based implementation of the KVStore trait.
#[derive(Debug, Clone)]
pub struct SledStore {
    db: Db,
}

impl SledStore {
    /// Open or create a new Sled database at the given path.
    pub fn open(path: impl AsRef<Path>) -> Result<Self> {
        let db = sled::open(path).map_err(StorageError::Sled)?;
        Ok(Self { db })
    }

    /// Open a temporary Sled database for testing.
    ///
    /// The database will be automatically deleted when dropped.
    /// Useful for unit tests in this and other crates.
    pub fn open_temp() -> Result<Self> {
        let config = sled::Config::new().temporary(true);
        let db = config.open().map_err(StorageError::Sled)?;
        Ok(Self { db })
    }
}

#[async_trait]
impl KVStore for SledStore {
    async fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>> {
        let result = self.db.get(key).map_err(StorageError::Sled)?;
        Ok(result.map(|ivec| ivec.to_vec()))
    }

    async fn put(&self, key: &[u8], value: &[u8]) -> Result<()> {
        self.db.insert(key, value).map_err(StorageError::Sled)?;
        Ok(())
    }

    async fn delete(&self, key: &[u8]) -> Result<()> {
        self.db.remove(key).map_err(StorageError::Sled)?;
        Ok(())
    }

    async fn scan_prefix(&self, prefix: &[u8]) -> Result<Vec<(Vec<u8>, Vec<u8>)>> {
        let iter = self.db.scan_prefix(prefix);
        let mut results = Vec::new();
        for item in iter {
            let (k, v) = item.map_err(StorageError::Sled)?;
            results.push((k.to_vec(), v.to_vec()));
        }
        Ok(results)
    }

    async fn flush(&self) -> Result<()> {
        self.db.flush_async().await.map_err(StorageError::Sled)?;
        Ok(())
    }

    async fn fetch_and_add_u64(&self, key: &[u8], delta: u64) -> Result<u64> {
        let result = self
            .db
            .update_and_fetch(key, |old| {
                let current = match old {
                    Some(bytes) => match <[u8; 8]>::try_from(bytes) {
                        Ok(arr) => u64::from_le_bytes(arr),
                        Err(_) => 0, // Corrupted data, start fresh
                    },
                    None => 0, // Key doesn't exist, start at 0
                };
                Some(current.saturating_add(delta).to_le_bytes().to_vec())
            })
            .map_err(StorageError::Sled)?;

        // Result is Some because our update_fn always returns Some
        let bytes = result.ok_or_else(|| {
            StorageError::Serialization("fetch_and_add_u64 returned None unexpectedly".to_string())
        })?;
        let arr: [u8; 8] = bytes.as_ref().try_into().map_err(|_| {
            StorageError::Serialization("fetch_and_add_u64 returned wrong size".to_string())
        })?;
        Ok(u64::from_le_bytes(arr))
    }

    async fn compare_and_swap_f32<F>(&self, key: &[u8], update_fn: F) -> Result<f32>
    where
        F: Fn(f32) -> f32 + Send + Sync,
    {
        let result = self
            .db
            .update_and_fetch(key, |old| {
                let current = match old {
                    Some(bytes) => match <[u8; 4]>::try_from(bytes) {
                        Ok(arr) => f32::from_le_bytes(arr),
                        Err(_) => 0.0, // Corrupted data, start fresh
                    },
                    None => 0.0, // Key doesn't exist, start at 0.0
                };
                let new_value = update_fn(current);
                Some(new_value.to_le_bytes().to_vec())
            })
            .map_err(StorageError::Sled)?;

        let bytes = result.ok_or_else(|| {
            StorageError::Serialization(
                "compare_and_swap_f32 returned None unexpectedly".to_string(),
            )
        })?;
        let arr: [u8; 4] = bytes.as_ref().try_into().map_err(|_| {
            StorageError::Serialization("compare_and_swap_f32 returned wrong size".to_string())
        })?;
        Ok(f32::from_le_bytes(arr))
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_sled_store_roundtrip() {
        let store = SledStore::open_temp().expect("Failed to create temp DB");
        let key = b"test_key";
        let value = b"test_value";

        // Put
        store.put(key, value).await.expect("Put failed");

        // Get
        let retrieved = store.get(key).await.expect("Get failed");
        assert_eq!(retrieved, Some(value.to_vec()));

        // Delete
        store.delete(key).await.expect("Delete failed");

        // Get after delete
        let deleted = store.get(key).await.expect("Get failed");
        assert_eq!(deleted, None);
    }

    #[tokio::test]
    async fn test_scan_prefix() {
        let store = SledStore::open_temp().expect("Failed to create temp DB");
        store.put(b"prefix:1", b"val1").await.unwrap();
        store.put(b"prefix:2", b"val2").await.unwrap();
        store.put(b"other:3", b"val3").await.unwrap();

        let results = store.scan_prefix(b"prefix:").await.unwrap();
        assert_eq!(results.len(), 2);
        assert_eq!(results[0], (b"prefix:1".to_vec(), b"val1".to_vec()));
        assert_eq!(results[1], (b"prefix:2".to_vec(), b"val2".to_vec()));
    }
}