Skip to main content

ferriorm_runtime/
client.rs

1//! Database connection pool wrapper.
2//!
3//! [`DatabaseClient`] is an enum that wraps either a PostgreSQL or SQLite
4//! connection pool (via sqlx). It provides auto-detection from the connection
5//! URL and exposes typed `fetch_all`, `fetch_optional`, `fetch_one`, and
6//! `execute` helpers used by the generated query builders.
7//!
8//! [`PoolConfig`] allows fine-grained control over the underlying connection
9//! pool (max/min connections, timeouts, etc.).
10
11use std::time::Duration;
12
13use crate::error::FerriormError;
14
15/// Configuration options for the database connection pool.
16///
17/// All fields are optional; when `None`, the sqlx defaults are used.
18///
19/// # Example
20///
21/// ```rust
22/// use ferriorm_runtime::client::PoolConfig;
23/// use std::time::Duration;
24///
25/// let config = PoolConfig {
26///     max_connections: Some(20),
27///     idle_timeout: Some(Duration::from_secs(300)),
28///     ..Default::default()
29/// };
30/// ```
31#[derive(Debug, Clone, Default)]
32pub struct PoolConfig {
33    /// Maximum number of connections in the pool.
34    pub max_connections: Option<u32>,
35    /// Minimum number of connections to keep open at all times.
36    pub min_connections: Option<u32>,
37    /// Maximum time a connection can sit idle before being closed.
38    pub idle_timeout: Option<Duration>,
39    /// Maximum lifetime of a connection before it is closed and replaced.
40    pub max_lifetime: Option<Duration>,
41    /// Maximum time to wait when acquiring a connection from the pool.
42    pub acquire_timeout: Option<Duration>,
43}
44
45/// The database client, wrapping an sqlx connection pool.
46///
47/// Supports PostgreSQL and SQLite via feature flags.
48#[derive(Debug, Clone)]
49pub enum DatabaseClient {
50    #[cfg(feature = "postgres")]
51    Postgres(sqlx::PgPool),
52    #[cfg(feature = "sqlite")]
53    Sqlite(sqlx::SqlitePool),
54}
55
56impl DatabaseClient {
57    /// Connect to a PostgreSQL database.
58    #[cfg(feature = "postgres")]
59    pub async fn connect_postgres(url: &str) -> Result<Self, FerriormError> {
60        let pool = sqlx::PgPool::connect(url).await?;
61        Ok(Self::Postgres(pool))
62    }
63
64    /// Connect to a PostgreSQL database with custom pool configuration.
65    #[cfg(feature = "postgres")]
66    pub async fn connect_postgres_with_config(
67        url: &str,
68        config: &PoolConfig,
69    ) -> Result<Self, FerriormError> {
70        let mut opts = sqlx::postgres::PgPoolOptions::new();
71        if let Some(max) = config.max_connections {
72            opts = opts.max_connections(max);
73        }
74        if let Some(min) = config.min_connections {
75            opts = opts.min_connections(min);
76        }
77        if let Some(timeout) = config.idle_timeout {
78            opts = opts.idle_timeout(timeout);
79        }
80        if let Some(lifetime) = config.max_lifetime {
81            opts = opts.max_lifetime(lifetime);
82        }
83        if let Some(timeout) = config.acquire_timeout {
84            opts = opts.acquire_timeout(timeout);
85        }
86        let pool = opts.connect(url).await?;
87        Ok(Self::Postgres(pool))
88    }
89
90    /// Connect to a SQLite database.
91    #[cfg(feature = "sqlite")]
92    pub async fn connect_sqlite(url: &str) -> Result<Self, FerriormError> {
93        let url = normalize_sqlite_url(url);
94        let pool = sqlx::SqlitePool::connect(&url).await?;
95        Ok(Self::Sqlite(pool))
96    }
97
98    /// Connect to a SQLite database with custom pool configuration.
99    #[cfg(feature = "sqlite")]
100    pub async fn connect_sqlite_with_config(
101        url: &str,
102        config: &PoolConfig,
103    ) -> Result<Self, FerriormError> {
104        let url = normalize_sqlite_url(url);
105        let mut opts = sqlx::sqlite::SqlitePoolOptions::new();
106        if let Some(max) = config.max_connections {
107            opts = opts.max_connections(max);
108        }
109        if let Some(min) = config.min_connections {
110            opts = opts.min_connections(min);
111        }
112        if let Some(timeout) = config.idle_timeout {
113            opts = opts.idle_timeout(timeout);
114        }
115        if let Some(lifetime) = config.max_lifetime {
116            opts = opts.max_lifetime(lifetime);
117        }
118        if let Some(timeout) = config.acquire_timeout {
119            opts = opts.acquire_timeout(timeout);
120        }
121        let pool = opts.connect(&url).await?;
122        Ok(Self::Sqlite(pool))
123    }
124
125    /// Connect by auto-detecting the database type from the URL.
126    pub async fn connect(url: &str) -> Result<Self, FerriormError> {
127        #[cfg(feature = "sqlite")]
128        if url.starts_with("sqlite:") || url.starts_with("file:") || url.ends_with(".db") {
129            return Self::connect_sqlite(url).await;
130        }
131
132        #[cfg(feature = "postgres")]
133        {
134            return Self::connect_postgres(url).await;
135        }
136
137        #[allow(unreachable_code)]
138        Err(FerriormError::Connection(
139            "No database backend enabled. Enable 'postgres' or 'sqlite' feature.".into(),
140        ))
141    }
142
143    /// Connect by auto-detecting the database type from the URL, using custom
144    /// pool configuration.
145    pub async fn connect_with_config(
146        url: &str,
147        config: &PoolConfig,
148    ) -> Result<Self, FerriormError> {
149        #[cfg(feature = "sqlite")]
150        if url.starts_with("sqlite:") || url.starts_with("file:") || url.ends_with(".db") {
151            return Self::connect_sqlite_with_config(url, config).await;
152        }
153
154        #[cfg(feature = "postgres")]
155        {
156            return Self::connect_postgres_with_config(url, config).await;
157        }
158
159        #[allow(unreachable_code)]
160        Err(FerriormError::Connection(
161            "No database backend enabled. Enable 'postgres' or 'sqlite' feature.".into(),
162        ))
163    }
164
165    /// Get a reference to the underlying PostgreSQL pool for raw queries.
166    ///
167    /// Returns an error if this client is not connected to PostgreSQL.
168    #[cfg(feature = "postgres")]
169    pub fn pg_pool(&self) -> Result<&sqlx::PgPool, FerriormError> {
170        match self {
171            Self::Postgres(pool) => Ok(pool),
172            #[cfg(feature = "sqlite")]
173            _ => Err(FerriormError::Connection(
174                "Expected PostgreSQL connection".into(),
175            )),
176        }
177    }
178
179    /// Get a reference to the underlying SQLite pool for raw queries.
180    ///
181    /// Returns an error if this client is not connected to SQLite.
182    #[cfg(feature = "sqlite")]
183    pub fn sqlite_pool(&self) -> Result<&sqlx::SqlitePool, FerriormError> {
184        match self {
185            Self::Sqlite(pool) => Ok(pool),
186            #[cfg(feature = "postgres")]
187            _ => Err(FerriormError::Connection(
188                "Expected SQLite connection".into(),
189            )),
190        }
191    }
192
193    /// Close the connection pool.
194    pub async fn disconnect(self) {
195        match self {
196            #[cfg(feature = "postgres")]
197            Self::Postgres(pool) => pool.close().await,
198            #[cfg(feature = "sqlite")]
199            Self::Sqlite(pool) => pool.close().await,
200        }
201    }
202
203    // ── Raw SQL helpers (PostgreSQL) ────────────────────────────────────
204
205    /// Execute raw SQL and return all rows mapped to `T` (PostgreSQL).
206    #[cfg(feature = "postgres")]
207    pub async fn raw_fetch_all_pg<T>(&self, sql: &str) -> Result<Vec<T>, FerriormError>
208    where
209        T: for<'r> sqlx::FromRow<'r, sqlx::postgres::PgRow> + Send + Unpin,
210    {
211        match self {
212            Self::Postgres(pool) => Ok(sqlx::query_as::<_, T>(sql).fetch_all(pool).await?),
213            #[cfg(feature = "sqlite")]
214            _ => Err(FerriormError::Query(
215                "Expected PostgreSQL connection".into(),
216            )),
217        }
218    }
219
220    /// Execute raw SQL and return exactly one row mapped to `T` (PostgreSQL).
221    #[cfg(feature = "postgres")]
222    pub async fn raw_fetch_one_pg<T>(&self, sql: &str) -> Result<T, FerriormError>
223    where
224        T: for<'r> sqlx::FromRow<'r, sqlx::postgres::PgRow> + Send + Unpin,
225    {
226        match self {
227            Self::Postgres(pool) => Ok(sqlx::query_as::<_, T>(sql).fetch_one(pool).await?),
228            #[cfg(feature = "sqlite")]
229            _ => Err(FerriormError::Query(
230                "Expected PostgreSQL connection".into(),
231            )),
232        }
233    }
234
235    /// Execute raw SQL and return an optional row mapped to `T` (PostgreSQL).
236    #[cfg(feature = "postgres")]
237    pub async fn raw_fetch_optional_pg<T>(&self, sql: &str) -> Result<Option<T>, FerriormError>
238    where
239        T: for<'r> sqlx::FromRow<'r, sqlx::postgres::PgRow> + Send + Unpin,
240    {
241        match self {
242            Self::Postgres(pool) => Ok(sqlx::query_as::<_, T>(sql).fetch_optional(pool).await?),
243            #[cfg(feature = "sqlite")]
244            _ => Err(FerriormError::Query(
245                "Expected PostgreSQL connection".into(),
246            )),
247        }
248    }
249
250    /// Execute raw SQL without returning rows (PostgreSQL). Returns the
251    /// number of rows affected.
252    #[cfg(feature = "postgres")]
253    pub async fn raw_execute_pg(&self, sql: &str) -> Result<u64, FerriormError> {
254        match self {
255            Self::Postgres(pool) => Ok(sqlx::query(sql).execute(pool).await?.rows_affected()),
256            #[cfg(feature = "sqlite")]
257            _ => Err(FerriormError::Query(
258                "Expected PostgreSQL connection".into(),
259            )),
260        }
261    }
262
263    // ── Raw SQL helpers (SQLite) ────────────────────────────────────────
264
265    /// Execute raw SQL and return all rows mapped to `T` (SQLite).
266    #[cfg(feature = "sqlite")]
267    pub async fn raw_fetch_all_sqlite<T>(&self, sql: &str) -> Result<Vec<T>, FerriormError>
268    where
269        T: for<'r> sqlx::FromRow<'r, sqlx::sqlite::SqliteRow> + Send + Unpin,
270    {
271        match self {
272            Self::Sqlite(pool) => Ok(sqlx::query_as::<_, T>(sql).fetch_all(pool).await?),
273            #[cfg(feature = "postgres")]
274            _ => Err(FerriormError::Query("Expected SQLite connection".into())),
275        }
276    }
277
278    /// Execute raw SQL and return exactly one row mapped to `T` (SQLite).
279    #[cfg(feature = "sqlite")]
280    pub async fn raw_fetch_one_sqlite<T>(&self, sql: &str) -> Result<T, FerriormError>
281    where
282        T: for<'r> sqlx::FromRow<'r, sqlx::sqlite::SqliteRow> + Send + Unpin,
283    {
284        match self {
285            Self::Sqlite(pool) => Ok(sqlx::query_as::<_, T>(sql).fetch_one(pool).await?),
286            #[cfg(feature = "postgres")]
287            _ => Err(FerriormError::Query("Expected SQLite connection".into())),
288        }
289    }
290
291    /// Execute raw SQL and return an optional row mapped to `T` (SQLite).
292    #[cfg(feature = "sqlite")]
293    pub async fn raw_fetch_optional_sqlite<T>(&self, sql: &str) -> Result<Option<T>, FerriormError>
294    where
295        T: for<'r> sqlx::FromRow<'r, sqlx::sqlite::SqliteRow> + Send + Unpin,
296    {
297        match self {
298            Self::Sqlite(pool) => Ok(sqlx::query_as::<_, T>(sql).fetch_optional(pool).await?),
299            #[cfg(feature = "postgres")]
300            _ => Err(FerriormError::Query("Expected SQLite connection".into())),
301        }
302    }
303
304    /// Execute raw SQL without returning rows (SQLite). Returns the number
305    /// of rows affected.
306    #[cfg(feature = "sqlite")]
307    pub async fn raw_execute_sqlite(&self, sql: &str) -> Result<u64, FerriormError> {
308        match self {
309            Self::Sqlite(pool) => Ok(sqlx::query(sql).execute(pool).await?.rows_affected()),
310            #[cfg(feature = "postgres")]
311            _ => Err(FerriormError::Query("Expected SQLite connection".into())),
312        }
313    }
314
315    /// Execute a query builder against the appropriate pool, returning all rows.
316    #[cfg(feature = "postgres")]
317    pub async fn fetch_all_pg<'q, T>(
318        &self,
319        mut qb: sqlx::QueryBuilder<'q, sqlx::Postgres>,
320    ) -> Result<Vec<T>, FerriormError>
321    where
322        T: for<'r> sqlx::FromRow<'r, sqlx::postgres::PgRow> + Send + Unpin,
323    {
324        match self {
325            Self::Postgres(pool) => Ok(qb.build_query_as::<T>().fetch_all(pool).await?),
326            #[cfg(feature = "sqlite")]
327            _ => Err(FerriormError::Query(
328                "Expected PostgreSQL connection".into(),
329            )),
330        }
331    }
332
333    #[cfg(feature = "postgres")]
334    pub async fn fetch_optional_pg<'q, T>(
335        &self,
336        mut qb: sqlx::QueryBuilder<'q, sqlx::Postgres>,
337    ) -> Result<Option<T>, FerriormError>
338    where
339        T: for<'r> sqlx::FromRow<'r, sqlx::postgres::PgRow> + Send + Unpin,
340    {
341        match self {
342            Self::Postgres(pool) => Ok(qb.build_query_as::<T>().fetch_optional(pool).await?),
343            #[cfg(feature = "sqlite")]
344            _ => Err(FerriormError::Query(
345                "Expected PostgreSQL connection".into(),
346            )),
347        }
348    }
349
350    #[cfg(feature = "postgres")]
351    pub async fn fetch_one_pg<'q, T>(
352        &self,
353        mut qb: sqlx::QueryBuilder<'q, sqlx::Postgres>,
354    ) -> Result<T, FerriormError>
355    where
356        T: for<'r> sqlx::FromRow<'r, sqlx::postgres::PgRow> + Send + Unpin,
357    {
358        match self {
359            Self::Postgres(pool) => Ok(qb.build_query_as::<T>().fetch_one(pool).await?),
360            #[cfg(feature = "sqlite")]
361            _ => Err(FerriormError::Query(
362                "Expected PostgreSQL connection".into(),
363            )),
364        }
365    }
366
367    #[cfg(feature = "postgres")]
368    pub async fn execute_pg<'q>(
369        &self,
370        mut qb: sqlx::QueryBuilder<'q, sqlx::Postgres>,
371    ) -> Result<u64, FerriormError> {
372        match self {
373            Self::Postgres(pool) => Ok(qb.build().execute(pool).await?.rows_affected()),
374            #[cfg(feature = "sqlite")]
375            _ => Err(FerriormError::Query(
376                "Expected PostgreSQL connection".into(),
377            )),
378        }
379    }
380
381    // SQLite variants
382    #[cfg(feature = "sqlite")]
383    pub async fn fetch_all_sqlite<'q, T>(
384        &self,
385        mut qb: sqlx::QueryBuilder<'q, sqlx::Sqlite>,
386    ) -> Result<Vec<T>, FerriormError>
387    where
388        T: for<'r> sqlx::FromRow<'r, sqlx::sqlite::SqliteRow> + Send + Unpin,
389    {
390        match self {
391            Self::Sqlite(pool) => Ok(qb.build_query_as::<T>().fetch_all(pool).await?),
392            #[cfg(feature = "postgres")]
393            _ => Err(FerriormError::Query("Expected SQLite connection".into())),
394        }
395    }
396
397    #[cfg(feature = "sqlite")]
398    pub async fn fetch_optional_sqlite<'q, T>(
399        &self,
400        mut qb: sqlx::QueryBuilder<'q, sqlx::Sqlite>,
401    ) -> Result<Option<T>, FerriormError>
402    where
403        T: for<'r> sqlx::FromRow<'r, sqlx::sqlite::SqliteRow> + Send + Unpin,
404    {
405        match self {
406            Self::Sqlite(pool) => Ok(qb.build_query_as::<T>().fetch_optional(pool).await?),
407            #[cfg(feature = "postgres")]
408            _ => Err(FerriormError::Query("Expected SQLite connection".into())),
409        }
410    }
411
412    #[cfg(feature = "sqlite")]
413    pub async fn fetch_one_sqlite<'q, T>(
414        &self,
415        mut qb: sqlx::QueryBuilder<'q, sqlx::Sqlite>,
416    ) -> Result<T, FerriormError>
417    where
418        T: for<'r> sqlx::FromRow<'r, sqlx::sqlite::SqliteRow> + Send + Unpin,
419    {
420        match self {
421            Self::Sqlite(pool) => Ok(qb.build_query_as::<T>().fetch_one(pool).await?),
422            #[cfg(feature = "postgres")]
423            _ => Err(FerriormError::Query("Expected SQLite connection".into())),
424        }
425    }
426
427    #[cfg(feature = "sqlite")]
428    pub async fn execute_sqlite<'q>(
429        &self,
430        mut qb: sqlx::QueryBuilder<'q, sqlx::Sqlite>,
431    ) -> Result<u64, FerriormError> {
432        match self {
433            Self::Sqlite(pool) => Ok(qb.build().execute(pool).await?.rows_affected()),
434            #[cfg(feature = "postgres")]
435            _ => Err(FerriormError::Query("Expected SQLite connection".into())),
436        }
437    }
438}
439
440/// Normalize a SQLite connection URL for sqlx.
441///
442/// Converts `file:` prefixed URLs (e.g. `file:./dev.db`) to the `sqlite:`
443/// scheme that sqlx expects, and appends `?mode=rwc` so the database file
444/// is auto-created if it does not exist.
445#[cfg(feature = "sqlite")]
446pub fn normalize_sqlite_url(url: &str) -> String {
447    let url = if let Some(path) = url.strip_prefix("file:") {
448        format!("sqlite:{}", path)
449    } else if !url.starts_with("sqlite:") {
450        format!("sqlite:{}", url)
451    } else {
452        url.to_string()
453    };
454    // Ensure mode=rwc for auto-creation
455    if !url.contains("mode=") {
456        if url.contains('?') {
457            format!("{}&mode=rwc", url)
458        } else {
459            format!("{}?mode=rwc", url)
460        }
461    } else {
462        url
463    }
464}