Add rocksdb

src/helpers/db.rs

@@ -0,0 +1,90 @@
+use crate::constants::{APP_DB_DATA_DIR, DB_COLUMNS};
+use rocksdb::{ColumnFamilyDescriptor, IteratorMode, OptimisticTransactionDB, Options};
+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 values = self.db.iterator_cf(&cf, IteratorMode::Start)
+            .filter_map(|res| res.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(())
+    }
+}