reinhardt_db/pool/

pool.rs

//! Connection pool implementation

use super::config::PoolConfig;
use super::errors::{PoolError, PoolResult};
use super::events::{PoolEvent, PoolEventListener};
use sqlx::{Database, MySql, Pool, Postgres, Sqlite};
use std::mem::ManuallyDrop;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::{Arc, OnceLock};
use tokio::sync::RwLock;

/// Mask the password in a database URL for safe display.
///
/// Handles standard URL formats like `scheme://user:password@host/db`
/// and replaces the password portion with `***`.
/// Correctly handles passwords containing `@` by using the last `@` as
/// the user-info delimiter.
pub(crate) fn mask_url_password(url: &str) -> String {
	// Try to parse as a standard URL with scheme://user:pass@host format
	if let Some(scheme_end) = url.find("://") {
		let after_scheme = &url[scheme_end + 3..];

		// Use the last @ as the user-info delimiter, since passwords may contain @
		if let Some(at_pos) = after_scheme.rfind('@') {
			let user_info = &after_scheme[..at_pos];

			// Find the first colon separating user from password
			if let Some(colon_pos) = user_info.find(':') {
				let scheme_and_user = &url[..scheme_end + 3 + colon_pos + 1];
				let rest = &url[scheme_end + 3 + at_pos..];
				return format!("{}***{}", scheme_and_user, rest);
			}
		}
	}

	// No password found, return as-is
	url.to_string()
}

fn generate_connection_id() -> String {
	uuid::Uuid::now_v7().to_string()
}

struct PoolEventHub {
	listeners: RwLock<Vec<Arc<dyn PoolEventListener>>>,
	listener_count: AtomicUsize,
}

impl PoolEventHub {
	fn new() -> Self {
		Self {
			listeners: RwLock::new(Vec::new()),
			listener_count: AtomicUsize::new(0),
		}
	}

	fn has_listeners(&self) -> bool {
		self.listener_count.load(Ordering::Acquire) > 0
	}

	async fn add_listener(&self, listener: Arc<dyn PoolEventListener>) {
		let mut listeners = self.listeners.write().await;
		listeners.push(listener);
		self.listener_count
			.store(listeners.len(), Ordering::Release);
	}

	async fn emit_event(&self, event: PoolEvent) {
		if !self.has_listeners() {
			return;
		}

		let listeners = self.listeners.read().await;
		for listener in listeners.iter() {
			listener.on_event(event.clone()).await;
		}
	}
}

/// A database connection pool
pub struct ConnectionPool<DB: Database> {
	pool: Pool<DB>,
	config: PoolConfig,
	url: String,
	events: Arc<PoolEventHub>,
	first_connect_fired: Arc<AtomicBool>,
}

impl ConnectionPool<Postgres> {
	/// Create a new PostgreSQL connection pool
	///
	/// # Examples
	///
	/// ```
	/// use reinhardt_db::pool::{ConnectionPool, PoolConfig};
	///
	/// # async fn example() {
	/// let config = PoolConfig::default();
	/// // For doctest purposes, using SQLite in-memory instead of PostgreSQL
	/// let pool = ConnectionPool::new_sqlite("sqlite::memory:", config).await.unwrap();
	/// assert!(pool.url().contains("memory"));
	/// assert_eq!(pool.config().max_connections, 10);
	/// # }
	/// # tokio::runtime::Runtime::new().unwrap().block_on(example());
	/// ```
	pub async fn new_postgres(url: &str, config: PoolConfig) -> PoolResult<Self> {
		config.validate().map_err(PoolError::Config)?;

		let pool = sqlx::postgres::PgPoolOptions::new()
			.min_connections(config.min_connections)
			.max_connections(config.max_connections)
			.acquire_timeout(config.acquire_timeout)
			.idle_timeout(config.idle_timeout)
			.max_lifetime(config.max_lifetime)
			.test_before_acquire(config.test_before_acquire)
			.connect(url)
			.await?;

		Ok(Self {
			pool,
			config,
			url: url.to_string(),
			events: Arc::new(PoolEventHub::new()),
			first_connect_fired: Arc::new(AtomicBool::new(false)),
		})
	}
}

