use crate::error::{DatabaseError, DatabaseResult, ErrorContext}; use crate::pool::DatabasePool; use config::database::TransactionConfig; use sqlx::FromRow; use std::future::Future; use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; use std::time::{Duration, Instant}; use tracing::{debug, error, info, warn}; use uuid::Uuid; // TransactionConfig is now imported from the config crate /// Transaction manager for handling database transactions #[derive(Debug)] pub struct TransactionManager { pool: DatabasePool, config: TransactionConfig, active_transactions: Arc, total_transactions: Arc, failed_transactions: Arc, retry_attempts: Arc, } impl TransactionManager { /// Create a new transaction manager pub fn new(pool: DatabasePool, config: TransactionConfig) -> Self { Self { pool, config, active_transactions: Arc::new(AtomicU64::new(0)), total_transactions: Arc::new(AtomicU64::new(0)), failed_transactions: Arc::new(AtomicU64::new(0)), retry_attempts: Arc::new(AtomicU64::new(0)), } } /// Begin a new transaction with default timeout pub async fn begin(&self) -> DatabaseResult { self.begin_with_timeout(Duration::from_secs(self.config.default_timeout_secs)) .await } /// Begin a new transaction with custom timeout pub async fn begin_with_timeout( &self, timeout_duration: Duration, ) -> DatabaseResult { let start_time = Instant::now(); self.total_transactions.fetch_add(1_u64, Ordering::Relaxed); self.active_transactions.fetch_add(1_u64, Ordering::Relaxed); debug!( "Beginning new database transaction with {}s timeout", timeout_duration.as_secs() ); // Get the transaction directly from the pool to avoid lifetime issues let transaction = self.pool.inner().begin().await?; let transaction_id = Uuid::new_v4(); debug!("Transaction {} started successfully", transaction_id); Ok(DatabaseTransaction { inner: Some(transaction), id: transaction_id, start_time, timeout: timeout_duration, savepoints: Vec::new(), config: self.config.clone(), manager_stats: TransactionManagerStats { active_transactions: self.active_transactions.clone(), failed_transactions: self.failed_transactions.clone(), }, }) } /// Execute a closure within a transaction with retry logic pub async fn with_transaction(&self, f: F) -> DatabaseResult where F: Fn(DatabaseTransaction) -> Fut + Send + Sync, Fut: Future> + Send, R: Send, { self.with_transaction_timeout(f, Duration::from_secs(self.config.default_timeout_secs)) .await } /// Execute a closure within a transaction with custom timeout and retry logic pub async fn with_transaction_timeout( &self, f: F, timeout_duration: Duration, ) -> DatabaseResult where F: Fn(DatabaseTransaction) -> Fut + Send + Sync, Fut: Future> + Send, R: Send, { let mut attempts = 0; let max_attempts = if self.config.enable_retry { self.config.max_retries + 1 } else { 1 }; loop { attempts += 1; let tx = self.begin_with_timeout(timeout_duration).await?; let tx_id = tx.id(); match f(tx).await { Ok((result, tx)) => match tx.commit().await { Ok(_) => { debug!( "Transaction {} completed successfully on attempt {}", tx_id, attempts ); return Ok(result); }, Err(e) if e.is_retryable() && attempts < max_attempts => { self.retry_attempts.fetch_add(1_u64, Ordering::Relaxed); warn!( "Transaction {} commit failed (attempt {}): {}. Retrying...", tx_id, attempts, e ); self.delay_retry(attempts).await; continue; }, Err(e) => { error!( "Transaction {} commit failed on attempt {}: {}", tx_id, attempts, e ); return Err(e); }, }, Err(e) if e.is_retryable() && attempts < max_attempts => { self.retry_attempts.fetch_add(1_u64, Ordering::Relaxed); warn!( "Transaction {} failed (attempt {}): {}. Retrying...", tx_id, attempts, e ); self.delay_retry(attempts).await; continue; }, Err(e) => { error!( "Transaction {} failed on attempt {}: {}", tx_id, attempts, e ); return Err(e); }, } } } /// Get transaction statistics pub fn stats(&self) -> TransactionStats { TransactionStats { active_transactions: self.active_transactions.load(Ordering::Relaxed), total_transactions: self.total_transactions.load(Ordering::Relaxed), failed_transactions: self.failed_transactions.load(Ordering::Relaxed), retry_attempts: self.retry_attempts.load(Ordering::Relaxed), } } /// Reset transaction statistics pub fn reset_stats(&self) { self.active_transactions.store(0_u64, Ordering::Relaxed); self.total_transactions.store(0_u64, Ordering::Relaxed); self.failed_transactions.store(0_u64, Ordering::Relaxed); self.retry_attempts.store(0_u64, Ordering::Relaxed); info!("Transaction manager statistics reset"); } /// Delay before retry with exponential backoff async fn delay_retry(&self, attempt: u32) { let delay = Duration::from_millis( self.config.retry_delay_ms * (2_u64.pow(attempt.saturating_sub(1))), ); debug!( "Waiting {}ms before retry attempt {}", delay.as_millis(), attempt + 1 ); tokio::time::sleep(delay).await; } } impl Clone for TransactionManager { fn clone(&self) -> Self { Self { pool: self.pool.clone(), config: self.config.clone(), active_transactions: self.active_transactions.clone(), total_transactions: self.total_transactions.clone(), failed_transactions: self.failed_transactions.clone(), retry_attempts: self.retry_attempts.clone(), } } } /// A database transaction wrapper with enhanced features pub struct DatabaseTransaction { inner: Option>, id: Uuid, start_time: Instant, timeout: Duration, savepoints: Vec, config: TransactionConfig, manager_stats: TransactionManagerStats, } /// Statistics tracking for transaction manager #[derive(Debug, Clone)] struct TransactionManagerStats { active_transactions: Arc, failed_transactions: Arc, } impl DatabaseTransaction { /// Get the transaction ID pub fn id(&self) -> Uuid { self.id } /// Get the elapsed time since transaction start pub fn elapsed(&self) -> Duration { self.start_time.elapsed() } /// Check if the transaction has timed out pub fn is_timed_out(&self) -> bool { self.elapsed() > self.timeout } /// Get remaining time before timeout pub fn remaining_time(&self) -> Duration { self.timeout.saturating_sub(self.elapsed()) } /// Create a savepoint pub async fn savepoint(&mut self, name: &str) -> DatabaseResult<()> { if self.savepoints.len() >= self.config.max_savepoints as usize { return Err(DatabaseError::Transaction { message: format!( "Maximum number of savepoints ({}) exceeded", self.config.max_savepoints ), }); } if self.is_timed_out() { return Err(DatabaseError::Timeout { operation: format!("savepoint_{}", name), timeout_secs: self.timeout.as_secs(), }); } let savepoint_name = format!("sp_{}_{}", self.id.simple(), name); let inner = self.inner.as_mut().ok_or_else(|| DatabaseError::Transaction { message: "Transaction already consumed".to_string(), })?; sqlx::query(&format!("SAVEPOINT {}", savepoint_name)) .execute(&mut **inner) .await .with_context(&format!("Failed to create savepoint {}", savepoint_name))?; self.savepoints.push(savepoint_name.clone()); debug!( "Created savepoint {} for transaction {}", savepoint_name, self.id ); Ok(()) } /// Rollback to a savepoint pub async fn rollback_to_savepoint(&mut self, name: &str) -> DatabaseResult<()> { let savepoint_name = format!("sp_{}_{}", self.id.simple(), name); if !self.savepoints.contains(&savepoint_name) { return Err(DatabaseError::Transaction { message: format!("Savepoint {} not found", savepoint_name), }); } if self.is_timed_out() { return Err(DatabaseError::Timeout { operation: format!("rollback_to_savepoint_{}", name), timeout_secs: self.timeout.as_secs(), }); } let inner = self.inner.as_mut().ok_or_else(|| DatabaseError::Transaction { message: "Transaction already consumed".to_string(), })?; sqlx::query(&format!("ROLLBACK TO SAVEPOINT {}", savepoint_name)) .execute(&mut **inner) .await .with_context(&format!( "Failed to rollback to savepoint {}", savepoint_name ))?; // Remove this savepoint and all subsequent ones if let Some(pos) = self.savepoints.iter().position(|sp| sp == &savepoint_name) { self.savepoints.truncate(pos); } debug!( "Rolled back to savepoint {} for transaction {}", savepoint_name, self.id ); Ok(()) } /// Release a savepoint pub async fn release_savepoint(&mut self, name: &str) -> DatabaseResult<()> { let savepoint_name = format!("sp_{}_{}", self.id.simple(), name); if !self.savepoints.contains(&savepoint_name) { return Err(DatabaseError::Transaction { message: format!("Savepoint {} not found", savepoint_name), }); } let inner = self.inner.as_mut().ok_or_else(|| DatabaseError::Transaction { message: "Transaction already consumed".to_string(), })?; sqlx::query(&format!("RELEASE SAVEPOINT {}", savepoint_name)) .execute(&mut **inner) .await .with_context(&format!("Failed to release savepoint {}", savepoint_name))?; self.savepoints.retain(|sp| sp != &savepoint_name); debug!( "Released savepoint {} for transaction {}", savepoint_name, self.id ); Ok(()) } /// Execute a query within the transaction pub async fn execute(&mut self, query: &str) -> DatabaseResult { if self.is_timed_out() { return Err(DatabaseError::Timeout { operation: "execute_query".to_string(), timeout_secs: self.timeout.as_secs(), }); } let inner = self.inner.as_mut().ok_or_else(|| DatabaseError::Transaction { message: "Transaction already consumed".to_string(), })?; let result = sqlx::query(query) .execute(&mut **inner) .await .with_query_context(query)?; Ok(result.rows_affected()) } /// Fetch all rows from a query pub async fn fetch_all(&mut self, query: &str) -> DatabaseResult> where T: for<'r> FromRow<'r, sqlx::postgres::PgRow> + Send + Unpin, { if self.is_timed_out() { return Err(DatabaseError::Timeout { operation: "fetch_all".to_string(), timeout_secs: self.timeout.as_secs(), }); } let inner = self.inner.as_mut().ok_or_else(|| DatabaseError::Transaction { message: "Transaction already consumed".to_string(), })?; let rows = sqlx::query_as(query) .fetch_all(&mut **inner) .await .with_query_context(query)?; Ok(rows) } /// Fetch one row from a query pub async fn fetch_one(&mut self, query: &str) -> DatabaseResult where T: for<'r> FromRow<'r, sqlx::postgres::PgRow> + Send + Unpin, { if self.is_timed_out() { return Err(DatabaseError::Timeout { operation: "fetch_one".to_string(), timeout_secs: self.timeout.as_secs(), }); } let inner = self.inner.as_mut().ok_or_else(|| DatabaseError::Transaction { message: "Transaction already consumed".to_string(), })?; let row = sqlx::query_as(query) .fetch_one(&mut **inner) .await .with_query_context(query)?; Ok(row) } /// Fetch optional row from a query pub async fn fetch_optional(&mut self, query: &str) -> DatabaseResult> where T: for<'r> FromRow<'r, sqlx::postgres::PgRow> + Send + Unpin, { if self.is_timed_out() { return Err(DatabaseError::Timeout { operation: "fetch_optional".to_string(), timeout_secs: self.timeout.as_secs(), }); } let inner = self.inner.as_mut().ok_or_else(|| DatabaseError::Transaction { message: "Transaction already consumed".to_string(), })?; let row = sqlx::query_as(query) .fetch_optional(&mut **inner) .await .with_query_context(query)?; Ok(row) } /// Commit the transaction pub async fn commit(mut self) -> DatabaseResult<()> { if self.is_timed_out() { return Err(DatabaseError::Timeout { operation: "commit_transaction".to_string(), timeout_secs: self.timeout.as_secs(), }); } let elapsed = self.elapsed(); let tx_id = self.id; let tx = self.inner.take().ok_or_else(|| DatabaseError::Transaction { message: "Transaction already consumed".to_string(), })?; match tx.commit().await { Ok(_) => { self.manager_stats .active_transactions .fetch_sub(1_u64, Ordering::Relaxed); info!( "Transaction {} committed successfully in {}ms", tx_id, elapsed.as_millis() ); Ok(()) }, Err(e) => { self.manager_stats .active_transactions .fetch_sub(1_u64, Ordering::Relaxed); self.manager_stats .failed_transactions .fetch_add(1_u64, Ordering::Relaxed); error!( "Transaction {} commit failed after {}ms: {}", tx_id, elapsed.as_millis(), e ); Err(DatabaseError::from(e)) }, } } /// Rollback the transaction pub async fn rollback(mut self) -> DatabaseResult<()> { let elapsed = self.elapsed(); let tx_id = self.id; let tx = self.inner.take().ok_or_else(|| DatabaseError::Transaction { message: "Transaction already consumed".to_string(), })?; match tx.rollback().await { Ok(_) => { self.manager_stats .active_transactions .fetch_sub(1_u64, Ordering::Relaxed); info!( "Transaction {} rolled back successfully in {}ms", tx_id, elapsed.as_millis() ); Ok(()) }, Err(e) => { self.manager_stats .active_transactions .fetch_sub(1_u64, Ordering::Relaxed); self.manager_stats .failed_transactions .fetch_add(1_u64, Ordering::Relaxed); error!( "Transaction {} rollback failed after {}ms: {}", tx_id, elapsed.as_millis(), e ); Err(DatabaseError::from(e)) }, } } } impl Drop for DatabaseTransaction { fn drop(&mut self) { // Only warn if transaction is still active (not None) if self.inner.is_some() { self.manager_stats .active_transactions .fetch_sub(1_u64, Ordering::Relaxed); warn!("Transaction {} was dropped without explicit commit/rollback - automatic rollback will occur", self.id); } } } /// Transaction statistics #[derive(Debug, Clone)] pub struct TransactionStats { /// Number of currently active transactions pub active_transactions: u64, /// Total number of transactions started pub total_transactions: u64, /// Number of failed transactions pub failed_transactions: u64, /// Number of retry attempts pub retry_attempts: u64, } impl TransactionStats { /// Calculate success rate as percentage pub fn success_rate(&self) -> f64 { if self.total_transactions == 0 { 0.0 } else { ((self.total_transactions - self.failed_transactions) as f64 / self.total_transactions as f64) * 100.0 } } /// Calculate retry rate as percentage pub fn retry_rate(&self) -> f64 { if self.total_transactions == 0 { 0.0 } else { (self.retry_attempts as f64 / self.total_transactions as f64) * 100.0 } } } #[cfg(test)] #[allow(clippy::unwrap_used, clippy::expect_used)] mod tests { use super::*; #[test] fn test_transaction_config_default() { let config = TransactionConfig::default(); assert_eq!(config.default_timeout_secs, 30); assert_eq!(config.max_savepoints, 10); assert!(config.enable_retry); assert_eq!(config.max_retries, 3); } #[test] fn test_transaction_stats_calculation() { let stats = TransactionStats { active_transactions: 2, total_transactions: 100, failed_transactions: 5, retry_attempts: 10, }; assert_eq!(stats.success_rate(), 95.0); assert_eq!(stats.retry_rate(), 10.0); } #[test] fn test_transaction_stats_zero_transactions() { let stats = TransactionStats { active_transactions: 0, total_transactions: 0, failed_transactions: 0, retry_attempts: 0, }; assert_eq!(stats.success_rate(), 0.0); assert_eq!(stats.retry_rate(), 0.0); } }