use crate::DbContext;
use anyhow::{Context, Result};
use rocksdb::IteratorMode;
use serde::{de::DeserializeOwned, Serialize};
fn serialize_to_bytes<T: Serialize>(value: &T) -> Result<Vec<u8>> {
serde_json::to_vec(value).context("Failed to serialize value")
}
fn deserialize_from_bytes<T: DeserializeOwned>(bytes: &[u8]) -> Result<T> {
serde_json::from_slice(bytes).context("Failed to deserialize value")
}
pub trait Keyable: serde::Serialize + serde::de::DeserializeOwned {
fn key(&self) -> String;
fn column_family() -> &'static str;
}
pub struct ColumnFamily<T: Keyable> {
_phantom: std::marker::PhantomData<T>,
}
impl<T: Keyable> Default for ColumnFamily<T> {
fn default() -> Self {
Self::new()
}
}
impl<T: Keyable> ColumnFamily<T> {
pub fn new() -> Self {
Self {
_phantom: std::marker::PhantomData,
}
}
pub fn get_all(&self) -> Result<Vec<T>> {
let cf_handle = DbContext::get_instance()
.db
.cf_handle(T::column_family())
.context(format!(
"Failed to get {} column family handle",
T::column_family()
))?;
let mut items = Vec::new();
let iter = DbContext::get_instance()
.db
.iterator_cf(&cf_handle, IteratorMode::Start);
for item in iter {
let (_key, value) = item.context("Failed to read database entry")?;
let item: T = deserialize_from_bytes(&value)?;
items.push(item);
}
Ok(items)
}
pub fn get(&self, key: &str) -> Result<Option<T>> {
let cf_handle = DbContext::get_instance()
.db
.cf_handle(T::column_family())
.context(format!(
"Failed to get {} column family handle",
T::column_family()
))?;
match DbContext::get_instance()
.db
.get_cf(&cf_handle, key)
.context("Failed to read database entry")?
{
Some(value) => {
let item: T = deserialize_from_bytes(&value)?;
Ok(Some(item))
}
None => Ok(None),
}
}
pub fn del(&self, key: &str) -> Result<()> {
let cf_handle = DbContext::get_instance()
.db
.cf_handle(T::column_family())
.context(format!(
"Failed to get {} column family handle",
T::column_family()
))?;
DbContext::get_instance()
.db
.delete_cf(&cf_handle, key)
.context("Failed to delete item")
}
pub fn set(&self, item: &T) -> Result<()> {
let cf_handle = DbContext::get_instance()
.db
.cf_handle(T::column_family())
.context(format!(
"Failed to get {} column family handle",
T::column_family()
))?;
let key = item.key();
let value = serialize_to_bytes(item)?;
DbContext::get_instance()
.db
.put_cf(&cf_handle, key, value)
.context("Failed to write item to database")
}
}