impl ConnectionPool<MySql> {
	/// Create a new MySQL connection pool
	///
	/// # Examples
	///
	/// ```
	/// use reinhardt_db::pool::{ConnectionPool, PoolConfig};
	///
	/// # async fn example() {
	/// let config = PoolConfig::default();
	/// // For doctest purposes, using SQLite in-memory instead of MySQL
	/// let pool = ConnectionPool::new_sqlite("sqlite::memory:", config).await.unwrap();
	/// assert!(pool.url().contains("memory"));
	/// assert_eq!(pool.config().max_connections, 10);
	/// # }
	/// # tokio::runtime::Runtime::new().unwrap().block_on(example());
	/// ```
	pub async fn new_mysql(url: &str, config: PoolConfig) -> PoolResult<Self> {
		config.validate().map_err(PoolError::Config)?;

		let pool = sqlx::mysql::MySqlPoolOptions::new()
			.min_connections(config.min_connections)
			.max_connections(config.max_connections)
			.acquire_timeout(config.acquire_timeout)
			.idle_timeout(config.idle_timeout)
			.max_lifetime(config.max_lifetime)
			.test_before_acquire(config.test_before_acquire)
			.connect(url)
			.await?;

		Ok(Self {
			pool,
			config,
			url: url.to_string(),
			events: Arc::new(PoolEventHub::new()),
			first_connect_fired: Arc::new(AtomicBool::new(false)),
		})
	}
}

impl ConnectionPool<Sqlite> {
	/// Create a new SQLite connection pool
	///
	/// # Examples
	///
	/// ```
	/// use reinhardt_db::pool::{ConnectionPool, PoolConfig};
	///
	/// # async fn example() {
	/// let config = PoolConfig::default();
	/// // Using in-memory SQLite for doctest
	/// let pool = ConnectionPool::new_sqlite("sqlite::memory:", config).await.unwrap();
	/// assert!(pool.url().contains("memory"));
	/// assert!(pool.config().max_connections > 0);
	/// # }
	/// # tokio::runtime::Runtime::new().unwrap().block_on(example());
	/// ```
	pub async fn new_sqlite(url: &str, config: PoolConfig) -> PoolResult<Self> {
		config.validate().map_err(PoolError::Config)?;

		let pool = sqlx::sqlite::SqlitePoolOptions::new()
			.min_connections(config.min_connections)
			.max_connections(config.max_connections)
			.acquire_timeout(config.acquire_timeout)
			.idle_timeout(config.idle_timeout)
			.max_lifetime(config.max_lifetime)
			.test_before_acquire(config.test_before_acquire)
			.connect(url)
			.await?;

		Ok(Self {
			pool,
			config,
			url: url.to_string(),
			events: Arc::new(PoolEventHub::new()),
			first_connect_fired: Arc::new(AtomicBool::new(false)),
		})
	}
}

