prost_sled/

lib.rs

//! prost-sled: An integration layer between [prost] and [sled].
//!
//! prost-sled makes it easy to store protobufs in a sled database because
//! it abstracts away the boilerplate of encoding and decoding the protobufs.
//! If for any reason you wish to interact with the raw bytes instead or
//! [sled::Db] implements a method that [ProtoDb] doesn't yet, you can simply
//! use the `from` and `into` methods of the corresponding types as [From] and
//! [Into] are implemented as a go between, between the two types.

use std::path::Path;

use prost::{bytes::BytesMut, Message};
use thiserror::Error;

/// Errors that can be returned by this library. It's really a simple
/// integration layer to encompass the possible errors returned by [sled] and
/// [prost].
#[derive(Debug, Error, PartialEq)]
pub enum Error {
    /// An error was returned by [sled].
    #[error(transparent)]
    SledError(#[from] sled::Error),
    /// A decoding error ([prost::DecodeError]) occurred in [prost].
    #[error(transparent)]
    ProstDecodeError(#[from] prost::DecodeError),
    /// An encoding error ([prost::EncodeError]) occurred in [prost].
    #[error(transparent)]
    ProstEncodeError(#[from] prost::EncodeError),
}

/// Result of a database action. That is, either some type `T` or an
/// [enum@Error].
pub type Result<T> = std::result::Result<T, Error>;

/// Wrapper around [sled::Db] that allows you to use types implementing
/// [prost::Message] instead of raw bytes.
#[derive(Clone)]
pub struct ProtoDb(sled::Db);

/// Escape hatch to get from a [ProtoDb] to a [sled::Db] in-case you need
/// something more low level.
impl From<ProtoDb> for sled::Db {
    fn from(db: ProtoDb) -> Self {
        db.0
    }
}

/// Convenience implementation to convert an existing [sled::Db] to a [ProtoDb].
impl From<sled::Db> for ProtoDb {
    fn from(db: sled::Db) -> Self {
        Self(db)
    }
}

/// Wrapper around [sled::Tree] that allows you to use types implementing [prost::Message] instead
/// of raw bytes.
#[derive(Clone)]
pub struct ProtoTree(sled::Tree);

/// Escape hatch to get from a [ProtoTree] to a [sled::Tree] in-case you need
/// something more low level.
impl From<ProtoTree> for sled::Tree {
    fn from(t: ProtoTree) -> Self {
        t.0
    }
}

/// Convenience implementation to convert an existing [sled::Db] to a [ProtoDb].
impl From<sled::Tree> for ProtoTree {
    fn from(t: sled::Tree) -> Self {
        Self(t)
    }
}

/// Create a [ProtoDb] and return it.
pub fn open<P: AsRef<Path>>(path: P) -> Result<ProtoDb> {
    let db = sled::open(path)?;
    Ok(db.into())
}

impl ProtoDb {
    pub fn open_tree<V: AsRef<[u8]>>(&self, name: V) -> Result<ProtoTree> {
        Ok(self.0.open_tree(name)?.into())
    }

    // Ewww. These methods are exactly the same as those in ProtoTree. I'm not sure of a better way
    // though. sled implements Db as a deref of Tree but we can't also implement it that way.
    pub fn contains_key<K>(&self, key: K) -> Result<bool>
    where
        K: AsRef<[u8]>,
    {
        let exists = self.0.contains_key(key)?;
        Ok(exists)
    }

    /// Get a value by its key.
    pub fn get<K, T>(&self, key: K) -> Result<Option<T>>
    where
        K: AsRef<[u8]>,
        T: Message + Default,
    {
        let maybe_data = self.0.get(key)?;
        if let Some(data) = maybe_data {
            let msg = T::decode(&*data)?;
            Ok(Some(msg))
        } else {
            Ok(None)
        }
    }

    /// Atomically retrieve then update a value.
    pub fn update_and_fetch<K, V, F, T>(&self, key: K, mut f: F) -> Result<Option<T>>
    where
        K: AsRef<[u8]>,
        F: FnMut(Option<T>) -> Option<V>,
        V: Into<T>,
        T: Message + Default,
    {
        // Escape the scoping of the closure so we can report the error.
        let mut err: Option<Error> = None;
        let maybe_data = self.0.update_and_fetch(key, |maybe_data| {
            let maybe_msg = if let Some(data) = maybe_data {
                match T::decode(data) {
                    Ok(value) => Some(value),
                    Err(e) => {
                        err = Some(e.into());
                        None
                    }
                }
            } else {
                None
            };
            if let Some(inserted) = f(maybe_msg) {
                let mut buf = BytesMut::default();
                let inserted_msg: T = inserted.into();
                if let Err(e) = inserted_msg.encode(&mut buf) {
                    err = Some(e.into());
                    None
                } else {
                    Some(buf.as_bytes())
                }
            } else {
                None
            }
        })?;

        if let Some(e) = err {
            return Err(e);
        }

        if let Some(data) = maybe_data {
            let msg = T::decode(&*data)?;
            Ok(Some(msg))
        } else {
            Ok(None)
        }
    }

    /// Insert a value into the database.
    pub fn insert<K, V, T>(&self, key: K, value: V) -> Result<Option<T>>
    where
        K: AsRef<[u8]>,
        V: Into<T>,
        T: Message + Default,
    {
        let mut buf = BytesMut::default();
        let msg: T = value.into();
        msg.encode(&mut buf)?;
        let maybe_inserted = self.0.insert(key, buf.as_bytes())?;
        if let Some(inserted) = maybe_inserted {
            let msg = T::decode(&*inserted)?;
            Ok(Some(msg))
        } else {
            Ok(None)
        }
    }
}

impl ProtoTree {
    pub fn contains_key<K>(&self, key: K) -> Result<bool>
    where
        K: AsRef<[u8]>,
    {
        let exists = self.0.contains_key(key)?;
        Ok(exists)
    }

    /// Get a value by its key.
    pub fn get<K, T>(&self, key: K) -> Result<Option<T>>
    where
        K: AsRef<[u8]>,
        T: Message + Default,
    {
        let maybe_data = self.0.get(key)?;
        if let Some(data) = maybe_data {
            let msg = T::decode(&*data)?;
            Ok(Some(msg))
        } else {
            Ok(None)
        }
    }

    /// Atomically retrieve then update a value.
    pub fn update_and_fetch<K, V, F, T>(&self, key: K, mut f: F) -> Result<Option<T>>
    where
        K: AsRef<[u8]>,
        F: FnMut(Option<T>) -> Option<V>,
        V: Into<T>,
        T: Message + Default,
    {
        // Escape the scoping of the closure so we can report the error.
        let mut err: Option<Error> = None;
        let maybe_data = self.0.update_and_fetch(key, |maybe_data| {
            let maybe_msg = if let Some(data) = maybe_data {
                match T::decode(data) {
                    Ok(value) => Some(value),
                    Err(e) => {
                        err = Some(e.into());
                        None
                    }
                }
            } else {
                None
            };
            if let Some(inserted) = f(maybe_msg) {
                let mut buf = BytesMut::default();
                let inserted_msg: T = inserted.into();
                if let Err(e) = inserted_msg.encode(&mut buf) {
                    err = Some(e.into());
                    None
                } else {
                    Some(buf.as_bytes())
                }
            } else {
                None
            }
        })?;

        if let Some(e) = err {
            return Err(e);
        }

        if let Some(data) = maybe_data {
            let msg = T::decode(&*data)?;
            Ok(Some(msg))
        } else {
            Ok(None)
        }
    }

    /// Insert a value into the database.
    pub fn insert<K, V, T>(&self, key: K, value: V) -> Result<Option<T>>
    where
        K: AsRef<[u8]>,
        V: Into<T>,
        T: Message + Default,
    {
        let mut buf = BytesMut::default();
        let msg: T = value.into();
        msg.encode(&mut buf)?;
        let maybe_inserted = self.0.insert(key, buf.as_bytes())?;
        if let Some(inserted) = maybe_inserted {
            let msg = T::decode(&*inserted)?;
            Ok(Some(msg))
        } else {
            Ok(None)
        }
    }
}

/// Convenience trait to convert to stdlib bytes. This is intended only for
/// Convenience trait to convert to stdlib bytes. This is intended only for
/// internal use (hence not `pub`) and is only implemented for `BytesMut`.
trait BytesMutAsBytes {
    fn as_bytes(&self) -> Vec<u8>;
}

/// Conversion between `BytesMut` and `Vec<u8>`. This is just a convenience as
/// `sled` works with `Vec<u8>` and `prost` uses `BytesMut`.
impl BytesMutAsBytes for BytesMut {
    fn as_bytes(&self) -> Vec<u8> {
        let bytes: &[u8] = &self;
        bytes.to_owned()
    }
}

#[cfg(test)]
mod test {
    use once_cell::sync::OnceCell;
    use rand::{distributions::Alphanumeric, thread_rng, Rng};

    use super::{ProtoDb, ProtoTree};

    mod messages {
        include!(concat!(env!("OUT_DIR"), "/messages.rs"));
    }

    fn random_key() -> String {
        thread_rng()
            .sample_iter(&Alphanumeric)
            .take(5)
            .map(char::from)
            .collect()
    }

    fn proto_db() -> ProtoDb {
        static DB: OnceCell<sled::Db> = OnceCell::new();
        let db = DB.get_or_init(|| {
            sled::open("/tmp/prost-sled-test").expect("failed to open sled DB for test")
        });
        ProtoDb(db.to_owned())
    }

    fn proto_tree() -> ProtoTree {
        static TREE: OnceCell<sled::Tree> = OnceCell::new();
        let db = proto_db();
        let tree = TREE.get_or_init(|| db.0.open_tree("test").expect("failed to create tree"));
        ProtoTree(tree.to_owned())
    }

    #[test]
    fn get_exists() {
        let db = proto_db();
        let thing = messages::Thing::default();
        let key = random_key();
        let _: Option<messages::Thing> = db.insert(&key, thing.clone()).unwrap();
        let retrieved = db.get(&key).unwrap().unwrap();
        assert_eq!(thing, retrieved);
    }

    #[test]
    fn get_no_exist() {
        let db = proto_db();
        let key = random_key();
        let retrieved: Option<messages::Thing> = db.get(&key).unwrap();
        assert!(retrieved.is_none());
    }

    #[test]
    fn update_and_fetch_existing() {
        let db = proto_db();
        let thing = messages::Thing::default();
        let key = random_key();
        let _: Option<messages::Thing> = db.insert(&key, thing).unwrap();
        let updated = db
            .update_and_fetch(&key, |maybe_msg| {
                let mut msg: messages::Thing = maybe_msg.unwrap();
                msg.i = "test".into();
                Some(msg)
            })
            .unwrap();
        let expected = messages::Thing {
            i: "test".into(),
            ..Default::default()
        };
        assert_eq!(updated, Some(expected));
    }

    #[test]
    fn update_and_fetch_no_existing() {
        let db = proto_db();
        let key = random_key();
        let updated = db
            .update_and_fetch(&key, |maybe_msg| {
                assert!(maybe_msg.is_none());
                Some(messages::Thing {
                    i: "test".into(),
                    ..Default::default()
                })
            })
            .unwrap();
        let expected = messages::Thing {
            i: "test".into(),
            ..Default::default()
        };
        assert_eq!(updated, Some(expected));
    }

    #[test]
    fn insert_not_existing() {
        let db = proto_db();
        let key = random_key();
        let thing = messages::Thing::default();
        let previous: Option<messages::Thing> = db.insert(&key, thing.clone()).unwrap();
        assert!(previous.is_none());
        let inserted: messages::Thing = db.get(&key).unwrap().unwrap();
        assert_eq!(inserted, thing);
    }

    #[test]
    fn insert_existing() {
        let db = proto_db();
        let key = random_key();
        let first_thing = messages::Thing::default();
        let previous: Option<messages::Thing> = db.insert(&key, first_thing.clone()).unwrap();
        assert!(previous.is_none());
        let inserted: messages::Thing = db.get(&key).unwrap().unwrap();
        assert_eq!(inserted, first_thing);

        let second_thing = messages::Thing {
            i: random_key(),
            ..Default::default()
        };
        let second_inserted: messages::Thing =
            db.insert(&key, second_thing.clone()).unwrap().unwrap();
        assert_eq!(first_thing, second_inserted);
        let retieved: messages::Thing = db.get(&key).unwrap().unwrap();
        assert_eq!(retieved, second_thing);
    }

    #[test]
    fn tree_get_exists() {
        let tree = proto_tree();
        let thing = messages::Thing::default();
        let key = random_key();
        let _: Option<messages::Thing> = tree.insert(&key, thing.clone()).unwrap();
        let retrieved = tree.get(&key).unwrap().unwrap();
        assert_eq!(thing, retrieved);
    }

    #[test]
    fn tree_get_no_exist() {
        let tree = proto_tree();
        let key = random_key();
        let retrieved: Option<messages::Thing> = tree.get(&key).unwrap();
        assert!(retrieved.is_none());
    }

    #[test]
    fn tree_update_and_fetch_existing() {
        let tree = proto_tree();
        let thing = messages::Thing::default();
        let key = random_key();
        let _: Option<messages::Thing> = tree.insert(&key, thing).unwrap();
        let updated = tree
            .update_and_fetch(&key, |maybe_msg| {
                let mut msg: messages::Thing = maybe_msg.unwrap();
                msg.i = "test".into();
                Some(msg)
            })
            .unwrap();
        let expected = messages::Thing {
            i: "test".into(),
            ..Default::default()
        };
        assert_eq!(updated, Some(expected));
    }

    #[test]
    fn tree_update_and_fetch_no_existing() {
        let tree = proto_tree();
        let key = random_key();
        let updated = tree
            .update_and_fetch(&key, |maybe_msg| {
                assert!(maybe_msg.is_none());
                Some(messages::Thing {
                    i: "test".into(),
                    ..Default::default()
                })
            })
            .unwrap();
        let expected = messages::Thing {
            i: "test".into(),
            ..Default::default()
        };
        assert_eq!(updated, Some(expected));
    }

    #[test]
    fn tree_insert_not_existing() {
        let tree = proto_tree();
        let key = random_key();
        let thing = messages::Thing::default();
        let previous: Option<messages::Thing> = tree.insert(&key, thing.clone()).unwrap();
        assert!(previous.is_none());
        let inserted: messages::Thing = tree.get(&key).unwrap().unwrap();
        assert_eq!(inserted, thing);
    }

    #[test]
    fn tree_insert_existing() {
        let tree = proto_tree();
        let key = random_key();
        let first_thing = messages::Thing::default();
        let previous: Option<messages::Thing> = tree.insert(&key, first_thing.clone()).unwrap();
        assert!(previous.is_none());
        let inserted: messages::Thing = tree.get(&key).unwrap().unwrap();
        assert_eq!(inserted, first_thing);

        let second_thing = messages::Thing {
            i: random_key(),
            ..Default::default()
        };
        let second_inserted: messages::Thing =
            tree.insert(&key, second_thing.clone()).unwrap().unwrap();
        assert_eq!(first_thing, second_inserted);
        let retieved: messages::Thing = tree.get(&key).unwrap().unwrap();
        assert_eq!(retieved, second_thing);
    }
}