//! Database connection and management //! //! This module provides database connection pooling and utilities for //! interacting with the PostgreSQL database. use sqlx::postgres::{PgPool, PgPoolOptions}; use std::time::Duration; use tracing::{info, warn}; use crate::config::DatabaseConfig; use crate::error::Result; /// Database connection pool #[derive(Debug, Clone)] pub struct Database { pool: PgPool, schema: String, } impl Database { /// Create a new database connection from configuration pub async fn new(config: &DatabaseConfig) -> Result { // Default to "attune" schema for production safety let schema = config .schema .clone() .unwrap_or_else(|| "attune".to_string()); // Validate schema name (prevent SQL injection) Self::validate_schema_name(&schema)?; // Log schema configuration prominently info!( "Connecting to database with max_connections={}, schema={}", config.max_connections, schema ); // Clone schema for use in closure let schema_for_hook = schema.clone(); let pool = PgPoolOptions::new() .max_connections(config.max_connections) .min_connections(config.min_connections) .acquire_timeout(Duration::from_secs(config.connect_timeout)) .idle_timeout(Duration::from_secs(config.idle_timeout)) .after_connect(move |conn, _meta| { let schema = schema_for_hook.clone(); Box::pin(async move { // Set search_path for every connection in the pool // Only include 'public' for production schemas (attune), not test schemas // This ensures test schemas have isolated migrations tables let search_path = if schema.starts_with("test_") { format!("SET search_path TO {}", schema) } else { format!("SET search_path TO {}, public", schema) }; sqlx::query(&search_path).execute(&mut *conn).await?; Ok(()) }) }) .connect(&config.url) .await?; // Run a test query to verify connection sqlx::query("SELECT 1").execute(&pool).await.map_err(|e| { warn!("Failed to verify database connection: {}", e); e })?; info!("Successfully connected to database"); Ok(Self { pool, schema }) } /// Validate schema name to prevent SQL injection fn validate_schema_name(schema: &str) -> Result<()> { if schema.is_empty() { return Err(crate::error::Error::Configuration( "Schema name cannot be empty".to_string(), )); } // Only allow alphanumeric and underscores if !schema.chars().all(|c| c.is_alphanumeric() || c == '_') { return Err(crate::error::Error::Configuration(format!( "Invalid schema name '{}': only alphanumeric and underscores allowed", schema ))); } // Prevent excessively long names (PostgreSQL limit is 63 chars) if schema.len() > 63 { return Err(crate::error::Error::Configuration(format!( "Schema name '{}' too long (max 63 characters)", schema ))); } Ok(()) } /// Get a reference to the connection pool pub fn pool(&self) -> &PgPool { &self.pool } /// Get the current schema name pub fn schema(&self) -> &str { &self.schema } /// Close the database connection pool pub async fn close(&self) { self.pool.close().await; info!("Database connection pool closed"); } /// Run database migrations /// Note: Migrations should be in the workspace root migrations directory pub async fn migrate(&self) -> Result<()> { info!("Running database migrations"); // TODO: Implement migrations when migration files are created // sqlx::migrate!("../../migrations") // .run(&self.pool) // .await?; info!("Database migrations will be implemented with migration files"); Ok(()) } /// Check if the database connection is healthy pub async fn health_check(&self) -> Result<()> { sqlx::query("SELECT 1").execute(&self.pool).await?; Ok(()) } /// Get pool statistics pub fn stats(&self) -> PoolStats { PoolStats { connections: self.pool.size(), idle_connections: self.pool.num_idle(), } } } /// Database pool statistics #[derive(Debug, Clone)] pub struct PoolStats { pub connections: u32, pub idle_connections: usize, } #[cfg(test)] mod tests { use super::*; #[test] fn test_pool_stats() { // Test that PoolStats can be created let stats = PoolStats { connections: 10, idle_connections: 5, }; assert_eq!(stats.connections, 10); assert_eq!(stats.idle_connections, 5); } }