Skip to main content

floz_orm/db/
pool.rs

1//! `Db` — polymorphic database connection pool.
2//!
3//! Wraps either a PostgreSQL or SQLite pool behind a single type.
4//! The backend is auto-detected from the `DATABASE_URL` scheme.
5//!
6//! ```ignore
7//! let db = Db::connect("postgres://user:pass@localhost/mydb").await?;
8//! let db = Db::connect("sqlite::memory:").await?;
9//! ```
10
11#[cfg(feature = "postgres")]
12use sqlx::postgres::{PgPool, PgPoolOptions};
13#[cfg(feature = "sqlite")]
14use sqlx::sqlite::{SqlitePool, SqlitePoolOptions};
15
16use tokio::sync::Mutex;
17
18use crate::error::FlozError;
19use super::placeholder::replace_pg_placeholders;
20use super::tx::{Tx, TxInner};
21
22/// The inner pool enum — one variant per enabled backend.
23#[derive(Debug, Clone)]
24pub(crate) enum DbInner {
25    #[cfg(feature = "postgres")]
26    Postgres(PgPool),
27    #[cfg(feature = "sqlite")]
28    Sqlite(SqlitePool),
29}
30
31/// A database connection pool that auto-selects the backend based on
32/// the connection URL scheme.
33///
34/// `Db` is cheap to clone (internally reference-counted) and safe to share
35/// across tokio tasks.
36#[derive(Debug, Clone)]
37pub struct Db {
38    pub(crate) inner: DbInner,
39}
40
41impl Db {
42    /// Connect to a database, auto-detecting the backend from the URL scheme.
43    ///
44    /// - `postgres://...` or `postgresql://...` → PostgreSQL
45    /// - `sqlite://...` or `sqlite::memory:` → SQLite
46    pub async fn connect(url: &str) -> Result<Self, FlozError> {
47        #[cfg(feature = "sqlite")]
48        if url.starts_with("sqlite:") {
49            let pool = SqlitePoolOptions::new()
50                .max_connections(5)
51                .connect(url)
52                .await?;
53            return Ok(Self {
54                inner: DbInner::Sqlite(pool),
55            });
56        }
57
58        #[cfg(feature = "postgres")]
59        if url.starts_with("postgres:") || url.starts_with("postgresql:") {
60            let pool = PgPoolOptions::new()
61                .max_connections(10)
62                .connect(url)
63                .await?;
64            return Ok(Self {
65                inner: DbInner::Postgres(pool),
66            });
67        }
68
69        Err(FlozError::Database(sqlx::Error::Configuration(
70            format!(
71                "Unsupported database URL scheme: {}. Enable the `postgres` or `sqlite` feature.",
72                url
73            )
74            .into(),
75        )))
76    }
77
78    /// Connect using an environment variable, falling back to a default URL.
79    pub async fn connect_env(env_var: &str, default_url: &str) -> Result<Self, FlozError> {
80        let url = std::env::var(env_var).unwrap_or_else(|_| default_url.to_string());
81        Self::connect(&url).await
82    }
83
84    /// Wrap an existing PostgreSQL pool.
85    #[cfg(feature = "postgres")]
86    pub fn from_pg_pool(pool: PgPool) -> Self {
87        Self {
88            inner: DbInner::Postgres(pool),
89        }
90    }
91
92    /// Wrap an existing SQLite pool.
93    #[cfg(feature = "sqlite")]
94    pub fn from_sqlite_pool(pool: SqlitePool) -> Self {
95        Self {
96            inner: DbInner::Sqlite(pool),
97        }
98    }
99
100    /// Begin a new transaction.
101    #[allow(unreachable_patterns)]
102    pub async fn begin(&self) -> Result<Tx, FlozError> {
103        match &self.inner {
104            #[cfg(feature = "postgres")]
105            DbInner::Postgres(pool) => {
106                let mut conn = pool.acquire().await?;
107                sqlx::query("BEGIN").execute(&mut *conn).await?;
108                Ok(Tx {
109                    inner: TxInner::Postgres(Mutex::new(conn)),
110                    committed: false,
111                })
112            }
113            #[cfg(feature = "sqlite")]
114            DbInner::Sqlite(pool) => {
115                let mut conn = pool.acquire().await?;
116                sqlx::query("BEGIN").execute(&mut *conn).await?;
117                Ok(Tx {
118                    inner: TxInner::Sqlite(Mutex::new(conn)),
119                    committed: false,
120                })
121            }
122            _ => unreachable!("no database backend enabled"),
123        }
124    }
125
126    /// Adjust SQL placeholder syntax for the current backend.
127    #[allow(unreachable_patterns)]
128    pub(crate) fn adjust_sql(&self, sql: &str) -> String {
129        match &self.inner {
130            #[cfg(feature = "postgres")]
131            DbInner::Postgres(_) => sql.to_string(),
132            #[cfg(feature = "sqlite")]
133            DbInner::Sqlite(_) => replace_pg_placeholders(sql),
134            _ => unreachable!("no database backend enabled"),
135        }
136    }
137}