Skip to main content

nestrs_workflows/
store.rs

1use async_trait::async_trait;
2use serde::{Deserialize, Serialize};
3use serde_json::Value;
4use std::collections::HashMap;
5use std::sync::Arc;
6use tokio::sync::Mutex;
7use uuid::Uuid;
8
9#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
10pub enum WorkflowStatus {
11    Running,
12    WaitingApproval,
13    Completed,
14    Failed,
15}
16
17#[derive(Debug, Clone, Serialize, Deserialize)]
18pub struct WorkflowInstance {
19    pub id: String,
20    pub name: String,
21    pub input: Value,
22    pub status: WorkflowStatus,
23    pub steps: HashMap<String, Value>,
24}
25
26#[async_trait]
27pub trait WorkflowStore: Send + Sync {
28    async fn create(&self, name: &str, input: Value) -> WorkflowInstance;
29    async fn get(&self, id: &str) -> Option<WorkflowInstance>;
30    async fn complete_step(
31        &self,
32        id: &str,
33        step: &str,
34        result: Value,
35    ) -> Result<WorkflowInstance, String>;
36    async fn set_status(
37        &self,
38        id: &str,
39        status: WorkflowStatus,
40    ) -> Result<WorkflowInstance, String>;
41}
42
43pub type SharedWorkflowStore = Arc<dyn WorkflowStore>;
44
45pub struct InMemoryWorkflowStore {
46    inner: Mutex<HashMap<String, WorkflowInstance>>,
47}
48
49impl InMemoryWorkflowStore {
50    pub fn shared() -> Arc<Self> {
51        Arc::new(Self {
52            inner: Mutex::new(HashMap::new()),
53        })
54    }
55}
56
57#[async_trait]
58impl WorkflowStore for InMemoryWorkflowStore {
59    async fn create(&self, name: &str, input: Value) -> WorkflowInstance {
60        let inst = WorkflowInstance {
61            id: Uuid::new_v4().to_string(),
62            name: name.to_string(),
63            input,
64            status: WorkflowStatus::Running,
65            steps: HashMap::new(),
66        };
67        self.inner
68            .lock()
69            .await
70            .insert(inst.id.clone(), inst.clone());
71        inst
72    }
73
74    async fn get(&self, id: &str) -> Option<WorkflowInstance> {
75        self.inner.lock().await.get(id).cloned()
76    }
77
78    async fn complete_step(
79        &self,
80        id: &str,
81        step: &str,
82        result: Value,
83    ) -> Result<WorkflowInstance, String> {
84        let mut guard = self.inner.lock().await;
85        let inst = guard
86            .get_mut(id)
87            .ok_or_else(|| "workflow not found".to_string())?;
88        inst.steps.insert(step.to_string(), result);
89        Ok(inst.clone())
90    }
91
92    async fn set_status(
93        &self,
94        id: &str,
95        status: WorkflowStatus,
96    ) -> Result<WorkflowInstance, String> {
97        let mut guard = self.inner.lock().await;
98        let inst = guard
99            .get_mut(id)
100            .ok_or_else(|| "workflow not found".to_string())?;
101        inst.status = status;
102        Ok(inst.clone())
103    }
104}