use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::collections::HashMap;
use std::path::Path;
use std::sync::Arc;
use tokio::sync::RwLock;
use crate::Result;
#[derive(Clone)]
pub struct DataStore {
inner: Arc<RwLock<StoreInner>>,
persistence_path: Option<String>,
}
#[derive(Default, Serialize, Deserialize)]
struct StoreInner {
collections: HashMap<String, Vec<Value>>,
values: HashMap<String, Value>,
}
impl DataStore {
pub fn new() -> Self {
Self {
inner: Arc::new(RwLock::new(StoreInner::default())),
persistence_path: None,
}
}
pub fn with_persistence(path: impl Into<String>) -> Result<Self> {
let path = path.into();
let inner = if Path::new(&path).exists() {
let content = std::fs::read_to_string(&path)?;
serde_json::from_str(&content)?
} else {
StoreInner::default()
};
Ok(Self {
inner: Arc::new(RwLock::new(inner)),
persistence_path: Some(path),
})
}
pub async fn load_collection(&self, name: &str, path: impl AsRef<Path>) -> Result<()> {
let content = std::fs::read_to_string(path)?;
let data: Vec<Value> = serde_json::from_str(&content)?;
let mut inner = self.inner.write().await;
inner.collections.insert(name.to_string(), data);
Ok(())
}
pub async fn get_collection(&self, name: &str) -> Option<Vec<Value>> {
let inner = self.inner.read().await;
inner.collections.get(name).cloned()
}
pub async fn get_item(&self, collection: &str, index: usize) -> Option<Value> {
let inner = self.inner.read().await;
inner.collections.get(collection)?.get(index).cloned()
}
pub async fn find_by(&self, collection: &str, field: &str, value: &Value) -> Vec<Value> {
let inner = self.inner.read().await;
inner
.collections
.get(collection)
.map(|items| {
items
.iter()
.filter(|item| item.get(field).map(|v| v == value).unwrap_or(false))
.cloned()
.collect()
})
.unwrap_or_default()
}
pub async fn find_one_by(&self, collection: &str, field: &str, value: &Value) -> Option<Value> {
self.find_by(collection, field, value)
.await
.into_iter()
.next()
}
pub async fn create(&self, collection: &str, mut item: Value) -> Result<Value> {
let mut inner = self.inner.write().await;
let items = inner
.collections
.entry(collection.to_string())
.or_insert_with(Vec::new);
if item.get("id").is_none() {
let max_id = items
.iter()
.filter_map(|i| i.get("id").and_then(|v| v.as_u64()))
.max()
.unwrap_or(0);
let id = max_id + 1;
if let Value::Object(ref mut map) = item {
map.insert("id".to_string(), Value::Number(id.into()));
}
}
if let Value::Object(ref mut map) = item {
map.entry("created_at".to_string()).or_insert_with(|| {
Value::String(
chrono::Local::now()
.format("%Y-%m-%dT%H:%M:%S%.3f")
.to_string(),
)
});
}
items.push(item.clone());
drop(inner);
self.persist().await?;
Ok(item)
}
pub async fn update(
&self,
collection: &str,
id: &Value,
updates: Value,
) -> Result<Option<Value>> {
let mut inner = self.inner.write().await;
if let Some(items) = inner.collections.get_mut(collection) {
for item in items.iter_mut() {
if item.get("id") == Some(id) {
if let (Value::Object(existing), Value::Object(new)) = (item, &updates) {
for (key, value) in new {
existing.insert(key.clone(), value.clone());
}
}
let updated = items.iter().find(|i| i.get("id") == Some(id)).cloned();
drop(inner);
self.persist().await?;
return Ok(updated);
}
}
}
Ok(None)
}
pub async fn delete(&self, collection: &str, id: &Value) -> Result<bool> {
let mut inner = self.inner.write().await;
if let Some(items) = inner.collections.get_mut(collection) {
let original_len = items.len();
items.retain(|item| item.get("id") != Some(id));
let deleted = items.len() < original_len;
drop(inner);
if deleted {
self.persist().await?;
}
return Ok(deleted);
}
Ok(false)
}
pub async fn set(&self, key: &str, value: Value) -> Result<()> {
let mut inner = self.inner.write().await;
inner.values.insert(key.to_string(), value);
drop(inner);
self.persist().await
}
pub async fn get(&self, key: &str) -> Option<Value> {
let inner = self.inner.read().await;
inner.values.get(key).cloned()
}
pub async fn atomic_modify<F>(&self, key: &str, f: F) -> Result<Value>
where
F: FnOnce(Option<&Value>) -> Value,
{
let mut inner = self.inner.write().await;
let new_value = f(inner.values.get(key));
inner.values.insert(key.to_string(), new_value.clone());
drop(inner);
self.persist().await?;
Ok(new_value)
}
pub async fn remove(&self, key: &str) -> Result<Option<Value>> {
let mut inner = self.inner.write().await;
let removed = inner.values.remove(key);
drop(inner);
self.persist().await?;
Ok(removed)
}
pub async fn as_context(&self) -> HashMap<String, Value> {
let inner = self.inner.read().await;
let mut context = HashMap::new();
for (name, items) in &inner.collections {
context.insert(name.clone(), Value::Array(items.clone()));
}
for (key, value) in &inner.values {
context.insert(key.clone(), value.clone());
}
context
}
pub async fn set_collection(&self, name: &str, items: Vec<Value>) -> Result<()> {
let mut inner = self.inner.write().await;
inner.collections.insert(name.to_string(), items);
drop(inner);
self.persist().await
}
async fn persist(&self) -> Result<()> {
if let Some(ref path) = self.persistence_path {
let inner = self.inner.read().await;
let content = serde_json::to_string_pretty(&*inner)?;
let tmp_path = format!("{}.tmp", path);
tokio::fs::write(&tmp_path, &content).await?;
tokio::fs::rename(&tmp_path, path).await?;
}
Ok(())
}
}
impl Default for DataStore {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[tokio::test]
async fn test_create_and_get() {
let store = DataStore::new();
let post = json!({
"title": "Hello World",
"content": "My first post"
});
let created = store.create("posts", post).await.unwrap();
assert_eq!(created.get("id"), Some(&json!(1)));
let posts = store.get_collection("posts").await.unwrap();
assert_eq!(posts.len(), 1);
}
#[tokio::test]
async fn test_find_by() {
let store = DataStore::new();
store
.create("users", json!({"name": "Alice", "role": "admin"}))
.await
.unwrap();
store
.create("users", json!({"name": "Bob", "role": "user"}))
.await
.unwrap();
store
.create("users", json!({"name": "Charlie", "role": "admin"}))
.await
.unwrap();
let admins = store.find_by("users", "role", &json!("admin")).await;
assert_eq!(admins.len(), 2);
}
#[tokio::test]
async fn test_update() {
let store = DataStore::new();
store
.create("posts", json!({"title": "Draft"}))
.await
.unwrap();
store
.update(
"posts",
&json!(1),
json!({"title": "Published", "status": "live"}),
)
.await
.unwrap();
let post = store.find_one_by("posts", "id", &json!(1)).await.unwrap();
assert_eq!(post.get("title"), Some(&json!("Published")));
assert_eq!(post.get("status"), Some(&json!("live")));
}
#[tokio::test]
async fn test_delete() {
let store = DataStore::new();
store
.create("posts", json!({"title": "To Delete"}))
.await
.unwrap();
let deleted = store.delete("posts", &json!(1)).await.unwrap();
assert!(deleted);
let posts = store.get_collection("posts").await.unwrap();
assert!(posts.is_empty());
}
}