1use std::collections::HashMap;
8
9use chrono::{DateTime, Utc};
10
11use crate::secrets_manager::SecretsError;
12
13#[derive(Debug, Clone, Copy, PartialEq, Eq)]
15pub enum IsolationLevel {
16 ReadUncommitted,
18 ReadCommitted,
20 RepeatableRead,
22 Serializable,
24}
25
26impl std::fmt::Display for IsolationLevel {
27 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
28 match self {
29 Self::ReadUncommitted => write!(f, "READ UNCOMMITTED"),
30 Self::ReadCommitted => write!(f, "READ COMMITTED"),
31 Self::RepeatableRead => write!(f, "REPEATABLE READ"),
32 Self::Serializable => write!(f, "SERIALIZABLE"),
33 }
34 }
35}
36
37#[derive(Debug, Clone, Copy, PartialEq, Eq)]
39pub enum TransactionState {
40 Active,
42 Committed,
44 RolledBack,
46 Error,
48}
49
50impl std::fmt::Display for TransactionState {
51 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
52 match self {
53 Self::Active => write!(f, "active"),
54 Self::Committed => write!(f, "committed"),
55 Self::RolledBack => write!(f, "rolled back"),
56 Self::Error => write!(f, "error"),
57 }
58 }
59}
60
61#[derive(Debug, Clone)]
63pub struct TransactionContext {
64 pub transaction_id: String,
66 pub user_id: String,
68 pub session_id: String,
70 pub request_id: String,
72 pub started_at: DateTime<Utc>,
74 pub isolation_level: IsolationLevel,
76 pub state: TransactionState,
78 pub key_version: u32,
80 pub operations: Vec<String>,
82 pub metadata: HashMap<String, String>,
84 pub user_role: Option<String>,
86 pub client_ip: Option<String>,
88}
89
90impl TransactionContext {
91 pub fn new(
93 user_id: impl Into<String>,
94 session_id: impl Into<String>,
95 request_id: impl Into<String>,
96 ) -> Self {
97 let transaction_id = format!(
99 "txn_{}_{}",
100 std::time::SystemTime::now()
101 .duration_since(std::time::UNIX_EPOCH)
102 .unwrap()
103 .as_micros(),
104 uuid::Uuid::new_v4().to_string()[..8].to_string()
105 );
106
107 Self {
108 transaction_id,
109 user_id: user_id.into(),
110 session_id: session_id.into(),
111 request_id: request_id.into(),
112 started_at: Utc::now(),
113 isolation_level: IsolationLevel::ReadCommitted,
114 state: TransactionState::Active,
115 key_version: 1,
116 operations: Vec::new(),
117 metadata: HashMap::new(),
118 user_role: None,
119 client_ip: None,
120 }
121 }
122
123 pub fn with_isolation(mut self, level: IsolationLevel) -> Self {
125 self.isolation_level = level;
126 self
127 }
128
129 pub fn with_key_version(mut self, version: u32) -> Self {
131 self.key_version = version;
132 self
133 }
134
135 pub fn add_operation(&mut self, operation: impl Into<String>) {
137 self.operations.push(operation.into());
138 }
139
140 pub fn with_metadata(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
142 self.metadata.insert(key.into(), value.into());
143 self
144 }
145
146 pub fn with_role(mut self, role: impl Into<String>) -> Self {
148 self.user_role = Some(role.into());
149 self
150 }
151
152 pub fn with_client_ip(mut self, ip: impl Into<String>) -> Self {
154 self.client_ip = Some(ip.into());
155 self
156 }
157
158 pub fn commit(&mut self) {
160 self.state = TransactionState::Committed;
161 }
162
163 pub fn rollback(&mut self) {
165 self.state = TransactionState::RolledBack;
166 self.operations.clear();
167 }
168
169 pub fn error(&mut self) {
171 self.state = TransactionState::Error;
172 }
173
174 pub fn duration(&self) -> chrono::Duration {
176 Utc::now() - self.started_at
177 }
178
179 pub fn is_active(&self) -> bool {
181 self.state == TransactionState::Active
182 }
183
184 pub fn operation_count(&self) -> usize {
186 self.operations.len()
187 }
188}
189
190#[derive(Debug, Clone)]
192pub struct Savepoint {
193 pub name: String,
195 pub transaction_id: String,
197 pub created_at: DateTime<Utc>,
199 pub operations_before: usize,
201}
202
203impl Savepoint {
204 pub fn new(
206 name: impl Into<String>,
207 transaction_id: impl Into<String>,
208 operations_count: usize,
209 ) -> Self {
210 Self {
211 name: name.into(),
212 transaction_id: transaction_id.into(),
213 created_at: Utc::now(),
214 operations_before: operations_count,
215 }
216 }
217}
218
219pub struct TransactionManager {
221 active_transactions: HashMap<String, TransactionContext>,
223 savepoints: HashMap<String, Vec<Savepoint>>,
225}
226
227impl TransactionManager {
228 pub fn new() -> Self {
230 Self {
231 active_transactions: HashMap::new(),
232 savepoints: HashMap::new(),
233 }
234 }
235
236 pub fn begin(&mut self, context: TransactionContext) -> Result<String, SecretsError> {
238 let txn_id = context.transaction_id.clone();
239
240 if self.active_transactions.contains_key(&txn_id) {
241 return Err(SecretsError::ValidationError(format!(
242 "Transaction {} already active",
243 txn_id
244 )));
245 }
246
247 self.active_transactions.insert(txn_id.clone(), context);
248 Ok(txn_id)
249 }
250
251 pub fn get_transaction(&self, txn_id: &str) -> Option<&TransactionContext> {
253 self.active_transactions.get(txn_id)
254 }
255
256 pub fn get_transaction_mut(&mut self, txn_id: &str) -> Option<&mut TransactionContext> {
258 self.active_transactions.get_mut(txn_id)
259 }
260
261 pub fn commit(&mut self, txn_id: &str) -> Result<(), SecretsError> {
263 if let Some(txn) = self.active_transactions.get_mut(txn_id) {
264 txn.commit();
265 self.savepoints.remove(txn_id);
266 Ok(())
267 } else {
268 Err(SecretsError::ValidationError(format!("Transaction {} not found", txn_id)))
269 }
270 }
271
272 pub fn rollback(&mut self, txn_id: &str) -> Result<(), SecretsError> {
274 if let Some(txn) = self.active_transactions.get_mut(txn_id) {
275 txn.rollback();
276 self.savepoints.remove(txn_id);
277 Ok(())
278 } else {
279 Err(SecretsError::ValidationError(format!("Transaction {} not found", txn_id)))
280 }
281 }
282
283 pub fn savepoint(&mut self, txn_id: &str, name: impl Into<String>) -> Result<(), SecretsError> {
285 if let Some(txn) = self.active_transactions.get(txn_id) {
286 let savepoint = Savepoint::new(name, txn_id, txn.operation_count());
287 self.savepoints
288 .entry(txn_id.to_string())
289 .or_insert_with(Vec::new)
290 .push(savepoint);
291 Ok(())
292 } else {
293 Err(SecretsError::ValidationError(format!("Transaction {} not found", txn_id)))
294 }
295 }
296
297 pub fn rollback_to_savepoint(&mut self, txn_id: &str, name: &str) -> Result<(), SecretsError> {
299 if let Some(savepoints) = self.savepoints.get_mut(txn_id) {
300 if let Some(sp_idx) = savepoints.iter().position(|sp| sp.name == name) {
301 let savepoint = savepoints.remove(sp_idx);
302
303 if let Some(txn) = self.active_transactions.get_mut(txn_id) {
304 txn.operations.truncate(savepoint.operations_before);
306 return Ok(());
307 }
308 }
309 Err(SecretsError::ValidationError(format!("Savepoint {} not found", name)))
310 } else {
311 Err(SecretsError::ValidationError(format!(
312 "Transaction {} has no savepoints",
313 txn_id
314 )))
315 }
316 }
317
318 pub fn active_transactions(&self) -> Vec<&str> {
320 self.active_transactions.keys().map(|s| s.as_str()).collect()
321 }
322
323 pub fn active_count(&self) -> usize {
325 self.active_transactions.len()
326 }
327
328 pub fn cleanup_completed(&mut self) {
330 self.active_transactions.retain(|_, txn| txn.is_active());
331 }
332}
333
334impl Default for TransactionManager {
335 fn default() -> Self {
336 Self::new()
337 }
338}
339
340#[cfg(test)]
341mod tests {
342 use super::*;
343
344 #[test]
345 fn test_isolation_level_display() {
346 assert_eq!(IsolationLevel::ReadUncommitted.to_string(), "READ UNCOMMITTED");
347 assert_eq!(IsolationLevel::ReadCommitted.to_string(), "READ COMMITTED");
348 assert_eq!(IsolationLevel::RepeatableRead.to_string(), "REPEATABLE READ");
349 assert_eq!(IsolationLevel::Serializable.to_string(), "SERIALIZABLE");
350 }
351
352 #[test]
353 fn test_transaction_state_display() {
354 assert_eq!(TransactionState::Active.to_string(), "active");
355 assert_eq!(TransactionState::Committed.to_string(), "committed");
356 assert_eq!(TransactionState::RolledBack.to_string(), "rolled back");
357 assert_eq!(TransactionState::Error.to_string(), "error");
358 }
359
360 #[test]
361 fn test_transaction_context_creation() {
362 let ctx = TransactionContext::new("user123", "sess456", "req789");
363 assert_eq!(ctx.user_id, "user123");
364 assert_eq!(ctx.session_id, "sess456");
365 assert_eq!(ctx.request_id, "req789");
366 assert_eq!(ctx.state, TransactionState::Active);
367 assert_eq!(ctx.isolation_level, IsolationLevel::ReadCommitted);
368 assert_eq!(ctx.key_version, 1);
369 assert!(ctx.transaction_id.starts_with("txn_"));
370 }
371
372 #[test]
373 fn test_transaction_context_with_isolation() {
374 let ctx = TransactionContext::new("user123", "sess456", "req789")
375 .with_isolation(IsolationLevel::Serializable);
376 assert_eq!(ctx.isolation_level, IsolationLevel::Serializable);
377 }
378
379 #[test]
380 fn test_transaction_context_with_key_version() {
381 let ctx = TransactionContext::new("user123", "sess456", "req789").with_key_version(2);
382 assert_eq!(ctx.key_version, 2);
383 }
384
385 #[test]
386 fn test_transaction_context_add_operation() {
387 let mut ctx = TransactionContext::new("user123", "sess456", "req789");
388 ctx.add_operation("INSERT users");
389 ctx.add_operation("UPDATE roles");
390 assert_eq!(ctx.operation_count(), 2);
391 }
392
393 #[test]
394 fn test_transaction_context_with_metadata() {
395 let ctx =
396 TransactionContext::new("user123", "sess456", "req789").with_metadata("source", "api");
397 assert_eq!(ctx.metadata.get("source"), Some(&"api".to_string()));
398 }
399
400 #[test]
401 fn test_transaction_context_with_role() {
402 let ctx = TransactionContext::new("user123", "sess456", "req789").with_role("admin");
403 assert_eq!(ctx.user_role, Some("admin".to_string()));
404 }
405
406 #[test]
407 fn test_transaction_context_with_client_ip() {
408 let ctx =
409 TransactionContext::new("user123", "sess456", "req789").with_client_ip("192.168.1.1");
410 assert_eq!(ctx.client_ip, Some("192.168.1.1".to_string()));
411 }
412
413 #[test]
414 fn test_transaction_context_commit() {
415 let mut ctx = TransactionContext::new("user123", "sess456", "req789");
416 assert_eq!(ctx.state, TransactionState::Active);
417 ctx.commit();
418 assert_eq!(ctx.state, TransactionState::Committed);
419 }
420
421 #[test]
422 fn test_transaction_context_rollback() {
423 let mut ctx = TransactionContext::new("user123", "sess456", "req789");
424 ctx.add_operation("INSERT users");
425 ctx.rollback();
426 assert_eq!(ctx.state, TransactionState::RolledBack);
427 assert_eq!(ctx.operation_count(), 0);
428 }
429
430 #[test]
431 fn test_transaction_context_error() {
432 let mut ctx = TransactionContext::new("user123", "sess456", "req789");
433 ctx.error();
434 assert_eq!(ctx.state, TransactionState::Error);
435 }
436
437 #[test]
438 fn test_transaction_context_is_active() {
439 let mut ctx = TransactionContext::new("user123", "sess456", "req789");
440 assert!(ctx.is_active());
441 ctx.commit();
442 assert!(!ctx.is_active());
443 }
444
445 #[test]
446 fn test_savepoint_creation() {
447 let sp = Savepoint::new("sp1", "txn123", 5);
448 assert_eq!(sp.name, "sp1");
449 assert_eq!(sp.transaction_id, "txn123");
450 assert_eq!(sp.operations_before, 5);
451 }
452
453 #[test]
454 fn test_transaction_manager_begin() {
455 let mut manager = TransactionManager::new();
456 let ctx = TransactionContext::new("user123", "sess456", "req789");
457 let txn_id = ctx.transaction_id.clone();
458
459 let result = manager.begin(ctx);
460 assert!(result.is_ok());
461 assert_eq!(result.unwrap(), txn_id);
462 assert_eq!(manager.active_count(), 1);
463 }
464
465 #[test]
466 fn test_transaction_manager_get_transaction() {
467 let mut manager = TransactionManager::new();
468 let ctx = TransactionContext::new("user123", "sess456", "req789");
469 let txn_id = ctx.transaction_id.clone();
470
471 manager.begin(ctx).unwrap();
472
473 let retrieved = manager.get_transaction(&txn_id);
474 assert!(retrieved.is_some());
475 assert_eq!(retrieved.unwrap().user_id, "user123");
476 }
477
478 #[test]
479 fn test_transaction_manager_commit() {
480 let mut manager = TransactionManager::new();
481 let ctx = TransactionContext::new("user123", "sess456", "req789");
482 let txn_id = ctx.transaction_id.clone();
483
484 manager.begin(ctx).unwrap();
485 let result = manager.commit(&txn_id);
486
487 assert!(result.is_ok());
488 let txn = manager.get_transaction(&txn_id);
489 assert_eq!(txn.unwrap().state, TransactionState::Committed);
490 }
491
492 #[test]
493 fn test_transaction_manager_rollback() {
494 let mut manager = TransactionManager::new();
495 let ctx = TransactionContext::new("user123", "sess456", "req789");
496 let txn_id = ctx.transaction_id.clone();
497
498 manager.begin(ctx).unwrap();
499 let result = manager.rollback(&txn_id);
500
501 assert!(result.is_ok());
502 let txn = manager.get_transaction(&txn_id);
503 assert_eq!(txn.unwrap().state, TransactionState::RolledBack);
504 }
505
506 #[test]
507 fn test_transaction_manager_savepoint() {
508 let mut manager = TransactionManager::new();
509 let ctx = TransactionContext::new("user123", "sess456", "req789");
510 let txn_id = ctx.transaction_id.clone();
511
512 manager.begin(ctx).unwrap();
513 let result = manager.savepoint(&txn_id, "sp1");
514
515 assert!(result.is_ok());
516 }
517
518 #[test]
519 fn test_transaction_manager_rollback_to_savepoint() {
520 let mut manager = TransactionManager::new();
521 let mut ctx = TransactionContext::new("user123", "sess456", "req789");
522 ctx.add_operation("OP1");
523 let txn_id = ctx.transaction_id.clone();
524
525 manager.begin(ctx).unwrap();
526 manager.savepoint(&txn_id, "sp1").unwrap();
527
528 {
529 let txn = manager.get_transaction_mut(&txn_id).unwrap();
530 txn.add_operation("OP2");
531 }
532
533 let result = manager.rollback_to_savepoint(&txn_id, "sp1");
534 assert!(result.is_ok());
535
536 let txn = manager.get_transaction(&txn_id).unwrap();
537 assert_eq!(txn.operation_count(), 1);
538 }
539
540 #[test]
541 fn test_transaction_manager_active_transactions() {
542 let mut manager = TransactionManager::new();
543 let ctx1 = TransactionContext::new("user1", "sess1", "req1");
544 let ctx2 = TransactionContext::new("user2", "sess2", "req2");
545
546 manager.begin(ctx1).unwrap();
547 manager.begin(ctx2).unwrap();
548
549 let active = manager.active_transactions();
550 assert_eq!(active.len(), 2);
551 }
552
553 #[test]
554 fn test_transaction_manager_cleanup_completed() {
555 let mut manager = TransactionManager::new();
556 let ctx1 = TransactionContext::new("user1", "sess1", "req1");
557 let ctx2 = TransactionContext::new("user2", "sess2", "req2");
558
559 let id1 = ctx1.transaction_id.clone();
560
561 manager.begin(ctx1).unwrap();
562 manager.begin(ctx2).unwrap();
563
564 manager.commit(&id1).unwrap();
565 manager.cleanup_completed();
566
567 assert_eq!(manager.active_count(), 1);
568 assert!(manager.get_transaction(&id1).is_none());
569 }
570}