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

alibi_sqlx/
pool.rs
