Skip to main content

renox_core/db/
conn.rs

1//! Renox's own handle on the database, so the same code runs on SQLite and
2//! PostgreSQL: [`Db`] (the pool), [`Transaction`], [`Row`] and raw SQL
3//! through [`sql()`].
4//!
5//! SQL written for Renox uses `?` placeholders. On PostgreSQL they become
6//! `$1`, `$2`, … before the query is sent (a `?` inside quotes or comments is
7//! left alone).
8
9#[cfg(any(feature = "postgres", test))]
10use std::borrow::Cow;
11use std::fmt;
12use std::sync::Arc;
13
14use cookie::Key;
15
16use sqlx::sqlite::{Sqlite, SqliteArguments, SqlitePool, SqliteRow};
17use sqlx::{AssertSqlSafe, Column, Row as _};
18
19use super::value::DbValue;
20use super::{DbError, ToDbValue};
21
22#[cfg(feature = "postgres")]
23use sqlx::postgres::{PgArguments, PgPool, PgRow, Postgres};
24
25/// Which database engine a [`Db`] talks to.
26#[derive(Debug, Clone, Copy, PartialEq, Eq)]
27pub enum Dialect {
28    /// SQLite.
29    Sqlite,
30    /// PostgreSQL (the `postgres` feature).
31    Postgres,
32}
33
34impl fmt::Display for Dialect {
35    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
36        f.write_str(match self {
37            Dialect::Sqlite => "sqlite",
38            Dialect::Postgres => "postgres",
39        })
40    }
41}
42
43/// The database pool. Take it in a handler with `State(db): State<Db>`.
44///
45/// Models, queries and [`sql()`] take `&db` (or `&mut tx` inside a
46/// [`Transaction`]). For something only sqlx can do, reach the underlying
47/// pool with [`Db::sqlite`] or `Db::postgres` (with the `postgres` feature).
48#[derive(Clone)]
49pub struct Db {
50    pool: Pool,
51    schema: SchemaEpoch,
52    /// Seals and opens `Encrypted` columns (`APP_KEY`, set at boot).
53    key: Option<Arc<Key>>,
54}
55
56/// When this process last changed the schema (migrations ran).
57///
58/// A pooled connection keeps the schema it read, and prepares statements
59/// against it until it next steps one. sqlx takes a statement's columns from
60/// that first prepare, so after `ALTER TABLE users ADD COLUMN …` a
61/// `SELECT *` on an older connection gets one column more than sqlx expects
62/// and panics (sqlx-sqlite `row.rs`), or PostgreSQL refuses a cached plan.
63/// The pools made by [`super::connect`] drop connections opened before the
64/// last change instead of reusing them.
65#[derive(Clone, Default)]
66pub(crate) struct SchemaEpoch(std::sync::Arc<std::sync::Mutex<Option<std::time::Instant>>>);
67
68impl SchemaEpoch {
69    pub(crate) fn changed(&self) {
70        *self.0.lock().unwrap_or_else(|e| e.into_inner()) = Some(std::time::Instant::now());
71    }
72
73    /// Whether a connection of this age was opened before the last change.
74    pub(crate) fn is_stale(&self, age: std::time::Duration) -> bool {
75        self.0
76            .lock()
77            .unwrap_or_else(|e| e.into_inner())
78            .is_some_and(|at| age >= at.elapsed())
79    }
80}
81
82#[derive(Clone)]
83enum Pool {
84    Sqlite(SqlitePool),
85    #[cfg(feature = "postgres")]
86    Postgres(PgPool),
87}
88
89impl fmt::Debug for Db {
90    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
91        f.debug_struct("Db")
92            .field("dialect", &self.dialect())
93            .finish_non_exhaustive()
94    }
95}
96
97impl From<SqlitePool> for Db {
98    fn from(pool: SqlitePool) -> Self {
99        Self::with_epoch(Pool::Sqlite(pool), SchemaEpoch::default())
100    }
101}
102
103#[cfg(feature = "postgres")]
104impl From<PgPool> for Db {
105    fn from(pool: PgPool) -> Self {
106        Self::with_epoch(Pool::Postgres(pool), SchemaEpoch::default())
107    }
108}
109
110impl Db {
111    fn with_epoch(pool: Pool, schema: SchemaEpoch) -> Self {
112        Self {
113            pool,
114            schema,
115            key: None,
116        }
117    }
118
119    /// The key `Encrypted` columns are sealed and opened with (the app's
120    /// `APP_KEY`, set at boot).
121    pub(crate) fn with_key(mut self, key: Key) -> Self {
122        self.key = Some(Arc::new(key));
123        self
124    }
125
126    pub(crate) fn from_sqlite(pool: SqlitePool, schema: SchemaEpoch) -> Self {
127        Self::with_epoch(Pool::Sqlite(pool), schema)
128    }
129
130    #[cfg(feature = "postgres")]
131    pub(crate) fn from_postgres(pool: PgPool, schema: SchemaEpoch) -> Self {
132        Self::with_epoch(Pool::Postgres(pool), schema)
133    }
134
135    /// Records that the schema changed, so older pooled connections aren't
136    /// reused (see `SchemaEpoch`).
137    pub(crate) fn schema_changed(&self) {
138        self.schema.changed();
139    }
140
141    /// The engine this pool talks to.
142    pub fn dialect(&self) -> Dialect {
143        match self.pool {
144            Pool::Sqlite(_) => Dialect::Sqlite,
145            #[cfg(feature = "postgres")]
146            Pool::Postgres(_) => Dialect::Postgres,
147        }
148    }
149
150    /// The sqlx pool, when the database is SQLite.
151    pub fn sqlite(&self) -> Option<&SqlitePool> {
152        match &self.pool {
153            Pool::Sqlite(pool) => Some(pool),
154            #[cfg(feature = "postgres")]
155            Pool::Postgres(_) => None,
156        }
157    }
158
159    /// The sqlx pool, when the database is PostgreSQL.
160    #[cfg(feature = "postgres")]
161    pub fn postgres(&self) -> Option<&PgPool> {
162        match &self.pool {
163            Pool::Postgres(pool) => Some(pool),
164            Pool::Sqlite(_) => None,
165        }
166    }
167
168    /// Starts a transaction. Pass `&mut tx` wherever a `&db` goes, then
169    /// `tx.commit()`; dropping it without committing rolls it back.
170    pub async fn begin(&self) -> Result<Transaction, DbError> {
171        let inner = match &self.pool {
172            Pool::Sqlite(pool) => TxInner::Sqlite(pool.begin().await?),
173            #[cfg(feature = "postgres")]
174            Pool::Postgres(pool) => TxInner::Postgres(pool.begin().await?),
175        };
176        Ok(Transaction {
177            inner,
178            key: self.key.clone(),
179            savepoints: 0,
180        })
181    }
182
183    /// Runs `work` in a transaction: committed when it returns `Ok`, rolled
184    /// back when it returns `Err`.
185    ///
186    /// ```
187    /// # use renox::prelude::*;
188    /// # async fn demo(db: Db) -> Result {
189    /// let moved = db
190    ///     .transaction_retrying(3, |tx| {
191    ///         Box::pin(async move {
192    ///             renox::db::sql("UPDATE accounts SET balance = balance - 100 WHERE id = 1")
193    ///                 .execute(&mut *tx)
194    ///                 .await?;
195    ///             renox::db::sql("UPDATE accounts SET balance = balance + 100 WHERE id = 2")
196    ///                 .execute(&mut *tx)
197    ///                 .await?;
198    ///             Ok(100)
199    ///         })
200    ///     })
201    ///     .await?;
202    /// # let _: i64 = moved; Ok(()) }
203    /// ```
204    pub async fn transaction<T, F>(&self, work: F) -> crate::Result<T>
205    where
206        F: for<'t> FnMut(
207            &'t mut Transaction,
208        ) -> futures_util::future::BoxFuture<'t, crate::Result<T>>,
209    {
210        self.transaction_retrying(1, work).await
211    }
212
213    /// Like [`Db::transaction`], trying up to `attempts` times when another
214    /// transaction got in the way (SQLite busy, PostgreSQL serialization
215    /// failure or deadlock), with a short, growing pause in between. `work`
216    /// must be safe to run again.
217    pub async fn transaction_retrying<T, F>(&self, attempts: u32, mut work: F) -> crate::Result<T>
218    where
219        F: for<'t> FnMut(
220            &'t mut Transaction,
221        ) -> futures_util::future::BoxFuture<'t, crate::Result<T>>,
222    {
223        let attempts = attempts.max(1);
224        let mut attempt = 1;
225        loop {
226            let outcome = async {
227                let mut tx = self.begin().await?;
228                let value = work(&mut tx).await?;
229                tx.commit().await?;
230                Ok::<T, crate::Error>(value)
231            }
232            .await;
233            match outcome {
234                Err(err) if attempt < attempts && err.is_retryable() => {
235                    let pause = std::time::Duration::from_millis(20 * u64::from(attempt));
236                    tokio::time::sleep(pause).await;
237                    attempt += 1;
238                }
239                other => return other,
240            }
241        }
242    }
243
244    /// Runs `work` again (up to `attempts` times, with a short, growing
245    /// pause) while it fails because another transaction got in the way
246    /// (SQLite busy, PostgreSQL serialization failure or deadlock). `work`
247    /// opens and commits its own transaction, so unlike
248    /// [`Db::transaction_retrying`] it can borrow from the caller, and it can
249    /// roll back and still return a value (drop or `rollback` the
250    /// transaction instead of committing).
251    ///
252    /// ```
253    /// # use renox::prelude::*;
254    /// enum Transfer { Done, Short }
255    ///
256    /// # async fn demo(db: Db, amount: i64, from: i64, to: i64) -> Result {
257    /// let outcome = db
258    ///     .retrying(3, || async {
259    ///         let mut tx = db.begin().await?;
260    ///         let taken = renox::db::sql(
261    ///             "UPDATE accounts SET balance = balance - ? WHERE id = ? AND balance >= ?",
262    ///         )
263    ///         .bind(amount) // borrowed from the caller
264    ///         .bind(from)
265    ///         .bind(amount)
266    ///         .execute(&mut tx)
267    ///         .await?;
268    ///         if taken == 0 {
269    ///             tx.rollback().await?;
270    ///             return Ok(Transfer::Short); // rolled back, with an answer
271    ///         }
272    ///         renox::db::sql("UPDATE accounts SET balance = balance + ? WHERE id = ?")
273    ///             .bind(amount)
274    ///             .bind(to)
275    ///             .execute(&mut tx)
276    ///             .await?;
277    ///         tx.commit().await?;
278    ///         Ok(Transfer::Done)
279    ///     })
280    ///     .await?;
281    /// # let _ = outcome; Ok(()) }
282    /// ```
283    pub fn retrying<'a, T, F, Fut>(
284        &'a self,
285        attempts: u32,
286        mut work: F,
287    ) -> impl Future<Output = crate::Result<T>> + Send + 'a
288    where
289        T: Send + 'a,
290        F: FnMut() -> Fut + Send + 'a,
291        Fut: Future<Output = crate::Result<T>> + Send + 'a,
292    {
293        let attempts = attempts.max(1);
294        async move {
295            let mut attempt = 1;
296            loop {
297                match work().await {
298                    Err(err) if attempt < attempts && err.is_retryable() => {
299                        let pause = std::time::Duration::from_millis(20 * u64::from(attempt));
300                        tokio::time::sleep(pause).await;
301                        attempt += 1;
302                    }
303                    other => return other,
304                }
305            }
306        }
307    }
308
309    /// A transaction that takes SQLite's write lock at once (`BEGIN
310    /// IMMEDIATE`), so what it reads can't change before it writes: the
311    /// SQLite counterpart of `lock_for_update`. Same as `begin` on PostgreSQL.
312    pub async fn begin_immediate(&self) -> Result<Transaction, DbError> {
313        let inner = match &self.pool {
314            Pool::Sqlite(pool) => TxInner::Sqlite(pool.begin_with("BEGIN IMMEDIATE").await?),
315            #[cfg(feature = "postgres")]
316            Pool::Postgres(pool) => TxInner::Postgres(pool.begin().await?),
317        };
318        Ok(Transaction {
319            inner,
320            key: self.key.clone(),
321            savepoints: 0,
322        })
323    }
324
325    /// Closes every connection; later queries fail.
326    pub async fn close(&self) {
327        match &self.pool {
328            Pool::Sqlite(pool) => pool.close().await,
329            #[cfg(feature = "postgres")]
330            Pool::Postgres(pool) => pool.close().await,
331        }
332    }
333}
334
335/// A database transaction, from [`Db::begin`].
336pub struct Transaction {
337    inner: TxInner,
338    key: Option<Arc<Key>>,
339    /// Savepoints open inside it (for their names).
340    savepoints: u32,
341}
342
343enum TxInner {
344    Sqlite(sqlx::Transaction<'static, Sqlite>),
345    #[cfg(feature = "postgres")]
346    Postgres(sqlx::Transaction<'static, Postgres>),
347}
348
349impl Transaction {
350    /// The engine this transaction runs on.
351    pub fn dialect(&self) -> Dialect {
352        match self.inner {
353            TxInner::Sqlite(_) => Dialect::Sqlite,
354            #[cfg(feature = "postgres")]
355            TxInner::Postgres(_) => Dialect::Postgres,
356        }
357    }
358
359    /// Commits the transaction.
360    pub async fn commit(self) -> Result<(), DbError> {
361        match self.inner {
362            TxInner::Sqlite(tx) => Ok(tx.commit().await?),
363            #[cfg(feature = "postgres")]
364            TxInner::Postgres(tx) => Ok(tx.commit().await?),
365        }
366    }
367
368    /// Rolls the transaction back (dropping it uncommitted does the same).
369    pub async fn rollback(self) -> Result<(), DbError> {
370        match self.inner {
371            TxInner::Sqlite(tx) => Ok(tx.rollback().await?),
372            #[cfg(feature = "postgres")]
373            TxInner::Postgres(tx) => Ok(tx.rollback().await?),
374        }
375    }
376
377    /// Runs `work` inside a savepoint (a transaction inside the
378    /// transaction): when it returns `Err`, only what it did is undone and
379    /// the transaction goes on; when it returns `Ok`, its changes stay, to
380    /// be committed (or rolled back) with the rest. Savepoints nest.
381    ///
382    /// ```
383    /// # use renox::prelude::*;
384    /// # async fn demo(db: Db) -> Result {
385    /// let mut tx = db.begin().await?;
386    /// renox::db::sql("INSERT INTO orders (total) VALUES (?)").bind(18_000).execute(&mut tx).await?;
387    /// // Optional: a voucher that may be used up already. If it is, the order stays.
388    /// let voucher = tx
389    ///     .savepoint(|tx| {
390    ///         Box::pin(async move {
391    ///             let used = renox::db::sql("UPDATE vouchers SET used = used + 1 WHERE code = ? AND used < max_uses")
392    ///                 .bind("KOPI10")
393    ///                 .execute(&mut *tx)
394    ///                 .await?;
395    ///             renox::abort_if(used == 0, StatusCode::CONFLICT, "used up")?;
396    ///             Ok(())
397    ///         })
398    ///     })
399    ///     .await;
400    /// tx.commit().await?;
401    /// # let _ = voucher; Ok(()) }
402    /// ```
403    ///
404    /// On PostgreSQL a failed statement stops the whole transaction until
405    /// it rolls back; inside a savepoint only the savepoint rolls back, so
406    /// the transaction can go on after an expected failure (a unique
407    /// violation you handle, say).
408    pub async fn savepoint<T, F>(&mut self, work: F) -> crate::Result<T>
409    where
410        F: for<'t> FnOnce(
411            &'t mut Transaction,
412        ) -> futures_util::future::BoxFuture<'t, crate::Result<T>>,
413    {
414        let name = format!("renox_savepoint_{}", self.savepoints + 1);
415        sql(format!("SAVEPOINT {name}")).execute(&mut *self).await?;
416        self.savepoints += 1;
417        let outcome = work(self).await;
418        self.savepoints -= 1;
419        match outcome {
420            Ok(value) => {
421                sql(format!("RELEASE SAVEPOINT {name}"))
422                    .execute(&mut *self)
423                    .await?;
424                Ok(value)
425            }
426            Err(err) => {
427                sql(format!("ROLLBACK TO SAVEPOINT {name}"))
428                    .execute(&mut *self)
429                    .await?;
430                sql(format!("RELEASE SAVEPOINT {name}"))
431                    .execute(&mut *self)
432                    .await?;
433                Err(err)
434            }
435        }
436    }
437}
438
439impl fmt::Debug for Transaction {
440    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
441        f.debug_struct("Transaction")
442            .field("dialect", &self.dialect())
443            .finish_non_exhaustive()
444    }
445}
446
447/// Where a query runs: `&db` or `&mut tx`. Sealed: only those.
448///
449/// An executor is used up by one statement. A helper that runs several
450/// statements on whichever it is given turns it into a [`Conn`] first and
451/// [reborrows](Conn::reborrow) it for each one:
452///
453/// ```
454/// # use renox::prelude::*;
455/// use renox::db::{Executor, sql};
456///
457/// /// Records a movement and updates the level, on `&db` or in `&mut tx`.
458/// async fn record<'c>(db: impl Executor<'c>, product: i64, quantity: i64) -> Result {
459///     let mut conn = db.into_conn();
460///     sql("INSERT INTO movements (product_id, quantity) VALUES (?, ?)")
461///         .bind(product)
462///         .bind(quantity)
463///         .execute(conn.reborrow())
464///         .await?;
465///     sql("UPDATE levels SET quantity = quantity + ? WHERE product_id = ?")
466///         .bind(quantity)
467///         .bind(product)
468///         .execute(conn.reborrow())
469///         .await?;
470///     Ok(())
471/// }
472/// # async fn demo(db: &renox::db::Db) -> Result {
473/// record(db, 1, 5).await?;
474/// let mut tx = db.begin().await?;
475/// record(&mut tx, 1, -2).await?;
476/// tx.commit().await?;
477/// # Ok(()) }
478/// ```
479pub trait Executor<'c>: Send + executor::Sealed {
480    /// This executor as a [`Conn`], which can run several statements.
481    fn into_conn(self) -> Conn<'c>;
482}
483
484mod executor {
485    pub trait Sealed {}
486    impl Sealed for &super::Db {}
487    impl Sealed for &mut super::Transaction {}
488    impl Sealed for super::Conn<'_> {}
489}
490
491/// A connection for several statements: the pool (`&db`) or a transaction
492/// (`&mut tx`), from [`Executor::into_conn`]. It is an `Executor` itself;
493/// give each statement [`Conn::reborrow`] to keep it for the next one.
494pub enum Conn<'c> {
495    /// Statements run on the pool, each on any free connection.
496    #[doc(hidden)]
497    Pool(&'c Db),
498    /// Statements run in the transaction.
499    #[doc(hidden)]
500    Tx(&'c mut Transaction),
501}
502
503impl<'c> Executor<'c> for &'c Db {
504    fn into_conn(self) -> Conn<'c> {
505        Conn::Pool(self)
506    }
507}
508
509impl<'c> Executor<'c> for &'c mut Transaction {
510    fn into_conn(self) -> Conn<'c> {
511        Conn::Tx(self)
512    }
513}
514
515impl<'c> Executor<'c> for Conn<'c> {
516    fn into_conn(self) -> Conn<'c> {
517        self
518    }
519}
520
521impl Conn<'_> {
522    /// The key for `Encrypted` columns of the database this runs on.
523    fn key(&self) -> Option<Arc<Key>> {
524        match self {
525            Conn::Pool(db) => db.key.clone(),
526            Conn::Tx(tx) => tx.key.clone(),
527        }
528    }
529
530    /// The same connection for one more statement (a model's `save`, a
531    /// query's `fetch_all`…), keeping this one for the next.
532    pub fn reborrow(&mut self) -> Conn<'_> {
533        match self {
534            Conn::Pool(db) => Conn::Pool(db),
535            Conn::Tx(tx) => Conn::Tx(tx),
536        }
537    }
538
539    /// Which engine the statements run on, for SQL that differs.
540    pub fn dialect(&self) -> Dialect {
541        match self {
542            Conn::Pool(db) => db.dialect(),
543            Conn::Tx(tx) => tx.dialect(),
544        }
545    }
546}
547
548/// Runs `$body` against whichever backend `$conn` is: `$build` makes the
549/// backend's query from `(sql, args)` and `$exec` is its sqlx executor.
550macro_rules! dispatch {
551    ($conn:expr, |$build:ident, $exec:ident| $body:expr) => {
552        match $conn {
553            Conn::Pool(db) => match &db.pool {
554                Pool::Sqlite(pool) => {
555                    #[allow(unused_variables)]
556                    let $build = sqlite_query;
557                    let $exec = pool;
558                    ($body).map_err(DbError::from)
559                }
560                #[cfg(feature = "postgres")]
561                Pool::Postgres(pool) => {
562                    #[allow(unused_variables)]
563                    let $build = postgres_query;
564                    let $exec = pool;
565                    ($body).map_err(DbError::from)
566                }
567            },
568            Conn::Tx(tx) => match &mut tx.inner {
569                TxInner::Sqlite(tx) => {
570                    #[allow(unused_variables)]
571                    let $build = sqlite_query;
572                    let $exec = &mut **tx;
573                    ($body).map_err(DbError::from)
574                }
575                #[cfg(feature = "postgres")]
576                TxInner::Postgres(tx) => {
577                    #[allow(unused_variables)]
578                    let $build = postgres_query;
579                    let $exec = &mut **tx;
580                    ($body).map_err(DbError::from)
581                }
582            },
583        }
584    };
585}
586
587/// Seals the values of `Encrypted` fields with the database's key.
588fn seal(key: Option<&Key>, args: Vec<DbValue>) -> Result<Vec<DbValue>, DbError> {
589    args.into_iter()
590        .map(|value| match value {
591            DbValue::Encrypted(plain) => match key {
592                Some(key) => Ok(DbValue::Text(super::encrypted::seal(key, &plain.0))),
593                None => Err(DbError::from(sqlx::Error::Encode(
594                    "an Encrypted value needs the app's Db (it has the APP_KEY); \
595                     this Db was made outside App"
596                        .into(),
597                ))),
598            },
599            other => Ok(other),
600        })
601        .collect()
602}
603
604fn sqlite_query(
605    sql: String,
606    args: Vec<DbValue>,
607) -> sqlx::query::Query<'static, Sqlite, SqliteArguments> {
608    args.into_iter().fold(
609        sqlx::query(AssertSqlSafe(sql)),
610        |query, value| match value.for_sqlite() {
611            DbValue::Integer(v) => query.bind(v),
612            DbValue::Real(v) => query.bind(v),
613            DbValue::Text(v) => query.bind(v),
614            DbValue::Blob(v) => query.bind(v),
615            _ => query.bind(None::<i64>),
616        },
617    )
618}
619
620#[cfg(feature = "postgres")]
621fn postgres_query(
622    sql: String,
623    args: Vec<DbValue>,
624) -> sqlx::query::Query<'static, Postgres, PgArguments> {
625    let sql = numbered_placeholders(&sql).into_owned();
626    args.into_iter().fold(
627        sqlx::query(AssertSqlSafe(sql)),
628        |query, value| match value {
629            DbValue::Null => query.bind(UntypedNull),
630            DbValue::Integer(v) => query.bind(v),
631            DbValue::Real(v) => query.bind(v),
632            DbValue::Text(v) => query.bind(v),
633            DbValue::Blob(v) => query.bind(v),
634            DbValue::Bool(v) => query.bind(v),
635            DbValue::DateTime(v) => query.bind(v),
636            DbValue::NaiveDateTime(v) => query.bind(v),
637            DbValue::Date(v) => query.bind(v),
638            DbValue::Time(v) => query.bind(v),
639            DbValue::Json(v) => query.bind(sqlx::types::Json(v)),
640            #[cfg(feature = "uuid")]
641            DbValue::Uuid(v) => query.bind(v),
642            // Sealed into `Text` before the statement is built.
643            DbValue::Encrypted(_) => query.bind(UntypedNull),
644        },
645    )
646}
647
648/// A `NULL` parameter without a type, so PostgreSQL takes the type from
649/// where it's used (a `NULL` sent as `BIGINT` can't go into a `TIMESTAMPTZ`).
650#[cfg(feature = "postgres")]
651struct UntypedNull;
652
653#[cfg(feature = "postgres")]
654impl sqlx::Type<Postgres> for UntypedNull {
655    fn type_info() -> sqlx::postgres::PgTypeInfo {
656        // OID 0 means "unspecified" in the protocol.
657        sqlx::postgres::PgTypeInfo::with_oid(sqlx::postgres::types::Oid(0))
658    }
659}
660
661#[cfg(feature = "postgres")]
662impl sqlx::Encode<'_, Postgres> for UntypedNull {
663    fn encode_by_ref(
664        &self,
665        _buf: &mut sqlx::postgres::PgArgumentBuffer,
666    ) -> Result<sqlx::encode::IsNull, sqlx::error::BoxDynError> {
667        Ok(sqlx::encode::IsNull::Yes)
668    }
669}
670
671/// Rewrites `?` placeholders as PostgreSQL's `$1`, `$2`, …, skipping quoted
672/// strings (`'…'`, `E'…'` with backslash escapes, dollar-quoted `$$…$$` and
673/// `$tag$…$tag$`), quoted identifiers and comments.
674#[cfg(any(feature = "postgres", test))]
675pub(crate) fn numbered_placeholders(sql: &str) -> Cow<'_, str> {
676    if !sql.contains('?') {
677        return Cow::Borrowed(sql);
678    }
679    // Every delimiter is ASCII, so byte offsets always fall between characters.
680    let bytes = sql.as_bytes();
681    let ident = |b: u8| b.is_ascii_alphanumeric() || b == b'_' || b == b'$';
682    let mut out = String::with_capacity(sql.len() + 8);
683    let mut n = 0;
684    let mut copied = 0; // `sql[..copied]` is in `out`
685    let mut i = 0;
686    while i < bytes.len() {
687        let at = i;
688        let end = match bytes[i] {
689            quote @ (b'\'' | b'"') => {
690                // `E'…'` (or `e'…'`): a backslash escapes the next character.
691                let escapes = quote == b'\''
692                    && i > 0
693                    && matches!(bytes[i - 1], b'E' | b'e')
694                    && (i < 2 || !ident(bytes[i - 2]));
695                let mut j = i + 1;
696                // A doubled quote inside is an escaped quote: the loop ends at
697                // the first one and the next one opens a new run.
698                while j < bytes.len() && bytes[j] != quote {
699                    j += if escapes && bytes[j] == b'\\' { 2 } else { 1 };
700                }
701                j + 1
702            }
703            b'-' if bytes.get(i + 1) == Some(&b'-') => {
704                sql[i..].find('\n').map_or(bytes.len(), |nl| i + nl + 1)
705            }
706            b'/' if bytes.get(i + 1) == Some(&b'*') => sql[i + 2..]
707                .find("*/")
708                .map_or(bytes.len(), |e| i + 2 + e + 2),
709            b'$' if i == 0 || !ident(bytes[i - 1]) => {
710                // `$tag$` opens a dollar-quoted string (the tag may be empty
711                // and doesn't start with a digit, unlike `$1`), closed by the
712                // same `$tag$`.
713                let tag_len = bytes[i + 1..]
714                    .iter()
715                    .take_while(|b| b.is_ascii_alphanumeric() || **b == b'_')
716                    .count();
717                let starts_with_digit = bytes.get(i + 1).is_some_and(u8::is_ascii_digit);
718                if !starts_with_digit && bytes.get(i + 1 + tag_len) == Some(&b'$') {
719                    let tag = &sql[i..i + tag_len + 2];
720                    sql[i + tag.len()..]
721                        .find(tag)
722                        .map_or(bytes.len(), |e| i + tag.len() + e + tag.len())
723                } else {
724                    i + 1
725                }
726            }
727            b'?' => {
728                n += 1;
729                out.push_str(&sql[copied..i]);
730                out.push_str(&format!("${n}"));
731                copied = i + 1;
732                i + 1
733            }
734            _ => i + 1,
735        };
736        i = end.min(bytes.len()).max(at + 1);
737    }
738    out.push_str(&sql[copied..]);
739    Cow::Owned(out)
740}
741
742/// A raw SQL statement with `?` placeholders.
743///
744/// ```
745/// # use renox::prelude::*;
746/// # async fn demo(db: Db) -> Result {
747/// let rows = renox::db::sql("SELECT name FROM products WHERE price < ?")
748///     .bind(20_000)
749///     .fetch_all(&db)
750///     .await?;
751/// let name: String = rows[0].try_get("name")?;
752///
753/// let total: i64 = renox::db::sql("SELECT COUNT(*) FROM products").scalar(&db).await?;
754/// # let _ = (name, total); Ok(()) }
755/// ```
756pub fn sql(sql: impl Into<String>) -> Sql {
757    Sql {
758        sql: sql.into(),
759        args: Vec::new(),
760    }
761}
762
763/// A statement built by [`sql()`].
764#[derive(Debug, Clone)]
765#[must_use = "a statement does nothing until it is run"]
766pub struct Sql {
767    sql: String,
768    args: Vec<DbValue>,
769}
770
771impl Sql {
772    /// Binds the next `?`.
773    pub fn bind(mut self, value: impl ToDbValue) -> Self {
774        self.args.push(value.to_db_value());
775        self
776    }
777
778    /// Binds several values in order.
779    pub fn bind_all(mut self, values: impl IntoIterator<Item = DbValue>) -> Self {
780        self.args.extend(values);
781        self
782    }
783
784    /// Runs the statement and returns every row.
785    pub async fn fetch_all<'c>(self, db: impl Executor<'c>) -> Result<Vec<Row>, DbError> {
786        let Self { sql, args } = self;
787        super::query_log::record(&sql);
788        let conn = db.into_conn();
789        let key = conn.key();
790        let args = seal(key.as_deref(), args)?;
791        let rows: Vec<Row> = dispatch!(conn, |build, exec| build(sql, args)
792            .fetch_all(exec)
793            .await
794            .map(|rows| rows.into_iter().map(Row::from).collect()))?;
795        Ok(rows
796            .into_iter()
797            .map(|row| row.with_key(key.clone()))
798            .collect())
799    }
800
801    /// Runs the statement and returns the first row, if any.
802    pub async fn fetch_optional<'c>(self, db: impl Executor<'c>) -> Result<Option<Row>, DbError> {
803        let Self { sql, args } = self;
804        super::query_log::record(&sql);
805        let conn = db.into_conn();
806        let key = conn.key();
807        let args = seal(key.as_deref(), args)?;
808        let row: Option<Row> = dispatch!(conn, |build, exec| build(sql, args)
809            .fetch_optional(exec)
810            .await
811            .map(|row| row.map(Row::from)))?;
812        Ok(row.map(|row| row.with_key(key)))
813    }
814
815    /// The first row; an error if there is none.
816    pub async fn fetch_one<'c>(self, db: impl Executor<'c>) -> Result<Row, DbError> {
817        self.fetch_optional(db)
818            .await?
819            .ok_or_else(|| DbError::from(sqlx::Error::RowNotFound))
820    }
821
822    /// Runs the statement and returns the number of rows it changed.
823    pub async fn execute<'c>(self, db: impl Executor<'c>) -> Result<u64, DbError> {
824        let Self { sql, args } = self;
825        super::query_log::record(&sql);
826        let conn = db.into_conn();
827        let args = seal(conn.key().as_deref(), args)?;
828        dispatch!(conn, |build, exec| build(sql, args)
829            .execute(exec)
830            .await
831            .map(|done| done.rows_affected()))
832    }
833
834    /// Every row, read into `T`: a model, a `#[derive(FromRow)]` struct or a
835    /// tuple; see [`super::FromRow`].
836    pub async fn fetch_as<'c, T: super::FromRow>(
837        self,
838        db: impl Executor<'c>,
839    ) -> Result<Vec<T>, DbError> {
840        self.fetch_all(db).await?.iter().map(T::from_row).collect()
841    }
842
843    /// The first row read into `T`, or `None` if there is no row.
844    pub async fn fetch_optional_as<'c, T: super::FromRow>(
845        self,
846        db: impl Executor<'c>,
847    ) -> Result<Option<T>, DbError> {
848        self.fetch_optional(db)
849            .await?
850            .as_ref()
851            .map(T::from_row)
852            .transpose()
853    }
854
855    /// The first row read into `T`; an error if there is none.
856    pub async fn fetch_one_as<'c, T: super::FromRow>(
857        self,
858        db: impl Executor<'c>,
859    ) -> Result<T, DbError> {
860        T::from_row(&self.fetch_one(db).await?)
861    }
862
863    /// The first column of the first row; an error if there is no row.
864    pub async fn scalar<'c, T: FromDb>(self, db: impl Executor<'c>) -> Result<T, DbError> {
865        self.fetch_one(db).await?.try_get(0)
866    }
867
868    /// The first column of the first row, or `None` if there is no row.
869    pub async fn scalar_optional<'c, T: FromDb>(
870        self,
871        db: impl Executor<'c>,
872    ) -> Result<Option<T>, DbError> {
873        self.fetch_optional(db)
874            .await?
875            .map(|row| row.try_get(0))
876            .transpose()
877    }
878
879    /// The first column of every row.
880    pub async fn scalars<'c, T: FromDb>(self, db: impl Executor<'c>) -> Result<Vec<T>, DbError> {
881        self.fetch_all(db)
882            .await?
883            .iter()
884            .map(|row| row.try_get(0))
885            .collect()
886    }
887}
888
889/// Runs SQL that may hold several statements and no parameters, e.g. a
890/// migration file. Returns the number of rows changed.
891pub(crate) async fn script<'c>(db: impl Executor<'c>, sql: &str) -> Result<u64, DbError> {
892    let sql = sql.to_owned();
893    dispatch!(db.into_conn(), |build, exec| sqlx::raw_sql(AssertSqlSafe(
894        sql
895    ))
896    .execute(exec)
897    .await
898    .map(|done| done.rows_affected()))
899}
900
901/// One result row. Read columns by name or position with [`Row::try_get`].
902pub struct Row(pub(crate) RowInner, Option<Arc<Key>>);
903
904pub(crate) enum RowInner {
905    Sqlite(SqliteRow),
906    #[cfg(feature = "postgres")]
907    Postgres(PgRow),
908}
909
910impl From<SqliteRow> for Row {
911    fn from(row: SqliteRow) -> Self {
912        Self(RowInner::Sqlite(row), None)
913    }
914}
915
916#[cfg(feature = "postgres")]
917impl From<PgRow> for Row {
918    fn from(row: PgRow) -> Self {
919        Self(RowInner::Postgres(row), None)
920    }
921}
922
923impl Row {
924    fn with_key(mut self, key: Option<Arc<Key>>) -> Self {
925        self.1 = key;
926        self
927    }
928
929    /// A column's value, by name (`"name"`) or position (`0`).
930    pub fn try_get<T: FromDb>(&self, index: impl RowIndex) -> Result<T, DbError> {
931        super::encrypted::reading(self.1.as_ref(), || match &self.0 {
932            RowInner::Sqlite(row) => Ok(row.try_get(index)?),
933            #[cfg(feature = "postgres")]
934            RowInner::Postgres(row) => Ok(row.try_get(index)?),
935        })
936    }
937
938    /// A column's value as JSON, whatever its type; `null` when it can't be
939    /// read as a number, boolean, text, date, time, timestamp or JSON.
940    pub(crate) fn json(&self, column: &str) -> serde_json::Value {
941        use serde_json::Value;
942        if let Ok(v) = self.try_get::<Option<i64>>(column) {
943            return v.map_or(Value::Null, Value::from);
944        }
945        if let Ok(v) = self.try_get::<Option<i32>>(column) {
946            return v.map_or(Value::Null, Value::from);
947        }
948        // PostgreSQL's SMALLINT and REAL decode only as these.
949        if let Ok(v) = self.try_get::<Option<i16>>(column) {
950            return v.map_or(Value::Null, Value::from);
951        }
952        if let Ok(v) = self.try_get::<Option<f64>>(column) {
953            return v.map_or(Value::Null, Value::from);
954        }
955        if let Ok(v) = self.try_get::<Option<f32>>(column) {
956            return v.map_or(Value::Null, |n| Value::from(f64::from(n)));
957        }
958        if let Ok(v) = self.try_get::<Option<bool>>(column) {
959            return v.map_or(Value::Null, Value::from);
960        }
961        if let Ok(v) = self.try_get::<Option<chrono::DateTime<chrono::Utc>>>(column) {
962            return v.map_or(Value::Null, |d| Value::from(d.to_rfc3339()));
963        }
964        if let Ok(v) = self.try_get::<Option<chrono::NaiveDateTime>>(column) {
965            return v.map_or(Value::Null, |d| {
966                Value::from(d.format("%Y-%m-%dT%H:%M:%S").to_string())
967            });
968        }
969        if let Ok(v) = self.try_get::<Option<chrono::NaiveDate>>(column) {
970            return v.map_or(Value::Null, |d| Value::from(d.to_string()));
971        }
972        if let Ok(v) = self.try_get::<Option<String>>(column) {
973            return v.map_or(Value::Null, Value::from);
974        }
975        // After text, so SQLite text that looks like a time stays as written;
976        // PostgreSQL's TIME decodes only as this.
977        if let Ok(v) = self.try_get::<Option<chrono::NaiveTime>>(column) {
978            return v.map_or(Value::Null, |t| Value::from(t.to_string()));
979        }
980        if let Ok(v) = self.try_get::<Option<serde_json::Value>>(column) {
981            return v.unwrap_or(Value::Null);
982        }
983        Value::Null
984    }
985
986    /// The column names, in order.
987    pub fn columns(&self) -> Vec<&str> {
988        match &self.0 {
989            RowInner::Sqlite(row) => row.columns().iter().map(Column::name).collect(),
990            #[cfg(feature = "postgres")]
991            RowInner::Postgres(row) => row.columns().iter().map(Column::name).collect(),
992        }
993    }
994
995    /// The sqlx row, when the database is SQLite.
996    pub fn sqlite(&self) -> Option<&SqliteRow> {
997        match &self.0 {
998            RowInner::Sqlite(row) => Some(row),
999            #[cfg(feature = "postgres")]
1000            RowInner::Postgres(_) => None,
1001        }
1002    }
1003
1004    /// The sqlx row, when the database is PostgreSQL.
1005    #[cfg(feature = "postgres")]
1006    pub fn postgres(&self) -> Option<&PgRow> {
1007        match &self.0 {
1008            RowInner::Postgres(row) => Some(row),
1009            RowInner::Sqlite(_) => None,
1010        }
1011    }
1012}
1013
1014impl fmt::Debug for Row {
1015    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1016        f.debug_struct("Row")
1017            .field("columns", &self.columns())
1018            .finish_non_exhaustive()
1019    }
1020}
1021
1022/// A Rust type that can be read from a column on every enabled database:
1023/// the integer and float types, `bool`, `String`, `Vec<u8>`, the chrono date
1024/// and time types, `Option<T>` of those, `Json<T>`, `Encrypted<T>`, `Ulid`,
1025/// `Uuid` (the `uuid` feature) and `#[derive(DbEnum)]` enums. Those are the
1026/// types Renox promises; the trait is implemented through sqlx's own
1027/// traits, so other types sqlx decodes work too, without that promise (see
1028/// docs/stability.md).
1029pub trait FromDb: for<'r> sqlx::Decode<'r, Sqlite> + sqlx::Type<Sqlite> + bounds::Postgres {}
1030
1031impl<T> FromDb for T where
1032    T: for<'r> sqlx::Decode<'r, Sqlite> + sqlx::Type<Sqlite> + bounds::Postgres
1033{
1034}
1035
1036/// A column name (`&str`) or position (`usize`).
1037pub trait RowIndex: sqlx::ColumnIndex<SqliteRow> + bounds::PostgresIndex {}
1038
1039impl<T> RowIndex for T where T: sqlx::ColumnIndex<SqliteRow> + bounds::PostgresIndex {}
1040
1041/// Extra bounds that only apply when the `postgres` feature is on.
1042#[doc(hidden)]
1043pub mod bounds {
1044    #[cfg(feature = "postgres")]
1045    mod on {
1046        use sqlx::postgres::PgRow;
1047
1048        pub trait Postgres:
1049            for<'r> sqlx::Decode<'r, sqlx::Postgres> + sqlx::Type<sqlx::Postgres>
1050        {
1051        }
1052        impl<T> Postgres for T where T: for<'r> sqlx::Decode<'r, sqlx::Postgres> + sqlx::Type<sqlx::Postgres>
1053        {}
1054
1055        pub trait PostgresIndex: sqlx::ColumnIndex<PgRow> {}
1056        impl<T: sqlx::ColumnIndex<PgRow>> PostgresIndex for T {}
1057    }
1058
1059    #[cfg(not(feature = "postgres"))]
1060    mod on {
1061        pub trait Postgres {}
1062        impl<T: ?Sized> Postgres for T {}
1063
1064        pub trait PostgresIndex {}
1065        impl<T: ?Sized> PostgresIndex for T {}
1066    }
1067
1068    pub use on::{Postgres, PostgresIndex};
1069}
1070
1071#[cfg(test)]
1072mod tests {
1073    use super::{numbered_placeholders, seal};
1074    use crate::db::DbValue;
1075
1076    #[test]
1077    fn numbers_placeholders_outside_quotes_and_comments() {
1078        assert_eq!(
1079            numbered_placeholders("SELECT * FROM t WHERE a = ? AND b IN (?, ?)"),
1080            "SELECT * FROM t WHERE a = $1 AND b IN ($2, $3)"
1081        );
1082        assert_eq!(
1083            numbered_placeholders("SELECT '?', \"a?\" -- why?\nFROM t /* ? */ WHERE x = ?"),
1084            "SELECT '?', \"a?\" -- why?\nFROM t /* ? */ WHERE x = $1"
1085        );
1086        assert_eq!(
1087            numbered_placeholders("SELECT 'it''s ?' WHERE y = ?"),
1088            "SELECT 'it''s ?' WHERE y = $1"
1089        );
1090        assert_eq!(numbered_placeholders("SELECT 1"), "SELECT 1");
1091    }
1092
1093    /// PostgreSQL's other string forms keep their `?` (#219).
1094    #[test]
1095    fn skips_dollar_quoted_and_escaped_strings() {
1096        assert_eq!(
1097            numbered_placeholders("SELECT $$why?$$, $fn$ a ? b $fn$ WHERE x = ?"),
1098            "SELECT $$why?$$, $fn$ a ? b $fn$ WHERE x = $1"
1099        );
1100        assert_eq!(
1101            numbered_placeholders(r"SELECT E'it\'s ?', e'\\' WHERE y = ? AND z = ?"),
1102            r"SELECT E'it\'s ?', e'\\' WHERE y = $1 AND z = $2"
1103        );
1104        // A plain string keeps its backslash; `$1` and `a$b` aren't quotes.
1105        assert_eq!(
1106            numbered_placeholders(r"SELECT '\', a$b$ FROM t WHERE c = ?"),
1107            r"SELECT '\', a$b$ FROM t WHERE c = $1"
1108        );
1109        assert_eq!(
1110            numbered_placeholders("SELECT 'é?' WHERE ü = ?"),
1111            "SELECT 'é?' WHERE ü = $1"
1112        );
1113    }
1114
1115    /// Unterminated quotes, comments and dollar quotes swallow the rest; a
1116    /// `$` at the very end or before a digit is no quote (#254).
1117    #[test]
1118    fn placeholders_after_unfinished_quotes_stay() {
1119        assert_eq!(numbered_placeholders("SELECT 'open ?"), "SELECT 'open ?");
1120        assert_eq!(numbered_placeholders("SELECT 1 -- ?"), "SELECT 1 -- ?");
1121        assert_eq!(numbered_placeholders("SELECT 1 /* ?"), "SELECT 1 /* ?");
1122        assert_eq!(numbered_placeholders("SELECT $tag$ ?"), "SELECT $tag$ ?");
1123        assert_eq!(numbered_placeholders("SELECT ? || $"), "SELECT $1 || $");
1124        // A `$1` doesn't open a quote, so the `?` after it is still seen.
1125        assert!(numbered_placeholders("SELECT $1, ?").ends_with(", $1"));
1126        assert_eq!(numbered_placeholders("?"), "$1");
1127    }
1128
1129    /// `retrying` and `transaction_retrying` try again after a busy error,
1130    /// stop at the last attempt, and return other errors at once.
1131    #[tokio::test]
1132    async fn retrying_follows_the_error_kind() {
1133        use std::sync::Arc;
1134        use std::sync::atomic::{AtomicU32, Ordering};
1135
1136        use crate::db::DbError;
1137        use crate::db::error::fake;
1138
1139        fn busy() -> crate::Error {
1140            crate::Error::from(DbError::from(fake::coded(Some("5"))))
1141        }
1142        let db = super::super::connect(&crate::Config::default())
1143            .await
1144            .unwrap();
1145
1146        let tries = Arc::new(AtomicU32::new(0));
1147        let counter = tries.clone();
1148        let value = db
1149            .retrying(3, move || {
1150                let counter = counter.clone();
1151                async move {
1152                    match counter.fetch_add(1, Ordering::SeqCst) {
1153                        0 | 1 => Err(busy()),
1154                        _ => Ok("done"),
1155                    }
1156                }
1157            })
1158            .await
1159            .unwrap();
1160        assert_eq!((value, tries.load(Ordering::SeqCst)), ("done", 3));
1161
1162        let tries = Arc::new(AtomicU32::new(0));
1163        let counter = tries.clone();
1164        let err = db
1165            .retrying(2, move || {
1166                counter.fetch_add(1, Ordering::SeqCst);
1167                async { Err::<(), _>(busy()) }
1168            })
1169            .await
1170            .unwrap_err();
1171        assert!(err.is_retryable());
1172        assert_eq!(
1173            tries.load(Ordering::SeqCst),
1174            2,
1175            "gives up at the last attempt"
1176        );
1177
1178        let tries = Arc::new(AtomicU32::new(0));
1179        let counter = tries.clone();
1180        let err = db
1181            .retrying(5, move || {
1182                counter.fetch_add(1, Ordering::SeqCst);
1183                async { Err::<(), _>(crate::Error::NotFound) }
1184            })
1185            .await
1186            .unwrap_err();
1187        assert!(matches!(err, crate::Error::NotFound));
1188        assert_eq!(tries.load(Ordering::SeqCst), 1, "not retried");
1189
1190        let tries = Arc::new(AtomicU32::new(0));
1191        let counter = tries.clone();
1192        let value = db
1193            .transaction_retrying(3, move |tx| {
1194                let first = counter.fetch_add(1, Ordering::SeqCst) == 0;
1195                Box::pin(async move {
1196                    let one: i64 = super::sql("SELECT CAST(1 AS BIGINT)").scalar(tx).await?;
1197                    if first { Err(busy()) } else { Ok(one) }
1198                })
1199            })
1200            .await
1201            .unwrap();
1202        assert_eq!((value, tries.load(Ordering::SeqCst)), (1, 2));
1203    }
1204
1205    // #254: an Encrypted value on a Db made outside App has no key to seal with.
1206    #[test]
1207    fn sealing_needs_the_apps_key() {
1208        let err = seal(
1209            None,
1210            vec![DbValue::Encrypted(super::super::encrypted::Unsealed(
1211                "x".into(),
1212            ))],
1213        )
1214        .unwrap_err();
1215        assert!(
1216            err.to_string()
1217                .contains("an Encrypted value needs the app's Db"),
1218            "{err}"
1219        );
1220        // Plain values pass through untouched.
1221        assert!(matches!(
1222            seal(None, vec![DbValue::Integer(1)]).unwrap()[..],
1223            [DbValue::Integer(1)]
1224        ));
1225    }
1226}