alibi_sqlx/
pool.rs

1//! Engine-tagged `SQLx` pools and transactions, and statement execution.
2
3use crate::error::map_sqlx_err;
4use crate::sql::Sql;
5use alibi_core::config::AuthConfig;
6use alibi_core::error::{AuthError, AuthResult};
7#[cfg(feature = "postgres")]
8use sqlx::postgres::{PgPool, PgRow, Postgres};
9#[cfg(feature = "sqlite")]
10use sqlx::sqlite::{Sqlite, SqlitePool, SqliteRow};
11use sqlx::{AssertSqlSafe, Decode, Type};
12
13#[cfg(not(any(feature = "sqlite", feature = "postgres")))]
14compile_error!("alibi-sqlx needs the `sqlite` or `postgres` feature");
15
16/// The database engine behind a pool.
17#[derive(Clone, Copy, Debug, PartialEq, Eq)]
18pub enum Engine {
19    Sqlite,
20    Postgres,
21}
22
23/// An application's `SQLx` pool. Better Auth shares it with the application.
24#[derive(Clone, Debug)]
25pub enum SqlxPool {
26    #[cfg(feature = "sqlite")]
27    Sqlite(SqlitePool),
28    #[cfg(feature = "postgres")]
29    Postgres(PgPool),
30}
31
32#[cfg(feature = "sqlite")]
33impl From<SqlitePool> for SqlxPool {
34    fn from(pool: SqlitePool) -> Self {
35        Self::Sqlite(pool)
36    }
37}
38
39#[cfg(feature = "postgres")]
40impl From<PgPool> for SqlxPool {
41    fn from(pool: PgPool) -> Self {
42        Self::Postgres(pool)
43    }
44}
45
46impl SqlxPool {
47    /// Connect a SQLite (`sqlite:`) or PostgreSQL (`postgres:`/`postgresql:`) URL.
48    ///
49    /// # Errors
50    ///
51    /// Returns the driver error, or a configuration error for a scheme whose
52    /// engine feature is not enabled.
53    pub async fn connect(url: &str) -> Result<Self, sqlx::Error> {
54        #[cfg(feature = "sqlite")]
55        if url.starts_with("sqlite:") {
56            return SqlitePool::connect(url).await.map(Self::Sqlite);
57        }
58        #[cfg(feature = "postgres")]
59        if url.starts_with("postgres:") || url.starts_with("postgresql:") {
60            return PgPool::connect(url).await.map(Self::Postgres);
61        }
62        Err(sqlx::Error::Configuration(
63            "the database URL scheme is not an enabled Better Auth SQLx engine".into(),
64        ))
65    }
66
67    #[must_use]
68    pub const fn engine(&self) -> Engine {
69        match self {
70            #[cfg(feature = "sqlite")]
71            Self::Sqlite(_) => Engine::Sqlite,
72            #[cfg(feature = "postgres")]
73            Self::Postgres(_) => Engine::Postgres,
74        }
75    }
76
77    #[cfg(feature = "sqlite")]
78    #[must_use]
79    pub const fn as_sqlite(&self) -> Option<&SqlitePool> {
80        match self {
81            Self::Sqlite(pool) => Some(pool),
82            #[cfg(feature = "postgres")]
83            Self::Postgres(_) => None,
84        }
85    }
86
87    #[cfg(feature = "postgres")]
88    #[must_use]
89    pub const fn as_postgres(&self) -> Option<&PgPool> {
90        match self {
91            Self::Postgres(pool) => Some(pool),
92            #[cfg(feature = "sqlite")]
93            Self::Sqlite(_) => None,
94        }
95    }
96
97    /// Close every pooled connection.
98    pub async fn close(&self) {
99        match self {
100            #[cfg(feature = "sqlite")]
101            Self::Sqlite(pool) => pool.close().await,
102            #[cfg(feature = "postgres")]
103            Self::Postgres(pool) => pool.close().await,
104        }
105    }
106
107    /// Run `statements` in order, without bound arguments. Generated
108    /// application migrations install their tables this way.
109    ///
110    /// # Errors
111    ///
112    /// Returns the first driver error.
113    pub async fn execute_batch(&self, statements: &[&str]) -> Result<(), sqlx::Error> {
114        for statement in statements {
115            let statement = AssertSqlSafe((*statement).to_owned());
116            match self {
117                #[cfg(feature = "sqlite")]
118                Self::Sqlite(pool) => _ = sqlx::raw_sql(statement).execute(pool).await?,
119                #[cfg(feature = "postgres")]
120                Self::Postgres(pool) => _ = sqlx::raw_sql(statement).execute(pool).await?,
121            }
122        }
123        Ok(())
124    }
125
126    /// Begin a transaction. SQLite uses `BEGIN IMMEDIATE` when `immediate` is set,
127    /// acquiring its writer reservation before the first read.
128    pub(crate) async fn begin(&self, immediate: bool) -> AuthResult<SqlxTransaction> {
129        let inner = match self {
130            #[cfg(feature = "sqlite")]
131            Self::Sqlite(pool) if immediate => pool
132                .begin_with("BEGIN IMMEDIATE")
133                .await
134                .map(TransactionKind::Sqlite),
135            #[cfg(feature = "sqlite")]
136            Self::Sqlite(pool) => pool.begin().await.map(TransactionKind::Sqlite),
137            #[cfg(feature = "postgres")]
138            Self::Postgres(pool) => {
139                _ = immediate;
140                pool.begin().await.map(TransactionKind::Postgres)
141            }
142        }
143        .map_err(map_sqlx_err)?;
144        Ok(SqlxTransaction {
145            engine: self.engine(),
146            config: None,
147            inner: tokio::sync::Mutex::new(Some(inner)),
148        })
149    }
150}
151
152/// An open driver transaction.
153#[derive(Debug)]
154pub(crate) enum TransactionKind {
155    #[cfg(feature = "sqlite")]
156    Sqlite(sqlx::Transaction<'static, Sqlite>),
157    #[cfg(feature = "postgres")]
158    Postgres(sqlx::Transaction<'static, Postgres>),
159}
160
161/// An auth transaction shared with lifecycle hooks.
162///
163/// Statements serialize on an internal lock; hooks run between statements, so a
164/// hook may [`lock`](Self::lock) the transaction to run its own queries in it.
165#[derive(Debug)]
166pub struct SqlxTransaction {
167    engine: Engine,
168    pub(crate) config: Option<std::sync::Arc<AuthConfig>>,
169    inner: tokio::sync::Mutex<Option<TransactionKind>>,
170}
171
172/// Exclusive access to an open auth transaction.
173#[derive(Debug)]
174pub struct SqlxTransactionGuard<'a>(tokio::sync::MutexGuard<'a, Option<TransactionKind>>);
175
176impl SqlxTransactionGuard<'_> {
177    /// The SQLite connection running this transaction.
178    #[cfg(feature = "sqlite")]
179    pub fn sqlite(&mut self) -> Option<&mut sqlx::SqliteConnection> {
180        match self.0.as_mut()? {
181            TransactionKind::Sqlite(transaction) => Some(&mut **transaction),
182            #[cfg(feature = "postgres")]
183            TransactionKind::Postgres(_) => None,
184        }
185    }
186
187    /// The PostgreSQL connection running this transaction.
188    #[cfg(feature = "postgres")]
189    pub fn postgres(&mut self) -> Option<&mut sqlx::PgConnection> {
190        match self.0.as_mut()? {
191            TransactionKind::Postgres(transaction) => Some(&mut **transaction),
192            #[cfg(feature = "sqlite")]
193            TransactionKind::Sqlite(_) => None,
194        }
195    }
196}
197
198impl SqlxTransaction {
199    #[must_use]
200    pub const fn engine(&self) -> Engine {
201        self.engine
202    }
203
204    /// Wait for the current statement, then run statements on this transaction.
205    pub async fn lock(&self) -> SqlxTransactionGuard<'_> {
206        SqlxTransactionGuard(self.inner.lock().await)
207    }
208
209    pub(crate) async fn commit(self) -> AuthResult<()> {
210        match self.inner.into_inner() {
211            #[cfg(feature = "sqlite")]
212            Some(TransactionKind::Sqlite(transaction)) => transaction.commit().await,
213            #[cfg(feature = "postgres")]
214            Some(TransactionKind::Postgres(transaction)) => transaction.commit().await,
215            None => return Err(AuthError::internal("Transaction already finished")),
216        }
217        .map_err(map_sqlx_err)
218    }
219
220    pub(crate) async fn rollback(self) -> AuthResult<()> {
221        match self.inner.into_inner() {
222            #[cfg(feature = "sqlite")]
223            Some(TransactionKind::Sqlite(transaction)) => transaction.rollback().await,
224            #[cfg(feature = "postgres")]
225            Some(TransactionKind::Postgres(transaction)) => transaction.rollback().await,
226            None => return Err(AuthError::internal("Transaction already finished")),
227        }
228        .map_err(map_sqlx_err)
229    }
230}
231
232macro_rules! decodable {
233    ($($row:tt)*) => {
234        /// A model decodable from the rows of every enabled engine.
235        pub trait SqlxRow: $($row)* + Send + Unpin {}
236
237        impl<T> SqlxRow for T where T: $($row)* + Send + Unpin {}
238    };
239}
240
241#[cfg(all(feature = "sqlite", feature = "postgres"))]
242decodable!(for<'r> sqlx::FromRow<'r, SqliteRow> + for<'r> sqlx::FromRow<'r, PgRow>);
243#[cfg(all(feature = "sqlite", not(feature = "postgres")))]
244decodable!(for<'r> sqlx::FromRow<'r, SqliteRow>);
245#[cfg(all(feature = "postgres", not(feature = "sqlite")))]
246decodable!(for<'r> sqlx::FromRow<'r, PgRow>);
247
248macro_rules! scalar {
249    ($($bound:tt)*) => {
250        /// A single-column value decodable from every enabled engine.
251        pub(crate) trait SqlxScalar: $($bound)* + Send + Unpin {}
252
253        impl<T> SqlxScalar for T where T: $($bound)* + Send + Unpin {}
254    };
255}
256
257#[cfg(all(feature = "sqlite", feature = "postgres"))]
258scalar!(for<'r> Decode<'r, Sqlite> + Type<Sqlite> + for<'r> Decode<'r, Postgres> + Type<Postgres>);
259#[cfg(all(feature = "sqlite", not(feature = "postgres")))]
260scalar!(for<'r> Decode<'r, Sqlite> + Type<Sqlite>);
261#[cfg(all(feature = "postgres", not(feature = "sqlite")))]
262scalar!(for<'r> Decode<'r, Postgres> + Type<Postgres>);
263
264/// Where a statement runs: the shared pool or an open transaction.
265#[derive(Clone, Copy, Debug)]
266enum ExecConnection<'a> {
267    Pool(&'a SqlxPool),
268    Tx(&'a SqlxTransaction),
269}
270
271#[derive(Clone, Copy, Debug)]
272pub(crate) struct Exec<'a> {
273    connection: ExecConnection<'a>,
274    config: Option<&'a AuthConfig>,
275}
276impl<'a> Exec<'a> {
277    pub(crate) const fn pool(pool: &'a SqlxPool) -> Self {
278        Self {
279            connection: ExecConnection::Pool(pool),
280            config: None,
281        }
282    }
283    pub(crate) fn tx(tx: &'a SqlxTransaction) -> Self {
284        Self {
285            connection: ExecConnection::Tx(tx),
286            config: tx.config.as_deref(),
287        }
288    }
289    pub(crate) const fn with_config(mut self, config: &'a AuthConfig) -> Self {
290        self.config = Some(config);
291        self
292    }
293    pub(crate) fn schema_name(self) -> Option<&'a str> {
294        self.config
295            .and_then(|config| config.advanced.database.schema_name.as_deref())
296    }
297    fn qualify(self, mut sql: String) -> AuthResult<String> {
298        if let Some(mapping) = self
299            .config
300            .and_then(|config| config.advanced.database.two_factor.as_ref())
301        {
302            sql = alibi_core::database_sql::map_two_factor(&sql, mapping)?;
303        }
304        if self.engine() == Engine::Postgres
305            && let Some(schema) = self.schema_name()
306        {
307            return alibi_core::database_sql::qualify_schema(&sql, schema);
308        }
309        Ok(sql)
310    }
311}
312
313/// Bind `$sql` for the engine behind `$exec` and run `$body` with `$query`
314/// (text and arguments) on `$executor`.
315macro_rules! dispatch {
316    ($exec:expr, $sql:expr, |$query:ident, $executor:ident| $body:expr) => {{
317        let exec = $exec;
318        let (text, args) = $sql.into_parts();
319        let text = exec.qualify(text)?;
320        match exec.connection {
321            #[cfg(feature = "sqlite")]
322            ExecConnection::Pool(SqlxPool::Sqlite(pool)) => {
323                match crate::sql::sqlite_arguments(args) {
324                    Ok(arguments) => {
325                        let $executor = pool;
326                        let $query = (AssertSqlSafe(text), arguments);
327                        $body
328                    }
329                    Err(error) => Err(sqlx::Error::Encode(error)),
330                }
331            }
332            #[cfg(feature = "postgres")]
333            ExecConnection::Pool(SqlxPool::Postgres(pool)) => {
334                match crate::sql::postgres_arguments(args) {
335                    Ok(arguments) => {
336                        let $executor = pool;
337                        let $query = (AssertSqlSafe(text), arguments);
338                        $body
339                    }
340                    Err(error) => Err(sqlx::Error::Encode(error)),
341                }
342            }
343            ExecConnection::Tx(transaction) => {
344                let mut guard = transaction.inner.lock().await;
345                match guard.as_mut() {
346                    #[cfg(feature = "sqlite")]
347                    Some(TransactionKind::Sqlite(open)) => match crate::sql::sqlite_arguments(args)
348                    {
349                        Ok(arguments) => {
350                            let $executor = &mut **open;
351                            let $query = (AssertSqlSafe(text), arguments);
352                            $body
353                        }
354                        Err(error) => Err(sqlx::Error::Encode(error)),
355                    },
356                    #[cfg(feature = "postgres")]
357                    Some(TransactionKind::Postgres(open)) => {
358                        match crate::sql::postgres_arguments(args) {
359                            Ok(arguments) => {
360                                let $executor = &mut **open;
361                                let $query = (AssertSqlSafe(text), arguments);
362                                $body
363                            }
364                            Err(error) => Err(sqlx::Error::Encode(error)),
365                        }
366                    }
367                    None => Err(sqlx::Error::PoolClosed),
368                }
369            }
370        }
371        .map_err(map_sqlx_err)
372    }};
373}
374
375impl Exec<'_> {
376    pub(crate) const fn engine(self) -> Engine {
377        match self.connection {
378            ExecConnection::Pool(pool) => pool.engine(),
379            ExecConnection::Tx(transaction) => transaction.engine,
380        }
381    }
382
383    pub(crate) async fn fetch_all<M: SqlxRow>(self, sql: Sql) -> AuthResult<Vec<M>> {
384        dispatch!(self, sql, |query, executor| {
385            sqlx::query_as_with(query.0, query.1)
386                .fetch_all(executor)
387                .await
388        })
389    }
390
391    pub(crate) async fn fetch_optional<M: SqlxRow>(self, sql: Sql) -> AuthResult<Option<M>> {
392        dispatch!(self, sql, |query, executor| {
393            sqlx::query_as_with(query.0, query.1)
394                .fetch_optional(executor)
395                .await
396        })
397    }
398
399    /// Run a statement, returning its affected row count.
400    pub(crate) async fn execute(self, sql: Sql) -> AuthResult<u64> {
401        dispatch!(self, sql, |query, executor| {
402            sqlx::query_with(query.0, query.1)
403                .execute(executor)
404                .await
405                .map(|result| result.rows_affected())
406        })
407    }
408
409    pub(crate) async fn fetch_all_scalar<T: SqlxScalar>(self, sql: Sql) -> AuthResult<Vec<T>> {
410        dispatch!(self, sql, |query, executor| {
411            sqlx::query_scalar_with(query.0, query.1)
412                .fetch_all(executor)
413                .await
414        })
415    }
416
417    pub(crate) async fn fetch_scalar<T: SqlxScalar>(self, sql: Sql) -> AuthResult<Option<T>> {
418        dispatch!(self, sql, |query, executor| {
419            sqlx::query_scalar_with(query.0, query.1)
420                .fetch_optional(executor)
421                .await
422        })
423    }
424
425    /// Run a fixed multi-statement script without bound arguments.
426    pub(crate) async fn execute_script(self, script: &'static str) -> AuthResult<()> {
427        let script = self.qualify(script.to_owned())?;
428        let result = match self.connection {
429            #[cfg(feature = "sqlite")]
430            ExecConnection::Pool(SqlxPool::Sqlite(pool)) => sqlx::raw_sql(AssertSqlSafe(script))
431                .execute(pool)
432                .await
433                .map(|_| ()),
434            #[cfg(feature = "postgres")]
435            ExecConnection::Pool(SqlxPool::Postgres(pool)) => sqlx::raw_sql(AssertSqlSafe(script))
436                .execute(pool)
437                .await
438                .map(|_| ()),
439            ExecConnection::Tx(transaction) => {
440                let mut guard = transaction.inner.lock().await;
441                match guard.as_mut() {
442                    #[cfg(feature = "sqlite")]
443                    Some(TransactionKind::Sqlite(open)) => sqlx::raw_sql(AssertSqlSafe(script))
444                        .execute(&mut **open)
445                        .await
446                        .map(|_| ()),
447                    #[cfg(feature = "postgres")]
448                    Some(TransactionKind::Postgres(open)) => sqlx::raw_sql(AssertSqlSafe(script))
449                        .execute(&mut **open)
450                        .await
451                        .map(|_| ()),
452                    None => Err(sqlx::Error::PoolClosed),
453                }
454            }
455        };
456        result.map_err(map_sqlx_err)
457    }
458}
alibi_sqlx/pool.rs

alibi_sqlx/
pool.rs
