1use std::time::Duration;
12
13use crate::error::FerriormError;
14
15#[derive(Debug, Clone, Default)]
32pub struct PoolConfig {
33 pub max_connections: Option<u32>,
35 pub min_connections: Option<u32>,
37 pub idle_timeout: Option<Duration>,
39 pub max_lifetime: Option<Duration>,
41 pub acquire_timeout: Option<Duration>,
43}
44
45#[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 #[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 #[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 #[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 #[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 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 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 #[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 #[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 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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 #[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#[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 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}