92 lines
3.3 KiB
Rust
92 lines
3.3 KiB
Rust
use crate::constants::{APP_DB_DATA_DIR, DB_COLUMNS};
|
|
use rocksdb::{ColumnFamilyDescriptor, IteratorMode, OptimisticTransactionDB, Options, ReadOptions};
|
|
use serde::{Serialize};
|
|
use serde::de::DeserializeOwned;
|
|
use crate::models::RocksColumn;
|
|
|
|
pub struct RocksDB {
|
|
db: OptimisticTransactionDB,
|
|
}
|
|
|
|
impl RocksDB {
|
|
pub fn new(db_opts: Options, cf_opts: Options) -> color_eyre::Result<Self> {
|
|
let cfs = DB_COLUMNS.iter()
|
|
.map(|cf| ColumnFamilyDescriptor::new(cf.to_string(), cf_opts.clone()))
|
|
.collect::<Vec<_>>();
|
|
let db = OptimisticTransactionDB::open_cf_descriptors(
|
|
&db_opts,
|
|
APP_DB_DATA_DIR.as_path(),
|
|
cfs
|
|
)?;
|
|
let rocks = Self {
|
|
db
|
|
};
|
|
Ok(rocks)
|
|
}
|
|
|
|
pub fn get_value<TValue, TColumn>(&self, id: TColumn::Id) -> color_eyre::Result<Option<TValue>>
|
|
where TColumn: RocksColumn, TValue: DeserializeOwned
|
|
{
|
|
let cf = self.db.cf_handle(TColumn::get_column_name().as_str()).unwrap();
|
|
let query_res = self.db.get_cf(&cf, serde_json::to_string(&id)?)?;
|
|
if query_res.is_none() {
|
|
return Ok(None);
|
|
}
|
|
Ok(Some(serde_json::from_slice(&query_res.unwrap())?))
|
|
}
|
|
|
|
pub fn set_value<TColumn>(&self, value: &TColumn) -> color_eyre::Result<()>
|
|
where TColumn: RocksColumn + Serialize
|
|
{
|
|
let cf = self.db.cf_handle(TColumn::get_column_name().as_str()).unwrap();
|
|
self.db.put_cf(&cf, serde_json::to_string(&value.get_id())?, serde_json::to_string(value)?)?;
|
|
Ok(())
|
|
}
|
|
|
|
pub fn get_values<TColumn>(&self, ids: &[TColumn::Id]) -> color_eyre::Result<Vec<TColumn>>
|
|
where TColumn: RocksColumn + DeserializeOwned
|
|
{
|
|
let transaction = self.db.transaction();
|
|
let cf = self.db.cf_handle(TColumn::get_column_name().as_str()).unwrap();
|
|
let mut values = Vec::new();
|
|
for id in ids {
|
|
let query_res = transaction.get_cf(&cf, serde_json::to_string(id)?)?;
|
|
if let Some(res) = query_res {
|
|
let value = serde_json::from_slice(&res)?;
|
|
values.push(value);
|
|
}
|
|
}
|
|
Ok(values)
|
|
}
|
|
|
|
pub fn get_all_values<TColumn>(&self) -> color_eyre::Result<Vec<(TColumn::Id, TColumn)>>
|
|
where TColumn: RocksColumn + DeserializeOwned
|
|
{
|
|
let cf = self.db.cf_handle(TColumn::get_column_name().as_str()).unwrap();
|
|
let mut options = ReadOptions::default();
|
|
options.set_async_io(true);
|
|
let values = self.db.iterator_cf_opt(&cf, options, IteratorMode::Start)
|
|
.filter_map(Result::ok)
|
|
.map(|(k, v)|
|
|
(
|
|
serde_json::from_slice::<TColumn::Id>(&k).unwrap(),
|
|
serde_json::from_slice::<TColumn>(&v).unwrap()
|
|
)
|
|
)
|
|
.collect::<Vec<_>>();
|
|
Ok(values)
|
|
}
|
|
|
|
|
|
pub fn set_values<TColumn>(&self, values: &[TColumn]) -> color_eyre::Result<()>
|
|
where TColumn: RocksColumn + Serialize
|
|
{
|
|
let transaction = self.db.transaction();
|
|
let cf = self.db.cf_handle(TColumn::get_column_name().as_str()).unwrap();
|
|
for value in values {
|
|
transaction.put_cf(&cf, serde_json::to_string(&value.get_id())?, serde_json::to_string(value)?)?;
|
|
}
|
|
transaction.commit()?;
|
|
Ok(())
|
|
}
|
|
} |