Skip to main content

arc_web/helpers/
database.rs

1use crate::helpers::config;
2use diesel::r2d2::{ConnectionManager, Pool, PooledConnection};
3use diesel::SqliteConnection;
4use diesel_migrations::{embed_migrations, EmbeddedMigrations};
5use std::sync::{OnceLock, RwLock};
6
7/// Compatibility migration set for framework consumers.
8///
9/// Application/domain migrations belong to the consuming app and should be
10/// embedded there. This constant remains for source compatibility with 0.2.2.
11#[deprecated(note = "embed and run migrations from the consuming application")]
12pub const MIGRATIONS: EmbeddedMigrations = embed_migrations!("migrations");
13
14/// Internal state for the connection pool singleton.
15struct PoolState {
16    pool: Pool<ConnectionManager<SqliteConnection>>,
17    database_url: String,
18}
19
20static POOL: OnceLock<RwLock<Option<PoolState>>> = OnceLock::new();
21
22/// Get a database connection from the shared connection pool.
23/// The pool is created lazily on first call and reused for subsequent calls.
24/// If the `DATABASE_URL` changes (e.g., between test runs), the pool is recreated.
25pub fn get_connection() -> PooledConnection<ConnectionManager<SqliteConnection>> {
26    let current_db_url = config::database_url();
27
28    let pool_state = POOL.get_or_init(|| RwLock::new(Some(create_pool_state(&current_db_url))));
29
30    // Check if we need to recreate the pool (e.g., for tests with different DB or after reset)
31    {
32        let state = pool_state.read().unwrap();
33        if let Some(ref ps) = *state {
34            if ps.database_url == current_db_url {
35                return ps.pool.get().expect("Failed to get connection from pool");
36            }
37        }
38    }
39
40    // Database URL changed or pool was reset, recreate pool
41    {
42        let mut state = pool_state.write().unwrap();
43        let needs_recreate = match *state {
44            None => true,
45            Some(ref ps) => ps.database_url != current_db_url,
46        };
47        if needs_recreate {
48            *state = Some(create_pool_state(&current_db_url));
49        }
50        state
51            .as_ref()
52            .unwrap()
53            .pool
54            .get()
55            .expect("Failed to get connection from pool")
56    }
57}
58
59/// Get a clone of the shared connection pool.
60/// Useful when you need to pass the pool to a different context.
61#[allow(dead_code)]
62pub fn get_connection_pool() -> Pool<ConnectionManager<SqliteConnection>> {
63    let current_db_url = config::database_url();
64
65    let pool_state = POOL.get_or_init(|| RwLock::new(Some(create_pool_state(&current_db_url))));
66
67    {
68        let state = pool_state.read().unwrap();
69        if let Some(ref ps) = *state {
70            if ps.database_url == current_db_url {
71                return ps.pool.clone();
72            }
73        }
74    }
75
76    {
77        let mut state = pool_state.write().unwrap();
78        let needs_recreate = match *state {
79            None => true,
80            Some(ref ps) => ps.database_url != current_db_url,
81        };
82        if needs_recreate {
83            *state = Some(create_pool_state(&current_db_url));
84        }
85        state.as_ref().unwrap().pool.clone()
86    }
87}
88
89/// Resets the connection pool. Call this when the database file is deleted/recreated.
90/// This is primarily used in tests after TestFinalizer deletes the database.
91#[cfg(any(test, feature = "test-utils"))]
92pub fn reset_pool() {
93    if let Some(pool_state) = POOL.get() {
94        let mut state = pool_state.write().unwrap();
95        *state = None;
96    }
97}
98
99fn create_pool_state(database_url: &str) -> PoolState {
100    let pool_limit: u32 = config::database_pool_limit();
101
102    let manager = ConnectionManager::<SqliteConnection>::new(database_url);
103
104    let pool = Pool::builder()
105        .max_size(pool_limit)
106        .test_on_check_out(true)
107        .build(manager)
108        .expect("Could not build connection pool");
109
110    PoolState {
111        pool,
112        database_url: database_url.to_string(),
113    }
114}