use std::path::Path; use ::diesel::{query_dsl::*, ExpressionMethods}; use anyhow::anyhow; use diesel::{Connection, SqliteConnection}; use lazy_static::lazy_static; use parking_lot::RwLock; use serde::de::DeserializeOwned; use serde::Serialize; use crate::kv::schema::{kv_table, kv_table::dsl, KV_SQL}; use crate::sqlite::{DBConnection, Database, PoolConfig}; const DB_NAME: &str = "kv.db"; lazy_static! { static ref KV_HOLDER: RwLock = RwLock::new(KV::new()); } /// [KV] uses a sqlite database to store key value pairs. /// Most of the time, it used to storage AppFlowy configuration. pub struct KV { database: Option, } impl KV { fn new() -> Self { KV { database: None } } #[tracing::instrument(level = "trace", err)] pub fn init(root: &str) -> Result<(), anyhow::Error> { if !Path::new(root).exists() { return Err(anyhow!("Init KV failed. {} not exists", root)); } let pool_config = PoolConfig::default(); let database = Database::new(root, DB_NAME, pool_config).unwrap(); let conn = database.get_connection().unwrap(); SqliteConnection::execute(&*conn, KV_SQL).unwrap(); tracing::trace!("Init kv with path: {}", root); KV_HOLDER.write().database = Some(database); Ok(()) } /// Set a string value of a key pub fn set_str(key: &str, value: T) { let _ = Self::set_key_value(key, Some(value.to_string())); } /// Set a bool value of a key pub fn set_bool(key: &str, value: bool) -> Result<(), anyhow::Error> { Self::set_key_value(key, Some(value.to_string())) } /// Set a object that implements [Serialize] trait of a key pub fn set_object(key: &str, value: T) -> Result<(), anyhow::Error> { let value = serde_json::to_string(&value)?; Self::set_key_value(key, Some(value))?; Ok(()) } /// Set a i64 value of a key pub fn set_i64(key: &str, value: i64) -> Result<(), anyhow::Error> { Self::set_key_value(key, Some(value.to_string())) } /// Get a string value of a key pub fn get_str(key: &str) -> Option { Self::get_key_value(key).and_then(|kv| kv.value) } /// Get a bool value of a key pub fn get_bool(key: &str) -> bool { Self::get_key_value(key) .and_then(|kv| kv.value) .and_then(|v| v.parse::().ok()) .unwrap_or(false) } /// Get a i64 value of a key pub fn get_i64(key: &str) -> Option { Self::get_key_value(key) .and_then(|kv| kv.value) .and_then(|v| v.parse::().ok()) } /// Get a object that implements [DeserializeOwned] trait of a key pub fn get_object(key: &str) -> Option { Self::get_str(key).and_then(|v| serde_json::from_str(&v).ok()) } #[allow(dead_code)] pub fn remove(key: &str) { if let Ok(conn) = get_connection() { let sql = dsl::kv_table.filter(kv_table::key.eq(key)); let _ = diesel::delete(sql).execute(&*conn); } } fn set_key_value(key: &str, value: Option) -> Result<(), anyhow::Error> { let conn = get_connection()?; diesel::replace_into(kv_table::table) .values(KeyValue { key: key.to_string(), value, }) .execute(&*conn)?; Ok(()) } fn get_key_value(key: &str) -> Option { let conn = get_connection().ok()?; dsl::kv_table .filter(kv_table::key.eq(key)) .first::(&*conn) .ok() } } fn get_connection() -> Result { let conn = KV_HOLDER .read() .database .as_ref() .expect("KVStore is not init") .get_connection() .map_err(|_e| anyhow!("Get KV connection error"))?; Ok(conn) } #[derive(Clone, Debug, Default, Queryable, Identifiable, Insertable, AsChangeset)] #[table_name = "kv_table"] #[primary_key(key)] pub struct KeyValue { pub key: String, pub value: Option, } #[cfg(test)] mod tests { use serde::{Deserialize, Serialize}; use tempfile::TempDir; use crate::kv::KV; #[derive(Serialize, Deserialize, Clone, Eq, PartialEq, Debug)] struct Person { name: String, age: i32, } #[test] fn kv_store_test() { let tempdir = TempDir::new().unwrap(); let path = tempdir.into_path(); KV::init(path.to_str().unwrap()).unwrap(); KV::set_str("1", "hello".to_string()); assert_eq!(KV::get_str("1").unwrap(), "hello"); assert_eq!(KV::get_str("2"), None); KV::set_bool("1", true).unwrap(); assert!(KV::get_bool("1")); assert!(!KV::get_bool("2")); KV::set_i64("1", 1).unwrap(); assert_eq!(KV::get_i64("1").unwrap(), 1); assert_eq!(KV::get_i64("2"), None); let person = Person { name: "nathan".to_string(), age: 30, }; KV::set_object("1", person.clone()).unwrap(); assert_eq!(KV::get_object::("1").unwrap(), person); } }