Skip to main content

ironflow_store/memory/
mod.rs

1//! In-memory [`Store`](crate::store::Store) implementation for development and testing.
2//!
3//! [`InMemoryStore`] uses `Arc<RwLock<..>>` internally, making it safe to share
4//! across tasks. Data is lost when the process exits.
5//!
6//! # Examples
7//!
8//! ```no_run
9//! use std::collections::HashMap;
10//! use ironflow_store::prelude::*;
11//! use serde_json::json;
12//!
13//! # async fn example() -> Result<(), ironflow_store::error::StoreError> {
14//! let store = InMemoryStore::new();
15//!
16//! let run = store.create_run(NewRun {
17//!     workflow_name: "test".to_string(),
18//!     trigger: TriggerKind::Manual,
19//!     payload: json!({}),
20//!     max_retries: 3,
21//!     handler_version: None,
22//!     labels: HashMap::new(),
23//!     scheduled_at: None,
24//!     created_by: None,
25//!     idempotency_key: None,
26//!     concurrency_key: None,
27//!     concurrency_limits: Vec::new(),
28//!     max_cost_usd: None,
29//! }).await?.into_run();
30//!
31//! assert_eq!(run.status.state, RunStatus::Pending);
32//! # Ok(())
33//! # }
34//! ```
35
36use std::collections::{BTreeSet, HashMap};
37use std::sync::Arc;
38
39use chrono::{DateTime, Utc};
40use tokio::sync::RwLock;
41use uuid::Uuid;
42
43use crate::entities::User;
44
45mod api_key_store;
46mod approval_delegation_store;
47mod artifact_store;
48mod audit_log_store;
49mod descendants;
50mod log_store;
51mod provider_account_store;
52mod run_store;
53mod schedule_store;
54mod secret_store;
55mod signal_store;
56mod stats_history;
57mod user_store;
58
59#[derive(Debug, Default)]
60pub(super) struct State {
61    pub(super) runs: HashMap<Uuid, crate::entities::Run>,
62    /// Idempotency key -> run holding it. Guarded by the same lock as `runs`,
63    /// so check-then-insert is atomic.
64    pub(super) idempotency_keys: HashMap<String, Uuid>,
65    pub(super) steps: HashMap<Uuid, crate::entities::Step>,
66    pub(super) step_dependencies: Vec<crate::entities::StepDependency>,
67    pub(super) artifacts: HashMap<Uuid, crate::entities::Artifact>,
68    pub(super) users: HashMap<Uuid, User>,
69    /// Group membership per user, kept sorted and deduplicated.
70    pub(super) user_groups: HashMap<Uuid, BTreeSet<String>>,
71    pub(super) api_keys: HashMap<Uuid, crate::entities::ApiKey>,
72    pub(super) secrets: HashMap<String, EncryptedSecret>,
73    pub(super) schedules: HashMap<Uuid, crate::entities::Schedule>,
74    pub(super) approval_delegations: HashMap<Uuid, crate::entities::ApprovalDelegation>,
75    pub(super) audit_logs: Vec<crate::entities::AuditLogEntry>,
76    pub(super) log_entries: Vec<crate::entities::LogEntry>,
77    pub(super) provider_accounts: HashMap<Uuid, crate::entities::ProviderAccount>,
78    /// Latest window per `(account_id, window, model_scope)`, `""` for no scope.
79    pub(super) provider_account_windows:
80        HashMap<(Uuid, String, String), crate::entities::ProviderAccountWindow>,
81    pub(super) provider_account_usage: Vec<crate::entities::ProviderAccountUsagePoint>,
82    pub(super) signals: Vec<crate::entities::Signal>,
83    /// Idempotency ID -> signal holding it. Guarded by the same lock as
84    /// `signals`, so check-then-insert is atomic.
85    pub(super) signal_idempotency: HashMap<String, Uuid>,
86    /// Issued refresh tokens, keyed by their SHA-256 hash.
87    pub(super) refresh_tokens: HashMap<String, StoredRefreshToken>,
88}
89
90#[derive(Debug, Clone)]
91pub(super) struct StoredRefreshToken {
92    pub(super) user_id: Uuid,
93    pub(super) expires_at: DateTime<Utc>,
94}
95
96#[derive(Debug, Clone)]
97pub(super) struct EncryptedSecret {
98    pub(super) id: Uuid,
99    pub(super) key: String,
100    #[cfg(feature = "secret-store")]
101    pub(super) encrypted_value: Vec<u8>,
102    #[cfg(feature = "secret-store")]
103    pub(super) nonce: Vec<u8>,
104    #[cfg(feature = "secret-store")]
105    pub(super) key_version: i32,
106    pub(super) created_at: chrono::DateTime<chrono::Utc>,
107    pub(super) updated_at: chrono::DateTime<chrono::Utc>,
108}
109
110/// In-memory store backed by `Arc<RwLock<..>>`.
111///
112/// Thread-safe and cheap to clone. All data is held in memory and lost on drop.
113/// Implements [`Store`](crate::store::Store) so a single `Arc<InMemoryStore>`
114/// covers runs, users, API keys, and secrets.
115///
116/// # Examples
117///
118/// ```
119/// use ironflow_store::memory::InMemoryStore;
120///
121/// let store = InMemoryStore::new();
122/// let store2 = store.clone(); // cheap Arc clone
123/// ```
124#[derive(Debug, Clone)]
125pub struct InMemoryStore {
126    pub(super) state: Arc<RwLock<State>>,
127    #[cfg(feature = "secret-store")]
128    pub(super) key_ring: Option<Arc<crate::crypto::KeyRing>>,
129}
130
131impl InMemoryStore {
132    /// Create a new empty in-memory store.
133    ///
134    /// # Examples
135    ///
136    /// ```
137    /// use ironflow_store::memory::InMemoryStore;
138    ///
139    /// let store = InMemoryStore::new();
140    /// ```
141    pub fn new() -> Self {
142        Self {
143            state: Arc::new(RwLock::new(State::default())),
144            #[cfg(feature = "secret-store")]
145            key_ring: None,
146        }
147    }
148
149    /// Set a single, unversioned master key for secret encryption.
150    ///
151    /// Shorthand for a key ring holding this key alone at
152    /// [`LEGACY_KEY_VERSION`](crate::crypto::LEGACY_KEY_VERSION).
153    ///
154    /// Required before using [`SecretStore`](crate::secret_store::SecretStore)
155    /// methods that read/write secret values. Without a key, those methods
156    /// return [`StoreError::Crypto`](crate::error::StoreError::Crypto).
157    ///
158    /// Listing and deleting secrets works without a key.
159    ///
160    /// # Examples
161    ///
162    /// ```
163    /// use ironflow_store::memory::InMemoryStore;
164    /// use ironflow_store::crypto::MasterKey;
165    ///
166    /// # fn example() -> Result<(), ironflow_store::crypto::CryptoError> {
167    /// let mut store = InMemoryStore::new();
168    /// let key = MasterKey::from_hex(
169    ///     "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef"
170    /// )?;
171    /// store.set_master_key(key);
172    /// # Ok(())
173    /// # }
174    /// ```
175    #[cfg(feature = "secret-store")]
176    pub fn set_master_key(&mut self, key: crate::crypto::MasterKey) {
177        self.set_key_ring(crate::crypto::KeyRing::single(key));
178    }
179
180    /// Set the versioned key ring for secret encryption.
181    ///
182    /// New secrets are encrypted with the ring's active version; existing ones
183    /// are decrypted with whichever version they were written with.
184    ///
185    /// # Examples
186    ///
187    /// ```
188    /// use ironflow_store::memory::InMemoryStore;
189    /// use ironflow_store::crypto::KeyRing;
190    ///
191    /// # fn example() -> Result<(), ironflow_store::crypto::CryptoError> {
192    /// let mut store = InMemoryStore::new();
193    /// let spec = format!("1:{},2:{}", "aa".repeat(32), "bb".repeat(32));
194    /// store.set_key_ring(KeyRing::from_spec(&spec, Some(2))?);
195    /// # Ok(())
196    /// # }
197    /// ```
198    #[cfg(feature = "secret-store")]
199    pub fn set_key_ring(&mut self, ring: crate::crypto::KeyRing) {
200        self.key_ring = Some(Arc::new(ring));
201    }
202
203    /// Override a run's `created_at` timestamp for testing retention policies.
204    ///
205    /// # Examples
206    ///
207    /// ```no_run
208    /// use chrono::{Utc, TimeDelta};
209    /// use ironflow_store::memory::InMemoryStore;
210    /// use uuid::Uuid;
211    ///
212    /// # async fn example(store: &InMemoryStore, run_id: Uuid) {
213    /// let old = Utc::now() - TimeDelta::days(100);
214    /// store.set_run_created_at(run_id, old).await;
215    /// # }
216    /// ```
217    pub async fn set_run_created_at(
218        &self,
219        run_id: Uuid,
220        created_at: chrono::DateTime<chrono::Utc>,
221    ) {
222        let mut state = self.state.write().await;
223        if let Some(run) = state.runs.get_mut(&run_id) {
224            run.created_at = created_at;
225        }
226    }
227}
228
229impl Default for InMemoryStore {
230    fn default() -> Self {
231        Self::new()
232    }
233}
234
235#[cfg(test)]
236mod tests {
237    use std::collections::HashMap;
238
239    use serde_json::json;
240
241    use crate::entities::{NewRun, TriggerKind};
242
243    use super::InMemoryStore;
244
245    pub(crate) fn new_run_req(name: &str) -> NewRun {
246        NewRun {
247            created_by: None,
248            workflow_name: name.to_string(),
249            trigger: TriggerKind::Manual,
250            payload: json!({}),
251            max_retries: 3,
252            handler_version: None,
253            labels: HashMap::new(),
254            scheduled_at: None,
255            idempotency_key: None,
256            concurrency_key: None,
257            concurrency_limits: Vec::new(),
258            max_cost_usd: None,
259        }
260    }
261
262    pub(crate) async fn create_terminal_run(
263        store: &InMemoryStore,
264        name: &str,
265        status: crate::entities::RunStatus,
266    ) -> crate::entities::Run {
267        use crate::store::RunStore;
268
269        let run = store
270            .create_run(new_run_req(name))
271            .await
272            .unwrap()
273            .into_run();
274        store
275            .update_run_status(run.id, crate::entities::RunStatus::Running)
276            .await
277            .unwrap();
278        store.update_run_status(run.id, status).await.unwrap();
279        store.get_run(run.id).await.unwrap().unwrap()
280    }
281}