lykin/src/db.rs

137 lines
3.7 KiB
Rust

// Key-value store using sled.
use std::{
fmt,
fmt::{Display, Formatter},
path::Path,
};
use log::{debug, info};
use serde::{Deserialize, Serialize};
use sled::{Batch, Db, IVec, Result, Tree};
// The text and metadata of a Scuttlebutt root post.
#[derive(Debug, Deserialize, Serialize)]
pub struct Post {
pub key: String,
pub text: String,
pub date: String,
pub sequence: u64,
pub read: bool,
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct IVecString {
pub bytes: IVec,
pub string: String,
}
impl Display for IVecString {
fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
write!(f, "{}", self.string)
}
}
impl From<IVec> for IVecString {
fn from(bytes: IVec) -> Self {
Self {
string: String::from_utf8_lossy(&bytes).into_owned(),
bytes,
}
}
}
#[derive(Clone)]
pub struct Database {
/// Stores the sled database instance.
db: Db,
/// Stores the public keys of all the peers we are subscribed to.
peer_tree: Tree,
/// Stores the posts (content and metadata) for all the feeds we are subscribed to.
pub post_tree: Tree,
}
impl Database {
pub fn init(path: &Path) -> Self {
// Open the database at the given path.
// The database will be created if it does not yet exist.
// This code will panic if an IO error is encountered.
info!("initialising the sled database");
let db = sled::open(path).expect("failed to open database");
debug!("opening the 'peers' database tree");
let peer_tree = db
.open_tree("peers")
.expect("failed to open database peers tree");
debug!("opening the 'posts' database tree");
let post_tree = db
.open_tree("posts")
.expect("failed to open database posts tree");
Database {
db,
peer_tree,
post_tree,
}
}
pub fn add_peer(&self, public_key: &str) -> Result<Option<IVec>> {
self.peer_tree.insert(&public_key, vec![0])
}
pub fn remove_peer(&self, public_key: &str) -> Result<Option<IVec>> {
self.peer_tree.remove(&public_key)
}
pub fn get_peers(&self) -> Vec<String> {
self.peer_tree
.iter()
.keys()
.map(|bytes| IVecString::from(bytes.unwrap()))
.map(|ivec_string| ivec_string.string)
.collect()
}
pub fn insert_post(&self, public_key: &str, post: Post) -> Result<Option<IVec>> {
let post_key = format!("{}_{}", public_key, post.key);
let post_bytes = bincode::serialize(&post).unwrap();
self.post_tree.insert(post_key.as_bytes(), post_bytes)
}
pub fn insert_post_batch(&self, public_key: &str, posts: Vec<Post>) -> Result<()> {
let mut post_batch = Batch::default();
for post in posts {
let post_key = format!("{}_{}", public_key, post.key);
let post_bytes = bincode::serialize(&post).unwrap();
post_batch.insert(post_key.as_bytes(), post_bytes)
}
self.post_tree.apply_batch(post_batch)
}
pub fn get_posts(&self, public_key: &str) -> Result<Vec<Post>> {
let mut posts = Vec::new();
self.post_tree
.scan_prefix(public_key.as_bytes())
.map(|post| post.unwrap())
.for_each(|post| posts.push(bincode::deserialize(&post.1).unwrap()));
Ok(posts)
}
pub fn get_post(&self, public_key: &str, msg_id: &str) -> Result<Option<Post>> {
let post_key = format!("{}_{}", public_key, msg_id);
let post = self
.post_tree
.get(post_key.as_bytes())
.unwrap()
.map(|post| bincode::deserialize(&post).unwrap());
Ok(post)
}
}