Skip to main content

kipuka/db/
mod.rs

1//! Database access layer.
2//!
3//! Provides connection pool initialization, migration execution, and
4//! query helpers for the three supported backends: SQLite, PostgreSQL,
5//! and MariaDB.
6//!
7//! The shared connection pool type [`Db`] is a runtime-dispatch `AnyPool`
8//! that routes queries to the configured backend.  All write transactions
9//! on SQLite should go through [`begin_write`] rather than `pool.begin()`
10//! to avoid `SQLITE_BUSY_SNAPSHOT` in WAL mode.
11
12pub mod schema;
13
14use crate::config::DbConfig;
15use crate::error::KipukaError;
16
17/// Type alias for the shared connection pool (runtime-dispatch Any backend).
18pub type Db = sqlx::AnyPool;
19
20/// Which database backend is active.
21///
22/// Drives `begin_write` only: SQLite requires `BEGIN IMMEDIATE` to prevent
23/// `SQLITE_BUSY_SNAPSHOT` (error 517) in WAL mode; PostgreSQL and MariaDB use
24/// standard MVCC and need only `BEGIN`.
25#[derive(Clone, Copy, Debug, PartialEq, Eq)]
26pub enum DbKind {
27    Sqlite,
28    Postgres,
29    MariaDb,
30}
31
32impl DbKind {
33    /// Infer the backend from the connection URL scheme.
34    pub fn from_url(url: &str) -> Self {
35        if url.starts_with("postgres") {
36            DbKind::Postgres
37        } else if url.starts_with("mariadb") || url.starts_with("mysql") {
38            DbKind::MariaDb
39        } else {
40            DbKind::Sqlite
41        }
42    }
43}
44
45/// Global PostgreSQL flag for parameter rewriting.
46///
47/// sqlx 0.8's `AnyPool` does not reliably rewrite `?` parameter placeholders
48/// to `$N` for PostgreSQL.  Call [`pg_sql`] on every SQL string that contains
49/// `?` before passing it to sqlx.
50static IS_POSTGRES: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
51
52/// Rewrite `?` → `$1`, `$2`, … for PostgreSQL; return unchanged for others.
53///
54/// The rewritten string is cached permanently (via `Box::leak`) keyed by the
55/// pointer identity of the static string literal, so each unique query string
56/// is rewritten at most once.
57#[allow(dead_code)]
58pub fn pg_sql(s: &'static str) -> &'static str {
59    if !IS_POSTGRES.get().copied().unwrap_or(false) {
60        return s;
61    }
62    static CACHE: std::sync::OnceLock<
63        std::sync::Mutex<std::collections::HashMap<usize, &'static str>>,
64    > = std::sync::OnceLock::new();
65    let cache = CACHE.get_or_init(|| std::sync::Mutex::new(std::collections::HashMap::new()));
66    let key = s.as_ptr() as usize;
67    {
68        let guard = cache.lock().unwrap();
69        if let Some(&cached) = guard.get(&key) {
70            return cached;
71        }
72    }
73    // Slow path: rewrite then store for the lifetime of the process.
74    let mut result = String::with_capacity(s.len() + 16);
75    let mut param_num = 0u32;
76    for ch in s.chars() {
77        if ch == '?' {
78            param_num += 1;
79            result.push('$');
80            result.push_str(&param_num.to_string());
81        } else {
82            result.push(ch);
83        }
84    }
85    let leaked: &'static str = Box::leak(result.into_boxed_str());
86    cache.lock().unwrap().insert(key, leaked);
87    leaked
88}
89
90/// Rewrite `?` → `$1`, `$2`, … for PostgreSQL on a dynamically-built SQL
91/// string.  Unlike [`pg_sql`], which operates on `&'static str` literals,
92/// this accepts an owned `String` and returns an owned `String`.
93///
94/// When the backend is not PostgreSQL the input is returned unchanged.
95#[allow(dead_code)]
96pub(crate) fn pg_sql_dynamic(s: String) -> String {
97    if !IS_POSTGRES.get().copied().unwrap_or(false) {
98        return s;
99    }
100    let mut result = String::with_capacity(s.len() + 16);
101    let mut param_num = 0u32;
102    for ch in s.chars() {
103        if ch == '?' {
104            param_num += 1;
105            result.push('$');
106            result.push_str(&param_num.to_string());
107        } else {
108            result.push(ch);
109        }
110    }
111    result
112}
113
114/// Initialize the primary (read-write) database connection pool.
115///
116/// The `resolved_url` is the pre-resolved database URL from `SecretResolver`.
117pub async fn init_pool(config: &DbConfig, resolved_url: &str) -> Result<(Db, DbKind), KipukaError> {
118    // The sqlx `Any` driver requires the backend drivers to be registered
119    // before any pool can connect.  This is idempotent (guarded internally by
120    // a `Once`), so calling it here makes `init_pool` self-sufficient for every
121    // caller — the binary, embedded users, and integration tests that never run
122    // `main` — instead of relying on a separate setup call elsewhere.
123    sqlx::any::install_default_drivers();
124
125    let url = resolved_url.to_string();
126
127    let kind = DbKind::from_url(&url);
128    let _ = IS_POSTGRES.set(kind == DbKind::Postgres);
129
130    let pool_opts = sqlx::any::AnyPoolOptions::new()
131        .acquire_timeout(std::time::Duration::from_secs(config.connect_timeout_secs))
132        .max_lifetime(std::time::Duration::from_secs(config.max_lifetime_secs));
133
134    let in_memory =
135        kind == DbKind::Sqlite && (url.contains(":memory:") || url.contains("mode=memory"));
136    let pool_opts = if in_memory {
137        // An ephemeral database must retain its one schema-bearing connection.
138        pool_opts
139            .max_connections(1)
140            .min_connections(1)
141            .max_lifetime(None)
142            .idle_timeout(None)
143    } else if let Some(max) = config.max_connections {
144        pool_opts.max_connections(max)
145    } else {
146        match kind {
147            DbKind::Sqlite => pool_opts.max_connections(1),
148            _ => pool_opts.max_connections(10),
149        }
150    };
151
152    let pool_opts = if let Some(min) = config.min_connections.filter(|_| !in_memory) {
153        pool_opts.min_connections(min)
154    } else {
155        pool_opts
156    };
157
158    let pool = pool_opts
159        .connect(&url)
160        .await
161        .map_err(|e| KipukaError::Db(format!("failed to connect to database: {e}")))?;
162
163    // Enable WAL mode for SQLite
164    if kind == DbKind::Sqlite && config.sqlite_wal {
165        sqlx::query("PRAGMA journal_mode=WAL")
166            .execute(&pool)
167            .await
168            .map_err(|e| KipukaError::Db(format!("failed to enable WAL mode: {e}")))?;
169    }
170
171    Ok((pool, kind))
172}
173
174/// Initialize a read-only connection pool for GET handlers.
175///
176/// For SQLite file-backed databases, this opens a `?mode=ro` pool that
177/// never acquires the write lock, enabling concurrent reads during writes
178/// (WAL concurrency benefit).  For `:memory:` and non-SQLite backends,
179/// returns a clone of the primary pool.
180pub async fn init_ro_pool(
181    primary: &Db,
182    _config: &DbConfig,
183    kind: DbKind,
184    resolved_url: &str,
185) -> Result<Db, KipukaError> {
186    let url = resolved_url.to_string();
187
188    // Only SQLite file-backed databases benefit from a separate RO pool
189    if kind != DbKind::Sqlite || (url.contains(":memory:") || url.contains("mode=memory")) {
190        return Ok(primary.clone());
191    }
192
193    // Build a read-only URL for SQLite
194    let ro_url = if url.contains('?') {
195        format!("{url}&mode=ro")
196    } else {
197        format!("{url}?mode=ro")
198    };
199
200    let pool = sqlx::any::AnyPoolOptions::new()
201        .max_connections(4)
202        .connect(&ro_url)
203        .await
204        .map_err(|e| KipukaError::Db(format!("failed to connect RO pool: {e}")))?;
205
206    Ok(pool)
207}
208
209/// Begin a write transaction.
210///
211/// SQLite uses `BEGIN IMMEDIATE` to avoid `SQLITE_BUSY_SNAPSHOT` under WAL
212/// mode; other backends use standard `BEGIN`.
213pub async fn begin_write(
214    pool: &Db,
215    kind: DbKind,
216) -> Result<sqlx::Transaction<'_, sqlx::Any>, KipukaError> {
217    let tx = if kind == DbKind::Sqlite {
218        pool.begin_with("BEGIN IMMEDIATE").await
219    } else {
220        pool.begin().await
221    }
222    .map_err(|e| KipukaError::Db(format!("begin transaction failed: {e}")))?;
223    Ok(tx)
224}
225
226/// Run pending database migrations.
227pub async fn run_migrations(pool: &Db, kind: DbKind) -> Result<(), KipukaError> {
228    crate::db::schema::run_migrations(pool, kind).await
229}