use crate::error::{DatabaseError, DatabaseResult}; use config::database::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 /// /// Provides validation methods for pool configuration to ensure /// settings are valid before creating a connection pool. trait PoolConfigValidation { /// Validate the pool configuration settings /// /// # Errors /// /// Returns `DatabaseError::Configuration` if any validation fails: /// - `min_connections` is greater than `max_connections` /// - `max_connections` is 0 /// - `acquire_timeout_secs` is 0 /// - `database_url` is not a valid PostgreSQL connection string 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 /// /// Provides comprehensive metrics about the database connection pool /// including usage patterns, health status, and performance indicators. #[derive(Debug, Clone, Default)] pub struct PoolStats { /// Total number of connections created since pool initialization pub total_connections_created: u64, /// Number of currently active connections in use pub active_connections: u32, /// Number of idle connections available for use pub idle_connections: u32, /// Total number of connection acquisitions requested pub total_acquisitions: u64, /// Number of failed connection acquisition attempts pub failed_acquisitions: u64, /// Number of connections closed due to reaching maximum 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 check attempts pub failed_health_checks: u64, } /// Enhanced database connection pool with monitoring and health checks /// /// Provides a high-level interface to PostgreSQL connection pooling with: /// - Automatic connection lifecycle management /// - Health monitoring and statistics collection /// - Configurable timeouts and pool sizing /// - Background health check tasks #[derive(Debug)] pub struct DatabasePool { /// The underlying SQLx PostgreSQL pool inner: PgPool, /// Pool configuration settings config: PoolConfig, /// Pool statistics protected by async RwLock stats: Arc>, /// Atomic counter for total connection acquisitions total_acquisitions: Arc, /// Atomic counter for failed connection acquisitions failed_acquisitions: Arc, /// Atomic counter for total connections created total_connections_created: Arc, } impl DatabasePool { /// Create a new database pool with the given configuration /// /// # Arguments /// /// * `config` - Pool configuration including connection limits, timeouts, and database URL /// /// # Returns /// /// A new `DatabasePool` instance ready for use /// /// # Errors /// /// Returns `DatabaseError::Configuration` if the configuration is invalid /// or `DatabaseError::ConnectionPool` if the pool cannot be created /// /// # Examples /// /// ```rust,no_run /// use database::{DatabasePool, PoolConfig}; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// let config = PoolConfig::default(); /// let pool = DatabasePool::new(config).await?; /// Ok(()) /// } /// ``` 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 /// /// Acquires a connection from the pool, waiting up to the configured /// acquire timeout if no connections are immediately available. /// /// # Returns /// /// A pooled connection that will be returned to the pool when dropped /// /// # Errors /// /// Returns `DatabaseError::Timeout` if the acquire timeout is exceeded /// or other connection-related errors from the underlying pool /// /// # Examples /// /// ```rust,no_run /// # use database::{DatabasePool, PoolConfig}; /// # #[tokio::main] /// # async fn main() -> Result<(), Box> { /// # let pool = DatabasePool::new(PoolConfig::default()).await?; /// let conn = pool.acquire().await?; /// // Use connection... /// // Connection automatically returned to pool when dropped /// # Ok(()) /// # } /// ``` 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 /// /// Provides direct access to the SQLx pool for operations that /// require the underlying pool interface. /// /// # Returns /// /// A reference to the underlying `PgPool` /// /// # Usage /// /// This method is typically used for executing queries directly /// without acquiring a connection explicitly. pub fn inner(&self) -> &PgPool { &self.inner } /// Get current pool statistics /// /// Returns a snapshot of the current pool state including /// connection counts, acquisition metrics, and health status. /// /// # Returns /// /// Current pool statistics including active/idle connections, /// total acquisitions, failed attempts, and health check results /// /// # Examples /// /// ```rust,no_run /// # use database::{DatabasePool, PoolConfig}; /// # #[tokio::main] /// # async fn main() -> Result<(), Box> { /// # let pool = DatabasePool::new(PoolConfig::default()).await?; /// let stats = pool.stats().await; /// println!("Active connections: {}", stats.active_connections); /// println!("Failed acquisitions: {}", stats.failed_acquisitions); /// # Ok(()) /// # } /// ``` 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 /// /// Performs a simple query to verify the database connection /// and pool functionality. Updates health check statistics. /// /// # Returns /// /// `true` if the health check passes, `false` otherwise /// /// # Errors /// /// This method does not return errors for failed health checks, /// instead it returns `false` and increments the failed health /// check counter in the statistics. /// /// # Examples /// /// ```rust,no_run /// # use database::{DatabasePool, PoolConfig}; /// # #[tokio::main] /// # async fn main() -> Result<(), Box> { /// # let pool = DatabasePool::new(PoolConfig::default()).await?; /// if pool.health_check().await? { /// println!("Database is healthy"); /// } else { /// println!("Database health check failed"); /// } /// # Ok(()) /// # } /// ``` 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 /// /// Returns a reference to the configuration used to create this pool. /// /// # Returns /// /// A reference to the `PoolConfig` used during pool creation pub fn config(&self) -> &PoolConfig { &self.config } /// Close the pool gracefully /// /// Closes all connections in the pool and prevents new connections /// from being created. This is an irreversible operation. /// /// # Examples /// /// ```rust,no_run /// # use database::{DatabasePool, PoolConfig}; /// # #[tokio::main] /// # async fn main() -> Result<(), Box> { /// # let pool = DatabasePool::new(PoolConfig::default()).await?; /// // Use pool for operations... /// pool.close().await; /// assert!(pool.is_closed()); /// # Ok(()) /// # } /// ``` pub async fn close(&self) { info!("Closing database pool"); self.inner.close().await; info!("Database pool closed"); } /// Check if the pool is closed /// /// Returns `true` if the pool has been closed and can no longer /// create new connections. /// /// # Returns /// /// `true` if the pool is closed, `false` if it's still active pub fn is_closed(&self) -> bool { self.inner.is_closed() } /// Start the health check background task /// /// Spawns a background task that periodically performs health checks /// on the database connection pool. The task runs until the pool is closed. /// /// The health check interval is configured via the pool configuration. /// Failed health checks are recorded in the pool statistics. 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 /// /// Performs a simple SELECT 1 query to verify database connectivity. /// This is similar to a health check but returns an error on failure. /// /// # Returns /// /// `Ok(())` if the ping succeeds /// /// # Errors /// /// Returns a `DatabaseError` if the ping query fails /// /// # Examples /// /// ```rust,no_run /// # use database::{DatabasePool, PoolConfig}; /// # #[tokio::main] /// # async fn main() -> Result<(), Box> { /// # let pool = DatabasePool::new(PoolConfig::default()).await?; /// pool.ping().await?; /// println!("Database connection is active"); /// # Ok(()) /// # } /// ``` pub async fn ping(&self) -> DatabaseResult<()> { sqlx::query("SELECT 1") .fetch_one(&self.inner) .await .map_err(DatabaseError::from)?; Ok(()) } /// Get current pool size /// /// Returns the total number of connections currently managed /// by the pool (both active and idle). /// /// # Returns /// /// The current total number of connections in the pool pub fn size(&self) -> u32 { self.inner.size() } /// Get number of idle connections /// /// Returns the number of connections that are currently idle /// and available for use. /// /// # Returns /// /// The number of idle connections available in the pool pub fn num_idle(&self) -> usize { self.inner.num_idle() } /// Reset pool statistics /// /// Resets all statistical counters to zero. This includes /// acquisition counts, failure counts, and other metrics. /// /// # Examples /// /// ```rust,no_run /// # use database::{DatabasePool, PoolConfig}; /// # #[tokio::main] /// # async fn main() -> Result<(), Box> { /// # let pool = DatabasePool::new(PoolConfig::default()).await?; /// // After some operations... /// pool.reset_stats().await; /// let stats = pool.stats().await; /// assert_eq!(stats.total_acquisitions, 0); /// # Ok(()) /// # } /// ``` 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); } }