impl<DB> ConnectionPool<DB>
where
	DB: sqlx::Database,
{
	/// Add an event listener
	///
	pub async fn add_listener(&self, listener: Arc<dyn PoolEventListener>) {
		self.events.add_listener(listener).await;
	}

	/// Emit an event to all listeners
	pub(crate) async fn emit_event(&self, event: PoolEvent) {
		self.events.emit_event(event).await;
	}
	/// Acquire a connection from the pool with event emission
	///
	/// # Examples
	///
	/// ```no_run
	/// use reinhardt_db::pool::{ConnectionPool, PoolConfig};
	///
	/// # async fn example() {
	/// let config = PoolConfig::default();
	/// let pool = ConnectionPool::new_postgres("postgresql://user:pass@localhost/test", config)
	///     .await
	///     .unwrap();
	///
	/// // Acquire a connection
	/// let conn = pool.acquire().await;
	/// assert!(conn.is_ok());
	/// # }
	/// ```
	pub async fn acquire(&self) -> PoolResult<PooledConnection<DB>> {
		// Check if this is the first connection
		let is_first = !self.first_connect_fired.swap(true, Ordering::SeqCst);

		let conn = self.pool.acquire().await?;
		let connection_id = OnceLock::new();

		if self.events.has_listeners() {
			let id = connection_id.get_or_init(generate_connection_id).clone();

			if is_first {
				// Emit first_connect event (using ConnectionCreated as proxy)
				self.emit_event(PoolEvent::connection_created(id.clone()))
					.await;
			}

			// Emit checkout event
			self.emit_event(PoolEvent::connection_acquired(id)).await;
		}

		Ok(PooledConnection {
			conn: ManuallyDrop::new(conn),
			events: self.events.clone(),
			connection_id,
		})
	}
	/// Get the underlying pool
	///
	pub fn inner(&self) -> &Pool<DB> {
		&self.pool
	}
	/// Get pool configuration
	///
	pub fn config(&self) -> &PoolConfig {
		&self.config
	}
	/// Close the pool
	///
	/// Attempts to gracefully close the pool with a 5-second timeout.
	/// If active connections are not returned within this time, the pool
	/// will be forcefully closed.
	pub async fn close(&self) {
		use tokio::time::{Duration, timeout};

		// Try to close gracefully with a timeout
		let close_future = self.pool.close();
		if timeout(Duration::from_secs(5), close_future).await.is_err() {
			// Timeout occurred - pool had active connections
			// The pool will be forcefully closed when dropped
		}
	}
	/// Get the database URL with password masked for safe display
	///
	/// Returns the database URL with any password replaced by `***`
	/// to prevent credential exposure in logs and debug output.
	/// Use `url_raw()` when the actual password is needed for reconnection.
	pub fn url(&self) -> String {
		mask_url_password(&self.url)
	}

	/// Get the raw database URL including credentials
	///
	/// This method returns the unmasked URL containing the actual password.
	/// Use with caution - prefer `url()` for logging and display purposes.
	// Allow dead_code: preserved for internal use by reconnection logic (e.g., `recreate()`)
	#[allow(dead_code)]
	pub(crate) fn url_raw(&self) -> &str {
		&self.url
	}
}

// Database-specific recreate implementations
impl ConnectionPool<Postgres> {
	/// Recreate the pool with the same configuration
	///
	/// # Examples
	///
	/// ```
	/// use reinhardt_db::pool::{ConnectionPool, PoolConfig};
	///
	/// # async fn example() {
	/// let config = PoolConfig::default();
	/// // For doctest purposes, using SQLite in-memory instead of PostgreSQL
	/// let mut pool = ConnectionPool::new_sqlite("sqlite::memory:", config)
	///     .await
	///     .unwrap();
	///
	/// // Recreate the pool
	/// pool.recreate().await.unwrap();
	/// assert_eq!(pool.config().max_connections, 10);
	/// # }
	/// # tokio::runtime::Runtime::new().unwrap().block_on(example());
	/// ```
	pub async fn recreate(&mut self) -> PoolResult<()> {
		// Close existing pool
		self.pool.close().await;

		// Create new pool with same configuration
		let new_pool = sqlx::postgres::PgPoolOptions::new()
			.min_connections(self.config.min_connections)
			.max_connections(self.config.max_connections)
			.acquire_timeout(self.config.acquire_timeout)
			.idle_timeout(self.config.idle_timeout)
			.max_lifetime(self.config.max_lifetime)
			.test_before_acquire(self.config.test_before_acquire)
			.connect(&self.url)
			.await?;

		self.pool = new_pool;
		self.first_connect_fired.store(false, Ordering::SeqCst);

		Ok(())
	}
}

