rust_viewflow/db/
memory.rs1use std::collections::HashMap;
2use std::sync::RwLock;
3
4use crate::core::{TaskState, WorkflowResult, WorkflowState};
5use crate::db::WorkflowDatabase;
6
7pub struct MemoryDatabase {
8 workflows: RwLock<HashMap<String, WorkflowState>>,
9 tasks: RwLock<HashMap<String, TaskState>>,
10}
11
12impl Default for MemoryDatabase {
13 fn default() -> Self {
14 Self::new()
15 }
16}
17
18impl MemoryDatabase {
19 pub fn new() -> Self {
20 Self {
21 workflows: RwLock::new(HashMap::new()),
22 tasks: RwLock::new(HashMap::new()),
23 }
24 }
25}
26
27#[async_trait::async_trait]
28impl WorkflowDatabase for MemoryDatabase {
29 async fn create_workflow(&self, workflow: &WorkflowState) -> WorkflowResult<()> {
30 self.workflows
31 .write()
32 .expect("memory workflow lock poisoned")
33 .insert(workflow.id.clone(), workflow.clone());
34 Ok(())
35 }
36
37 async fn get_workflow(&self, workflow_id: &str) -> WorkflowResult<WorkflowState> {
38 self.workflows
39 .read()
40 .expect("memory workflow lock poisoned")
41 .get(workflow_id)
42 .cloned()
43 .ok_or_else(|| format!("Workflow '{}' not found", workflow_id).into())
44 }
45
46 async fn update_workflow(&self, workflow: &WorkflowState) -> WorkflowResult<()> {
47 self.workflows
48 .write()
49 .expect("memory workflow lock poisoned")
50 .insert(workflow.id.clone(), workflow.clone());
51 Ok(())
52 }
53
54 async fn create_task(&self, task: &TaskState) -> WorkflowResult<()> {
55 self.tasks
56 .write()
57 .expect("memory task lock poisoned")
58 .insert(task.id.clone(), task.clone());
59 Ok(())
60 }
61
62 async fn get_task(&self, task_id: &str) -> WorkflowResult<TaskState> {
63 self.tasks
64 .read()
65 .expect("memory task lock poisoned")
66 .get(task_id)
67 .cloned()
68 .ok_or_else(|| format!("Task '{}' not found", task_id).into())
69 }
70
71 async fn update_task(&self, task: &TaskState) -> WorkflowResult<()> {
72 self.tasks
73 .write()
74 .expect("memory task lock poisoned")
75 .insert(task.id.clone(), task.clone());
76 Ok(())
77 }
78
79 async fn get_tasks_by_workflow(&self, workflow_id: &str) -> WorkflowResult<Vec<TaskState>> {
80 let tasks = self
81 .tasks
82 .read()
83 .expect("memory task lock poisoned")
84 .values()
85 .filter(|task| task.workflow_id == workflow_id)
86 .cloned()
87 .collect();
88 Ok(tasks)
89 }
90}