Skip to main content

rust_viewflow/db/
memory.rs

1use 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}