impl ConnectionPool<MySql> {
	/// Recreate the pool with the same configuration
	///
	/// # Examples
	///
	/// ```
	/// use reinhardt_db::pool::{ConnectionPool, PoolConfig};
	///
	/// # async fn example() {
	/// let config = PoolConfig::default();
	/// // For doctest purposes, using SQLite in-memory instead of MySQL
	/// let mut pool = ConnectionPool::new_sqlite("sqlite::memory:", config)
	///     .await
	///     .unwrap();
	///
	/// // Recreate the pool
	/// pool.recreate().await.unwrap();
	/// assert_eq!(pool.config().max_connections, 10);
	/// # }
	/// # tokio::runtime::Runtime::new().unwrap().block_on(example());
	/// ```
	pub async fn recreate(&mut self) -> PoolResult<()> {
		// Close existing pool
		self.pool.close().await;

		// Create new pool with same configuration
		let new_pool = sqlx::mysql::MySqlPoolOptions::new()
			.min_connections(self.config.min_connections)
			.max_connections(self.config.max_connections)
			.acquire_timeout(self.config.acquire_timeout)
			.idle_timeout(self.config.idle_timeout)
			.max_lifetime(self.config.max_lifetime)
			.test_before_acquire(self.config.test_before_acquire)
			.connect(&self.url)
			.await?;

		self.pool = new_pool;
		self.first_connect_fired.store(false, Ordering::SeqCst);

		Ok(())
	}
}

impl ConnectionPool<Sqlite> {
	/// Recreate the pool with the same configuration
	///
	/// # Examples
	///
	/// ```
	/// use reinhardt_db::pool::{ConnectionPool, PoolConfig};
	///
	/// # async fn example() {
	/// let config = PoolConfig::default();
	/// let mut pool = ConnectionPool::new_sqlite("sqlite::memory:", config)
	///     .await
	///     .unwrap();
	///
	/// // Recreate the pool
	/// pool.recreate().await.unwrap();
	/// assert!(pool.url().contains("memory"));
	/// # }
	/// # tokio::runtime::Runtime::new().unwrap().block_on(example());
	/// ```
	pub async fn recreate(&mut self) -> PoolResult<()> {
		// Close existing pool
		self.pool.close().await;

		// Create new pool with same configuration
		let new_pool = sqlx::sqlite::SqlitePoolOptions::new()
			.min_connections(self.config.min_connections)
			.max_connections(self.config.max_connections)
			.acquire_timeout(self.config.acquire_timeout)
			.idle_timeout(self.config.idle_timeout)
			.max_lifetime(self.config.max_lifetime)
			.test_before_acquire(self.config.test_before_acquire)
			.connect(&self.url)
			.await?;

		self.pool = new_pool;
		self.first_connect_fired.store(false, Ordering::SeqCst);

		Ok(())
	}
}

/// A pooled connection wrapper with event emission
pub struct PooledConnection<DB: sqlx::Database> {
	// Wrapped in ManuallyDrop so we can take ownership in Drop.
	// When no tokio runtime is available, we detach the connection
	// to avoid sqlx's PoolConnection::Drop calling rt::spawn().
	conn: ManuallyDrop<sqlx::pool::PoolConnection<DB>>,
	events: Arc<PoolEventHub>,
	connection_id: OnceLock<String>,
}

