1use crate::Result;
7use mockforge_core::VirtualClock;
8use std::sync::Arc;
9
10#[derive(Debug, Clone)]
12pub struct AgingRule {
13 pub entity_name: String,
15 pub expiration_field: String,
17 pub expiration_duration: u64,
19 pub action: AgingAction,
21}
22
23#[derive(Debug, Clone, Copy, PartialEq, Eq)]
25pub enum AgingAction {
26 Delete,
28 MarkExpired,
30 Archive,
32}
33
34pub struct AgingManager {
36 rules: Vec<AgingRule>,
38 virtual_clock: Option<Arc<VirtualClock>>,
40}
41
42impl AgingManager {
43 pub fn new() -> Self {
45 Self {
46 rules: Vec::new(),
47 virtual_clock: None,
48 }
49 }
50
51 pub fn with_virtual_clock(clock: Arc<VirtualClock>) -> Self {
53 Self {
54 rules: Vec::new(),
55 virtual_clock: Some(clock),
56 }
57 }
58
59 pub fn set_virtual_clock(&mut self, clock: Option<Arc<VirtualClock>>) {
61 self.virtual_clock = clock;
62 }
63
64 fn now(&self) -> chrono::DateTime<chrono::Utc> {
66 if let Some(ref clock) = self.virtual_clock {
67 clock.now()
68 } else {
69 chrono::Utc::now()
70 }
71 }
72
73 pub fn add_rule(&mut self, rule: AgingRule) {
75 self.rules.push(rule);
76 }
77
78 pub async fn cleanup_expired(
82 &self,
83 database: &dyn crate::database::VirtualDatabase,
84 registry: &crate::entities::EntityRegistry,
85 ) -> Result<usize> {
86 let mut total_cleaned = 0;
87
88 for rule in &self.rules {
89 let entity = match registry.get(&rule.entity_name) {
91 Some(e) => e,
92 None => continue, };
94
95 let table_name = entity.table_name();
96 let now = self.now();
97
98 let query = format!("SELECT * FROM {}", table_name);
100 let records = database.query(&query, &[]).await?;
101
102 for record in records {
103 if let Some(expiration_value) = record.get(&rule.expiration_field) {
105 let expiration_time = match expiration_value {
107 serde_json::Value::String(s) => {
108 match chrono::DateTime::parse_from_rfc3339(s) {
110 Ok(dt) => dt.with_timezone(&chrono::Utc),
111 Err(_) => continue, }
113 }
114 serde_json::Value::Number(n) => {
115 if let Some(ts) = n.as_i64() {
117 chrono::DateTime::from_timestamp(ts, 0)
118 .unwrap_or_else(|| self.now())
119 } else {
120 continue; }
122 }
123 _ => continue, };
125
126 let age = now.signed_duration_since(expiration_time);
128 if age.num_seconds() > rule.expiration_duration as i64 {
129 match rule.action {
131 AgingAction::Delete => {
132 let pk_field = entity
134 .schema
135 .primary_key
136 .first()
137 .map(|s| s.as_str())
138 .unwrap_or("id");
139 if let Some(pk_value) = record.get(pk_field) {
140 let delete_query = format!(
141 "DELETE FROM {} WHERE {} = ?",
142 table_name, pk_field
143 );
144 database.execute(&delete_query, &[pk_value.clone()]).await?;
145 total_cleaned += 1;
146 }
147 }
148 AgingAction::MarkExpired => {
149 let pk_field = entity
151 .schema
152 .primary_key
153 .first()
154 .map(|s| s.as_str())
155 .unwrap_or("id");
156 if let Some(pk_value) = record.get(pk_field) {
157 let update_query = format!(
158 "UPDATE {} SET status = ? WHERE {} = ?",
159 table_name, pk_field
160 );
161 database
162 .execute(
163 &update_query,
164 &[
165 serde_json::Value::String("expired".to_string()),
166 pk_value.clone(),
167 ],
168 )
169 .await?;
170 total_cleaned += 1;
171 }
172 }
173 AgingAction::Archive => {
174 let pk_field = entity
176 .schema
177 .primary_key
178 .first()
179 .map(|s| s.as_str())
180 .unwrap_or("id");
181 if let Some(pk_value) = record.get(pk_field) {
182 let update_query = format!(
183 "UPDATE {} SET archived = ? WHERE {} = ?",
184 table_name, pk_field
185 );
186 database
187 .execute(
188 &update_query,
189 &[serde_json::Value::Bool(true), pk_value.clone()],
190 )
191 .await?;
192 total_cleaned += 1;
193 }
194 }
195 }
196 }
197 }
198 }
199 }
200
201 Ok(total_cleaned)
202 }
203
204 pub async fn update_timestamps(
208 &self,
209 database: &dyn crate::database::VirtualDatabase,
210 table: &str,
211 primary_key_field: &str,
212 primary_key_value: &serde_json::Value,
213 ) -> Result<()> {
214 let now = self.now().to_rfc3339();
216 let update_query =
217 format!("UPDATE {} SET updated_at = ? WHERE {} = ?", table, primary_key_field);
218
219 let _ = database
221 .execute(&update_query, &[serde_json::Value::String(now), primary_key_value.clone()])
222 .await;
223
224 Ok(())
225 }
226}
227
228impl Default for AgingManager {
229 fn default() -> Self {
230 Self::new()
231 }
232}
233
234#[cfg(test)]
235mod tests {
236 use super::*;
237 use mockforge_core::VirtualClock;
238 use std::sync::Arc;
239
240 #[test]
241 fn test_aging_with_virtual_clock() {
242 let clock = Arc::new(VirtualClock::new());
244 let initial_time = chrono::Utc::now();
245 clock.enable_and_set(initial_time);
246
247 let aging_manager = AgingManager::with_virtual_clock(clock.clone());
248
249 let now = aging_manager.now();
251 assert!((now - initial_time).num_seconds().abs() < 1);
252
253 clock.advance(chrono::Duration::hours(2));
255
256 let advanced_now = aging_manager.now();
258 let elapsed = advanced_now - initial_time;
259 assert!(elapsed.num_hours() >= 1 && elapsed.num_hours() <= 3);
260 }
261
262 #[test]
263 fn test_aging_timestamps_with_virtual_clock() {
264 let clock = Arc::new(VirtualClock::new());
265 let initial_time = chrono::Utc::now();
266 clock.enable_and_set(initial_time);
267
268 let aging_manager = AgingManager::with_virtual_clock(clock.clone());
269
270 clock.advance(chrono::Duration::days(30));
272
273 let now = aging_manager.now();
276 let elapsed = now - initial_time;
277 assert!(elapsed.num_days() >= 29 && elapsed.num_days() <= 31);
278 }
279
280 #[test]
281 fn test_one_month_aging_scenario() {
282 let clock = Arc::new(VirtualClock::new());
284 let initial_time = chrono::Utc::now();
285 clock.enable_and_set(initial_time);
286
287 let aging_manager = AgingManager::with_virtual_clock(clock.clone());
288
289 let start_time = aging_manager.now();
291 assert!((start_time - initial_time).num_seconds().abs() < 1);
292
293 clock.advance(chrono::Duration::days(30));
295
296 let month_later = aging_manager.now();
298 let elapsed = month_later - start_time;
299
300 assert!(
302 elapsed.num_days() >= 29 && elapsed.num_days() <= 31,
303 "Expected ~30 days, got {} days",
304 elapsed.num_days()
305 );
306 }
307}