zirv_db/
lib.rs

1use diesel::mysql::MysqlConnection;
2use diesel::r2d2::{ConnectionManager, Pool, PooledConnection};
3use std::env;
4use std::sync::OnceLock;
5
6static GLOBAL_POOL: OnceLock<Pool<ConnectionManager<MysqlConnection>>> = OnceLock::new();
7
8pub struct DB;
9
10impl DB {
11    /// Initialize the global pool with optional overrides.
12    ///
13    /// # Errors
14    /// Returns an error if the pool cannot be created, or if the pool is already initialized.
15    pub fn init_pool() -> Result<(), Box<dyn std::error::Error>> {
16        // Read from environment if not supplied
17        let db_url = env::var("DATABASE_URL")
18                .expect("DATABASE_URL must be set in the environment or passed explicitly");
19
20        // Create a connection manager
21        let manager = ConnectionManager::<MysqlConnection>::new(db_url);
22
23        // Build the pool with a configurable max size
24        let pool_size = match env::var("DATABASE_POOL_SIZE") {
25            Ok(size) => size.parse::<u32>()?,
26            Err(_) => 10,
27        };
28
29        let pool = Pool::builder()
30            .max_size(pool_size)
31            .build(manager)?;
32
33        // Initialize the global pool (only once)
34        GLOBAL_POOL
35            .set(pool)
36            .map_err(|_| "Pool is already initialized!")?;
37
38        Ok(())
39    }
40
41    /// Retrieve a pooled connection.
42    ///
43    /// # Returns
44    /// A `PooledConnection` if successful, or an error if the pool is not initialized
45    /// or cannot allocate a new connection.
46    pub fn get_conn() 
47        -> Result<PooledConnection<ConnectionManager<MysqlConnection>>, String> 
48    {
49        // Safely access the global pool
50        let pool = GLOBAL_POOL
51            .get()
52            .ok_or("Couldn't retrieve connection from database pool")?;
53        
54        // Get a connection from the pool
55        match pool.get() {
56            Ok(conn) => Ok(conn),
57            Err(e) => Err(format!("Failed to get connection from pool: {}", e)),
58        }
59    }
60}
61
62#[cfg(test)]
63mod tests {
64    use super::*;
65    use std::env;
66
67    /// Helper function to reset environment variables used in tests.
68    fn clear_env_vars() {
69        // Unset any vars that might conflict with test expectations.
70        // Use `remove_var` to ensure a clean environment for each test.
71        env::remove_var("DATABASE_URL");
72        env::remove_var("DATABASE_POOL_SIZE");
73    }
74
75    #[test]
76    fn test_init_pool_and_get_conn() {
77        clear_env_vars();
78
79        // Set environment variables to valid values for the test.
80        // Replace this URL with a valid MySQL connection string pointing to your test database.
81        env::set_var("DATABASE_URL", "mysql://root:password@localhost/zirv-db-test");
82        env::set_var("DATABASE_POOL_SIZE", "2");
83
84        // Initialize the global pool.
85        let init_result = DB::init_pool();
86        assert!(
87            init_result.is_ok(),
88            "Expected successful pool initialization, got error: {:?}",
89            init_result.err()
90        );
91
92        // Retrieve a connection.
93        let conn_result = DB::get_conn();
94        assert!(
95            conn_result.is_ok(),
96            "Expected successful connection retrieval, got error: {:?}",
97            conn_result.err()
98        );
99    }
100
101    #[test]
102    fn test_get_conn_without_init() {
103        clear_env_vars();
104        // We do not set DATABASE_URL or call DB::init_pool()
105
106        // Attempting to get a connection without initializing the pool should fail.
107        let conn_result = DB::get_conn();
108        assert!(
109            conn_result.is_err(),
110            "Expected get_conn to fail without initialization, but it succeeded."
111        );
112    }
113}
zirv_db/lib.rs

zirv_db/
lib.rs