impl<DB: sqlx::Database> PooledConnection<DB> {
	/// Documentation for `inner`
	///
	pub fn inner(&mut self) -> &mut sqlx::pool::PoolConnection<DB> {
		&mut self.conn
	}
	/// Get the unique identifier for this connection
	///
	/// # Examples
	///
	/// ```no_run
	/// use reinhardt_db::pool::{ConnectionPool, PoolConfig};
	///
	/// # async fn example() {
	/// let config = PoolConfig::default();
	/// let pool = ConnectionPool::new_postgres("postgresql://user:pass@localhost/test", config)
	///     .await
	///     .unwrap();
	///
	/// let mut conn = pool.acquire().await.unwrap();
	/// let id = conn.connection_id();
	/// assert!(!id.is_empty());
	/// # }
	/// ```
	pub fn connection_id(&self) -> &str {
		self.connection_id
			.get_or_init(generate_connection_id)
			.as_str()
	}
	/// Invalidate this connection (hard invalidation - connection is unusable)
	///
	pub async fn invalidate(self, reason: String) {
		if self.events.has_listeners() {
			let connection_id = self
				.connection_id
				.get_or_init(generate_connection_id)
				.clone();
			self.events
				.emit_event(PoolEvent::connection_invalidated(connection_id, reason))
				.await;
		}
		// Connection will be dropped and not returned to pool
	}
	/// Soft invalidate this connection (can complete current operation)
	///
	pub async fn soft_invalidate(&mut self) {
		if self.events.has_listeners() {
			let connection_id = self
				.connection_id
				.get_or_init(generate_connection_id)
				.clone();
			self.events
				.emit_event(PoolEvent::connection_soft_invalidated(connection_id))
				.await;
		}
	}
	/// Reset this connection
	///
	pub async fn reset(&mut self) {
		if self.events.has_listeners() {
			let connection_id = self
				.connection_id
				.get_or_init(generate_connection_id)
				.clone();
			self.events
				.emit_event(PoolEvent::connection_reset(connection_id))
				.await;
		}
	}
}

impl<DB: sqlx::Database> Drop for PooledConnection<DB> {
	fn drop(&mut self) {
		// SAFETY: ManuallyDrop::take is called exactly once (in drop).
		let conn = unsafe { ManuallyDrop::take(&mut self.conn) };

		match tokio::runtime::Handle::try_current() {
			Ok(handle) => {
				// Runtime available: drop the connection normally (returns to pool)
				// and emit the connection-returned event.
				drop(conn);

				if self.events.has_listeners() {
					let events = self.events.clone();
					let connection_id = self
						.connection_id
						.get_or_init(generate_connection_id)
						.clone();

					handle.spawn(async move {
						events
							.emit_event(PoolEvent::connection_returned(connection_id))
							.await;
					});
				}
			}
			Err(_) => {
				// No runtime available: prevent sqlx's PoolConnection::Drop
				// from running, as it calls crate::rt::spawn() which panics
				// without a tokio runtime. The connection is intentionally
				// leaked to avoid the panic.
				std::mem::forget(conn);
			}
		}
	}
}

#[cfg(test)]
mod tests {
	use super::*;
	use rstest::rstest;
	use std::sync::Mutex;

	struct RecordingListener {
		events: Arc<Mutex<Vec<&'static str>>>,
	}

	#[async_trait::async_trait]
	impl PoolEventListener for RecordingListener {
		async fn on_event(&self, event: PoolEvent) {
			let name = match event {
				PoolEvent::ConnectionAcquired { .. } => "acquired",
				PoolEvent::ConnectionReturned { .. } => "returned",
				PoolEvent::ConnectionCreated { .. } => "created",
				PoolEvent::ConnectionClosed { .. } => "closed",
				PoolEvent::ConnectionTestFailed { .. } => "test_failed",
				PoolEvent::ConnectionInvalidated { .. } => "invalidated",
				PoolEvent::ConnectionSoftInvalidated { .. } => "soft_invalidated",
				PoolEvent::ConnectionReset { .. } => "reset",
			};
			self.events
				.lock()
				.expect("events mutex should not be poisoned")
				.push(name);
		}
	}

