use crate::error::{DatabaseError, DatabaseResult}; use config::PoolConfig; use sqlx::postgres::{PgPool, PgPoolOptions}; use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; use std::time::Duration; use tokio::sync::RwLock; use tracing::{debug, error, info, warn}; // PoolConfig is now imported from the config crate /// Extension trait for PoolConfig validation trait PoolConfigValidation { fn validate(&self) -> DatabaseResult<()>; } impl PoolConfigValidation for PoolConfig { fn validate(&self) -> DatabaseResult<()> { if self.min_connections > self.max_connections { return Err(DatabaseError::Configuration { message: "min_connections cannot be greater than max_connections".to_string(), }); } if self.max_connections == 0 { return Err(DatabaseError::Configuration { message: "max_connections must be greater than 0".to_string(), }); } if self.acquire_timeout_secs == 0 { return Err(DatabaseError::Configuration { message: "acquire_timeout_secs must be greater than 0".to_string(), }); } if !self.database_url.starts_with("postgres://") && !self.database_url.starts_with("postgresql://") { return Err(DatabaseError::Configuration { message: "database_url must be a valid PostgreSQL connection string".to_string(), }); } Ok(()) } } /// Connection pool statistics #[derive(Debug, Clone)] pub struct PoolStats { /// Total number of connections created pub total_connections_created: u64, /// Number of active connections pub active_connections: u32, /// Number of idle connections pub idle_connections: u32, /// Total number of connection acquisitions pub total_acquisitions: u64, /// Number of failed acquisitions pub failed_acquisitions: u64, /// Number of connections closed due to max lifetime pub connections_closed_max_lifetime: u64, /// Number of connections closed due to idle timeout pub connections_closed_idle_timeout: u64, /// Number of failed health checks pub failed_health_checks: u64, } impl Default for PoolStats { fn default() -> Self { Self { total_connections_created: 0, active_connections: 0, idle_connections: 0, total_acquisitions: 0, failed_acquisitions: 0, connections_closed_max_lifetime: 0, connections_closed_idle_timeout: 0, failed_health_checks: 0, } } } /// Enhanced database connection pool with monitoring and health checks #[derive(Debug)] pub struct DatabasePool { /// The underlying SQLx pool inner: PgPool, /// Pool configuration config: PoolConfig, /// Pool statistics stats: Arc>, /// Atomic counters for thread-safe metrics total_acquisitions: Arc, failed_acquisitions: Arc, total_connections_created: Arc, } impl DatabasePool { /// Create a new database pool with the given configuration pub async fn new(config: PoolConfig) -> DatabaseResult { // Validate configuration config.validate()?; info!( "Creating database pool with {} min connections, {} max connections", config.min_connections, config.max_connections ); // Build the pool let pool = PgPoolOptions::new() .min_connections(config.min_connections) .max_connections(config.max_connections) .acquire_timeout(Duration::from_secs(config.acquire_timeout_secs)) .max_lifetime(Duration::from_secs(config.max_lifetime_secs)) .idle_timeout(Duration::from_secs(config.idle_timeout_secs)) .test_before_acquire(config.test_before_acquire) .connect(&config.database_url) .await .map_err(|e| DatabaseError::ConnectionPool { message: format!("Failed to create connection pool: {}", e), })?; let database_pool = Self { inner: pool, config: config.clone(), stats: Arc::new(RwLock::new(PoolStats::default())), total_acquisitions: Arc::new(AtomicU64::new(0)), failed_acquisitions: Arc::new(AtomicU64::new(0)), total_connections_created: Arc::new(AtomicU64::new(0)), }; // Start health check task if enabled if config.health_check_enabled { database_pool.start_health_check_task().await; } info!("Database pool created successfully"); Ok(database_pool) } /// Get a connection from the pool pub async fn acquire(&self) -> DatabaseResult> { self.total_acquisitions.fetch_add(1, Ordering::Relaxed); debug!("Acquiring connection from pool"); match self.inner.acquire().await { Ok(conn) => { debug!("Successfully acquired connection from pool"); Ok(conn) } Err(e) => { self.failed_acquisitions.fetch_add(1, Ordering::Relaxed); error!("Failed to acquire connection from pool: {}", e); Err(DatabaseError::from(e)) } } } /// Get a reference to the underlying pool for direct use pub fn inner(&self) -> &PgPool { &self.inner } /// Get pool statistics pub async fn stats(&self) -> PoolStats { let mut stats = self.stats.read().await.clone(); // Update atomic counters stats.total_acquisitions = self.total_acquisitions.load(Ordering::Relaxed); stats.failed_acquisitions = self.failed_acquisitions.load(Ordering::Relaxed); stats.total_connections_created = self.total_connections_created.load(Ordering::Relaxed); // Get current pool state stats.active_connections = self.inner.size(); stats.idle_connections = self.inner.num_idle() as u32; stats } /// Check if the pool is healthy pub async fn health_check(&self) -> DatabaseResult { debug!("Performing pool health check"); match sqlx::query("SELECT 1").fetch_one(&self.inner).await { Ok(_) => { debug!("Pool health check passed"); Ok(true) } Err(e) => { warn!("Pool health check failed: {}", e); let mut stats = self.stats.write().await; stats.failed_health_checks += 1; Ok(false) } } } /// Get the pool configuration pub fn config(&self) -> &PoolConfig { &self.config } /// Close the pool gracefully pub async fn close(&self) { info!("Closing database pool"); self.inner.close().await; info!("Database pool closed"); } /// Check if the pool is closed pub fn is_closed(&self) -> bool { self.inner.is_closed() } /// Start the health check background task async fn start_health_check_task(&self) { if !self.config.health_check_enabled { return; } let pool = self.inner.clone(); let stats = self.stats.clone(); let interval = Duration::from_secs(self.config.health_check_interval_secs); tokio::spawn(async move { let mut interval_timer = tokio::time::interval(interval); loop { interval_timer.tick().await; if pool.is_closed() { debug!("Pool is closed, stopping health check task"); break; } match sqlx::query("SELECT 1").fetch_one(&pool).await { Ok(_) => { debug!("Scheduled health check passed"); } Err(e) => { warn!("Scheduled health check failed: {}", e); let mut stats = stats.write().await; stats.failed_health_checks += 1; } } } }); debug!( "Health check task started with {}s interval", self.config.health_check_interval_secs ); } /// Execute a test query to validate the connection pub async fn ping(&self) -> DatabaseResult<()> { sqlx::query("SELECT 1") .fetch_one(&self.inner) .await .map_err(DatabaseError::from)?; Ok(()) } /// Get current pool size pub fn size(&self) -> u32 { self.inner.size() } /// Get number of idle connections pub fn num_idle(&self) -> usize { self.inner.num_idle() } /// Reset pool statistics pub async fn reset_stats(&self) { let mut stats = self.stats.write().await; *stats = PoolStats::default(); self.total_acquisitions.store(0, Ordering::Relaxed); self.failed_acquisitions.store(0, Ordering::Relaxed); self.total_connections_created.store(0, Ordering::Relaxed); info!("Pool statistics reset"); } } impl Clone for DatabasePool { fn clone(&self) -> Self { Self { inner: self.inner.clone(), config: self.config.clone(), stats: self.stats.clone(), total_acquisitions: self.total_acquisitions.clone(), failed_acquisitions: self.failed_acquisitions.clone(), total_connections_created: self.total_connections_created.clone(), } } } #[cfg(test)] mod tests { use super::*; #[test] fn test_pool_config_validation() { let mut config = PoolConfig::default(); assert!(config.validate().is_ok()); // Test min > max config.min_connections = 10; config.max_connections = 5; assert!(config.validate().is_err()); // Test max = 0 config.min_connections = 0; config.max_connections = 0; assert!(config.validate().is_err()); // Test invalid URL config.max_connections = 10; config.database_url = "invalid://url".to_string(); assert!(config.validate().is_err()); } #[test] fn test_pool_config_default() { let config = PoolConfig::default(); assert_eq!(config.min_connections, 5); assert_eq!(config.max_connections, 100); assert!(config.test_before_acquire); assert!(config.health_check_enabled); } #[tokio::test] async fn test_pool_stats_default() { let stats = PoolStats::default(); assert_eq!(stats.total_connections_created, 0); assert_eq!(stats.active_connections, 0); assert_eq!(stats.failed_acquisitions, 0); } }