Skip to main content

rust_ef_sqlite/
pool_strategy.rs

1//! SQLite connection management strategy.
2//!
3//! Separates pool/memory dispatch from the public provider so `provider.rs`
4//! only contains `SqliteProvider` and its trait impls.
5use r2d2::CustomizeConnection;
6use r2d2_sqlite::SqliteConnectionManager;
7use rusqlite::Connection;
8use rust_ef::error::{EFError, EFResult};
9use std::sync::Arc;
10use tokio::sync::Mutex;
11
12/// Initializes each pooled connection with WAL mode and a busy timeout.
13///
14/// WAL allows concurrent readers while a writer holds the lock; the busy
15/// timeout makes writers wait up to 5 seconds for `SQLITE_BUSY` to clear
16/// instead of failing immediately. Applied to every connection at acquisition
17/// time so pool growth doesn't lose the PRAGMAs.
18#[derive(Debug)]
19pub(crate) struct SqliteConnectionCustomizer;
20
21impl CustomizeConnection<Connection, rusqlite::Error> for SqliteConnectionCustomizer {
22    fn on_acquire(&self, conn: &mut Connection) -> Result<(), rusqlite::Error> {
23        conn.execute_batch("PRAGMA journal_mode=WAL; PRAGMA busy_timeout=5000;")
24    }
25}
26
27/// Default max pool size for file-based SQLite databases.
28pub(crate) const SQLITE_DEFAULT_POOL_SIZE: u32 = 8;
29
30/// The connection management strategy.
31///
32/// `Pooled` is used for file-based databases — r2d2 maintains a pool of
33/// connections that can be checked out concurrently.
34///
35/// `Single` is used for `:memory:` databases — SQLite `:memory:` databases
36/// are per-connection, so a single shared connection (behind a `Mutex`) is
37/// the only way to keep all operations on the same database. This matches
38/// the pre-v1.4 behavior and preserves test isolation (each
39/// `new_in_memory()` call gets a fresh, independent database).
40pub(crate) enum SqliteProviderInner {
41    Pooled(r2d2::Pool<SqliteConnectionManager>),
42    Single(Arc<Mutex<rusqlite::Connection>>),
43}
44
45impl SqliteProviderInner {
46    pub(crate) async fn get_connection(
47        &self,
48    ) -> EFResult<Box<dyn rust_ef::provider::IAsyncConnection>> {
49        match self {
50            SqliteProviderInner::Pooled(pool) => {
51                let conn = pool.get().map_err(|e| {
52                    EFError::connection(format!("SQLite pool acquire failed: {}", e))
53                })?;
54                Ok(Box::new(crate::connection::SqliteConnection::new_pooled(
55                    conn,
56                )))
57            }
58            SqliteProviderInner::Single(conn) => Ok(Box::new(
59                crate::connection::SqliteConnection::new_shared(Arc::clone(conn)),
60            )),
61        }
62    }
63
64    pub(crate) async fn execute_migration_command(&self, sql: &str) -> EFResult<()> {
65        match self {
66            SqliteProviderInner::Pooled(pool) => {
67                let conn = pool.get().map_err(|e| {
68                    EFError::connection(format!("SQLite pool acquire failed: {}", e))
69                })?;
70                conn.execute_batch(sql).map_err(|e| {
71                    EFError::migration(format!("Migration execution failed: {}", e))
72                })?;
73                Ok(())
74            }
75            SqliteProviderInner::Single(conn) => {
76                let guard = conn.lock().await;
77                guard.execute_batch(sql).map_err(|e| {
78                    EFError::migration(format!("Migration execution failed: {}", e))
79                })?;
80                Ok(())
81            }
82        }
83    }
84}