	#[rstest]
	#[case(
		"postgresql://user:secret@localhost:5432/mydb",
		"postgresql://user:***@localhost:5432/mydb"
	)]
	#[case(
		"mysql://admin:p@ssw0rd@db.example.com/app",
		"mysql://admin:***@db.example.com/app"
	)]
	#[case(
		"postgres://user:pass@host:5432/db?sslmode=require",
		"postgres://user:***@host:5432/db?sslmode=require"
	)]
	fn test_mask_url_password_with_credentials(#[case] input: &str, #[case] expected: &str) {
		// Arrange
		// (input provided by case parameters)

		// Act
		let masked = mask_url_password(input);

		// Assert
		assert_eq!(masked, expected);
	}

	#[rstest]
	#[case("sqlite::memory:")]
	#[case("sqlite:///path/to/db.sqlite")]
	#[case("postgresql://user@localhost:5432/mydb")]
	fn test_mask_url_password_without_password(#[case] input: &str) {
		// Arrange
		// (input provided by case parameter)

		// Act
		let masked = mask_url_password(input);

		// Assert
		assert_eq!(masked, input, "URL without password should be unchanged");
	}

	#[rstest]
	fn test_mask_url_password_empty_password() {
		// Arrange
		let url = "postgresql://user:@localhost:5432/mydb";

		// Act
		let masked = mask_url_password(url);

		// Assert
		assert_eq!(masked, "postgresql://user:***@localhost:5432/mydb");
	}

	#[rstest]
	fn test_mask_url_password_special_chars_in_password() {
		// Arrange
		let url = "postgresql://user:p%40ss%3Aw0rd@localhost:5432/mydb";

		// Act
		let masked = mask_url_password(url);

		// Assert
		assert_eq!(masked, "postgresql://user:***@localhost:5432/mydb");
		assert!(
			!masked.contains("p%40ss"),
			"Password should be fully masked"
		);
	}

	#[rstest]
	fn test_mask_url_password_preserves_non_url() {
		// Arrange
		let non_url = "not-a-url-just-a-string";

		// Act
		let masked = mask_url_password(non_url);

		// Assert
		assert_eq!(
			masked, non_url,
			"Non-URL strings should pass through unchanged"
		);
	}

	#[rstest]
	fn test_handle_try_current_returns_err_outside_runtime() {
		// Arrange & Act & Assert
		// Run on a fresh thread to avoid inheriting runtime context
		// from the test runner's worker thread.
		let handle = std::thread::spawn(|| {
			let result = tokio::runtime::Handle::try_current();
			assert!(
				result.is_err(),
				"Handle::try_current() should return Err outside of a tokio runtime"
			);
		});
		handle.join().expect("thread should not panic");
	}

	#[rstest]
	fn test_drop_pooled_connection_outside_runtime_does_not_panic() {
		// Arrange
		// Create a Tokio runtime and acquire a pooled connection inside it.
		let rt = tokio::runtime::Runtime::new().expect("failed to create Tokio runtime");

		let (pool, conn) = rt.block_on(async {
			let config = PoolConfig::default();
			let pool = ConnectionPool::new_sqlite("sqlite::memory:", config)
				.await
				.expect("failed to create ConnectionPool");

			let conn = pool.acquire().await.expect("failed to acquire connection");

			(pool, conn)
		});

		// Drop the runtime so there is no active Tokio runtime.
		drop(rt);

		// Act & Assert
		// Dropping the connection outside any runtime should not panic.
		drop(conn);

		// Also drop the pool to ensure cleanup does not panic outside a runtime.
		drop(pool);
	}

	#[tokio::test]
	async fn test_pool_events_are_emitted_when_listener_registered() {
		// Arrange
		let events = Arc::new(Mutex::new(Vec::new()));
		let listener = Arc::new(RecordingListener {
			events: events.clone(),
		});
		let pool = ConnectionPool::new_sqlite("sqlite::memory:", PoolConfig::default())
			.await
			.expect("failed to create ConnectionPool");
		pool.add_listener(listener).await;

		// Act
		let conn = pool.acquire().await.expect("failed to acquire connection");
		drop(conn);
		tokio::task::yield_now().await;

		// Assert
		let recorded = events.lock().expect("events mutex should not be poisoned");
		assert_eq!(
			recorded.as_slice(),
			["created", "acquired", "returned"],
			"pool listener should observe the first acquire and return events"
		);
	}
}