axa_sqlx/
lib.rs

1//! `PostgreSQL` pooling, migrations, and test resets over sqlx.
2//!
3//! This crate holds the operations every service repeats. It does not wrap
4//! queries: a service writes its own SQL with `sqlx::query!`, so a reviewer
5//! reads exactly what runs.
6//!
7//! It has the same shape as `axa-seaorm`, `axa-diesel`, and `axa-toasty`: a [`Pool`],
8//! a [`PoolConfig`], [`connect`], [`ping`], [`migrate`], [`revert`], [`pending`],
9//! [`truncate`], [`dangerously_recreate`], [`is_unique_violation`], and
10//! [`quote_identifier`], with [`status`] besides.
11//!
12//! ```no_run
13//! use axa_sqlx::{PoolConfig, connect, migrate, sqlx};
14//!
15//! static MIGRATOR: sqlx::migrate::Migrator = sqlx::migrate::Migrator::DEFAULT;
16//!
17//! # async fn run() -> Result<(), Box<dyn std::error::Error>> {
18//! let pool = connect(&PoolConfig::new("postgres://localhost/app")).await?;
19//! migrate(&pool, &MIGRATOR).await?;
20//! # Ok(())
21//! # }
22//! ```
23
24use std::{
25    ops::{Deref, DerefMut},
26    time::Duration,
27};
28
29pub use sqlx;
30use sqlx::{
31    AssertSqlSafe, PgConnection, PgPool, Postgres, Row,
32    migrate::{MigrateError, Migrator},
33    pool::PoolConnection,
34    postgres::PgPoolOptions,
35};
36
37/// The handle a service holds: the pool.
38pub type Pool = PgPool;
39
40/// What this crate's connection and query functions return on failure.
41pub type Error = sqlx::Error;
42
43/// Pool settings.
44#[derive(Clone, PartialEq, Eq)]
45pub struct PoolConfig {
46    /// Connection string.
47    pub uri: String,
48    /// Connections kept open.
49    pub min_connections: u32,
50    /// Largest number of connections.
51    pub max_connections: u32,
52    /// How long to wait for a connection, in milliseconds.
53    pub connect_timeout_ms: u64,
54    /// How long an idle connection stays open, in milliseconds.
55    pub idle_timeout_ms: u64,
56}
57
58impl std::fmt::Debug for PoolConfig {
59    /// Leaves the connection string out, because it can carry a password.
60    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
61        f.debug_struct("PoolConfig")
62            .field("min_connections", &self.min_connections)
63            .field("max_connections", &self.max_connections)
64            .field("connect_timeout_ms", &self.connect_timeout_ms)
65            .field("idle_timeout_ms", &self.idle_timeout_ms)
66            .finish_non_exhaustive()
67    }
68}
69
70impl PoolConfig {
71    /// Settings for `uri` with a small pool.
72    #[must_use]
73    pub fn new(uri: impl Into<String>) -> Self {
74        Self {
75            uri: uri.into(),
76            min_connections: 1,
77            max_connections: 10,
78            connect_timeout_ms: 5_000,
79            idle_timeout_ms: 600_000,
80        }
81    }
82
83    /// The sqlx options these settings describe.
84    #[must_use]
85    pub fn options(&self) -> PgPoolOptions {
86        PgPoolOptions::new()
87            .min_connections(self.min_connections)
88            .max_connections(self.max_connections)
89            .acquire_timeout(Duration::from_millis(self.connect_timeout_ms))
90            .idle_timeout(Duration::from_millis(self.idle_timeout_ms))
91    }
92}
93
94/// Open a pool and prove it works with one connection.
95///
96/// # Errors
97///
98/// Fails when the database cannot be reached.
99pub async fn connect(config: &PoolConfig) -> Result<PgPool, sqlx::Error> {
100    config.options().connect(&config.uri).await
101}
102
103/// Check that the database answers.
104///
105/// # Errors
106///
107/// Fails when it does not.
108pub async fn ping(pool: &PgPool) -> Result<(), sqlx::Error> {
109    sqlx::query("SELECT 1").execute(pool).await?;
110    Ok(())
111}
112
113/// Where a function gets its connection: the pool, or a connection the caller
114/// holds, usually an open transaction, so one function serves both.
115///
116/// `impl sqlx::Acquire` says the same, but its associated types stop the future
117/// being `Send`, which Axum handlers and `tokio::spawn` need.
118///
119/// ```no_run
120/// use axa_sqlx::{Db, sqlx};
121///
122/// async fn count(db: impl Into<Db<'_>>) -> Result<i64, sqlx::Error> {
123///     let mut conn = db.into().acquire().await?;
124///     sqlx::query_scalar("SELECT count(*) FROM items")
125///         .fetch_one(&mut *conn)
126///         .await
127/// }
128/// ```
129pub enum Db<'a> {
130    /// Take a connection from the pool for the length of the call.
131    Pool(&'a PgPool),
132    /// Use the caller's connection.
133    Conn(&'a mut PgConnection),
134}
135
136impl<'a> From<&'a PgPool> for Db<'a> {
137    fn from(pool: &'a PgPool) -> Self {
138        Self::Pool(pool)
139    }
140}
141
142impl<'a> From<&'a mut PgConnection> for Db<'a> {
143    fn from(conn: &'a mut PgConnection) -> Self {
144        Self::Conn(conn)
145    }
146}
147
148impl<'a> Db<'a> {
149    /// The connection to use.
150    ///
151    /// # Errors
152    ///
153    /// Fails when the pool has none to give.
154    pub async fn acquire(self) -> Result<Conn<'a>, sqlx::Error> {
155        match self {
156            Self::Pool(pool) => Ok(Conn::Pooled(pool.acquire().await?)),
157            Self::Conn(conn) => Ok(Conn::Borrowed(conn)),
158        }
159    }
160}
161
162/// A connection for the length of one call. It derefs to [`PgConnection`].
163pub enum Conn<'a> {
164    /// Taken from the pool, and given back when dropped.
165    Pooled(PoolConnection<Postgres>),
166    /// The caller's.
167    Borrowed(&'a mut PgConnection),
168}
169
170impl Deref for Conn<'_> {
171    type Target = PgConnection;
172
173    fn deref(&self) -> &PgConnection {
174        match self {
175            Self::Pooled(conn) => conn,
176            Self::Borrowed(conn) => conn,
177        }
178    }
179}
180
181impl DerefMut for Conn<'_> {
182    fn deref_mut(&mut self) -> &mut PgConnection {
183        match self {
184            Self::Pooled(conn) => conn,
185            Self::Borrowed(conn) => conn,
186        }
187    }
188}
189
190/// Whether an error is a violated unique constraint, which a service usually
191/// reports as `409` rather than `500`.
192#[must_use]
193pub fn is_unique_violation(error: &sqlx::Error) -> bool {
194    error
195        .as_database_error()
196        .is_some_and(sqlx::error::DatabaseError::is_unique_violation)
197}
198
199/// Quote a table or type name for use in SQL.
200///
201/// Identifiers cannot be bound as parameters, so this is the one place names
202/// reach SQL as text. Only lowercase ASCII letters, digits, and underscores
203/// are accepted, starting with a letter or underscore. That is stricter than
204/// Postgres, and it leaves nothing for an attacker to work with.
205///
206/// ```
207/// assert_eq!(axa_sqlx::quote_identifier("users").as_deref(), Some("\"users\""));
208/// assert_eq!(axa_sqlx::quote_identifier("users; DROP TABLE x"), None);
209/// ```
210#[must_use]
211pub fn quote_identifier(name: &str) -> Option<String> {
212    let mut chars = name.chars();
213    let first = chars.next()?;
214    let valid = (first.is_ascii_lowercase() || first == '_')
215        && chars.all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == '_');
216    valid.then(|| format!("\"{name}\""))
217}
218
219/// Quote every name, failing on the first one [`quote_identifier`] refuses.
220fn quote_all<S: AsRef<str>>(names: &[S]) -> Result<Vec<String>, sqlx::Error> {
221    names
222        .iter()
223        .map(|name| {
224            quote_identifier(name.as_ref()).ok_or_else(|| {
225                sqlx::Error::Protocol(format!("invalid identifier: {}", name.as_ref()))
226            })
227        })
228        .collect()
229}
230
231/// Apply pending migrations. Sqlx holds a lock while it does, so services
232/// booting together take turns.
233///
234/// # Errors
235///
236/// Fails when a migration fails, or an applied one no longer matches.
237pub async fn migrate(pool: &PgPool, migrator: &Migrator) -> Result<(), MigrateError> {
238    migrator.run(pool).await
239}
240
241/// The newest applied versions, newest first. Empty on a database that has
242/// never been migrated.
243async fn applied_versions(pool: &PgPool) -> Result<Vec<i64>, sqlx::Error> {
244    let exists: bool = sqlx::query("SELECT to_regclass('_sqlx_migrations') IS NOT NULL AS found")
245        .fetch_one(pool)
246        .await?
247        .try_get("found")?;
248    if !exists {
249        return Ok(Vec::new());
250    }
251    let rows =
252        sqlx::query("SELECT version FROM _sqlx_migrations WHERE success ORDER BY version DESC")
253            .fetch_all(pool)
254            .await?;
255    rows.iter().map(|row| row.try_get("version")).collect()
256}
257
258/// Revert the newest `steps` migrations. Each needs a `.down.sql` file.
259///
260/// # Errors
261///
262/// Fails on a database error or a failed down migration.
263pub async fn revert(pool: &PgPool, migrator: &Migrator, steps: u32) -> Result<(), MigrateError> {
264    let applied = applied_versions(pool).await?;
265    let steps = usize::try_from(steps).unwrap_or(usize::MAX);
266    // `undo` reverts everything newer than its target, so the target is the
267    // version just past the last one to revert, or 0 to revert them all.
268    let target = applied.get(steps).copied().unwrap_or(0);
269    migrator.undo(pool, target).await
270}
271
272/// One migration and whether it has run.
273#[derive(Debug, Clone, PartialEq, Eq)]
274pub struct MigrationStatus {
275    /// Its version.
276    pub version: i64,
277    /// Its description.
278    pub description: String,
279    /// Whether it has been applied.
280    pub applied: bool,
281}
282
283/// Every migration in `migrator`, oldest first, with whether it has run.
284///
285/// # Errors
286///
287/// Fails on a database error.
288pub async fn status(
289    pool: &PgPool,
290    migrator: &Migrator,
291) -> Result<Vec<MigrationStatus>, sqlx::Error> {
292    let applied = applied_versions(pool).await?;
293    Ok(migrator
294        .iter()
295        .filter(|migration| !migration.migration_type.is_down_migration())
296        .map(|migration| MigrationStatus {
297            version: migration.version,
298            description: migration.description.to_string(),
299            applied: applied.contains(&migration.version),
300        })
301        .collect())
302}
303
304/// The names of the migrations in `migrator` not yet applied, in order, each as
305/// `version_description`.
306///
307/// # Errors
308///
309/// Fails on a database error.
310pub async fn pending(pool: &PgPool, migrator: &Migrator) -> Result<Vec<String>, sqlx::Error> {
311    Ok(status(pool, migrator)
312        .await?
313        .into_iter()
314        .filter(|migration| !migration.applied)
315        .map(|migration| format!("{}_{}", migration.version, migration.description))
316        .collect())
317}
318
319/// Empty `tables`, restarting their sequences. Each name goes through
320/// [`quote_identifier`].
321///
322/// # Errors
323///
324/// Fails on an invalid name or a database error.
325pub async fn truncate<S: AsRef<str>>(pool: &PgPool, tables: &[S]) -> Result<(), sqlx::Error> {
326    if tables.is_empty() {
327        return Ok(());
328    }
329    let tables = quote_all(tables)?.join(", ");
330    let sql = format!("TRUNCATE {tables} RESTART IDENTITY CASCADE");
331    sqlx::query(AssertSqlSafe(sql)).execute(pool).await?;
332    Ok(())
333}
334
335/// Drop every table and enum type in `public`, then migrate from nothing.
336/// This gives each test an empty schema, whatever tables it touched.
337///
338/// **This destroys all data.** It exists for tests, and a service must never
339/// reach it outside the `test` environment.
340///
341/// It drops tables and types by name rather than dropping the schema, because
342/// `DROP SCHEMA public CASCADE` would also remove the extensions a database
343/// was prepared with, such as `citext`. Functions are left in place, so a
344/// migration that creates one should use `CREATE OR REPLACE`.
345///
346/// # Errors
347///
348/// Fails on a database error, or a name [`quote_identifier`] refuses.
349pub async fn dangerously_recreate(pool: &PgPool, migrator: &Migrator) -> Result<(), MigrateError> {
350    let tables: Vec<String> =
351        sqlx::query("SELECT tablename::text AS name FROM pg_tables WHERE schemaname = 'public'")
352            .fetch_all(pool)
353            .await?
354            .iter()
355            .map(|row| row.try_get("name"))
356            .collect::<Result<_, _>>()?;
357    // Types owned by an extension (`deptype = 'e'`) belong to that extension.
358    let types: Vec<String> = sqlx::query(
359        "SELECT t.typname::text AS name FROM pg_type t
360         JOIN pg_namespace n ON n.oid = t.typnamespace
361         WHERE t.typtype = 'e' AND n.nspname = 'public'
362           AND NOT EXISTS (
363               SELECT 1 FROM pg_depend d WHERE d.objid = t.oid AND d.deptype = 'e'
364           )",
365    )
366    .fetch_all(pool)
367    .await?
368    .iter()
369    .map(|row| row.try_get("name"))
370    .collect::<Result<_, _>>()?;
371
372    if !tables.is_empty() {
373        let names = quote_all(&tables)?.join(", ");
374        let sql = format!("DROP TABLE {names} CASCADE");
375        sqlx::query(AssertSqlSafe(sql)).execute(pool).await?;
376    }
377    for name in quote_all(&types)? {
378        let sql = format!("DROP TYPE {name} CASCADE");
379        sqlx::query(AssertSqlSafe(sql)).execute(pool).await?;
380    }
381    migrate(pool, migrator).await
382}
383
384#[cfg(test)]
385mod tests {
386    use super::*;
387
388    #[test]
389    fn debug_text_leaves_out_the_connection_string() {
390        let text = format!("{:?}", PoolConfig::new("postgres://me:hunter2@db/app"));
391        assert!(!text.contains("hunter2"), "{text}");
392        assert!(text.contains("max_connections"), "{text}");
393    }
394
395    #[test]
396    fn identifiers_are_quoted() {
397        assert_eq!(quote_identifier("users").as_deref(), Some("\"users\""));
398        assert_eq!(quote_identifier("_a1").as_deref(), Some("\"_a1\""));
399    }
400
401    #[test]
402    fn unsafe_identifiers_are_refused() {
403        for name in [
404            "", "1users", "Users", "a b", "a\"b", "a;b", "a.b", "users--", "é",
405        ] {
406            assert_eq!(quote_identifier(name), None, "{name:?}");
407        }
408    }
409
410    #[test]
411    fn quote_all_fails_on_any_bad_name() {
412        assert!(quote_all(&["good", "also_good"]).is_ok());
413        assert!(quote_all(&["good", "bad name"]).is_err());
414    }
415
416    #[test]
417    fn pool_settings_reach_the_options() {
418        let config = PoolConfig {
419            min_connections: 2,
420            max_connections: 7,
421            connect_timeout_ms: 1_500,
422            idle_timeout_ms: 9_000,
423            ..PoolConfig::new("postgres://x")
424        };
425        let options = config.options();
426        assert_eq!(options.get_min_connections(), 2);
427        assert_eq!(options.get_max_connections(), 7);
428        assert_eq!(options.get_acquire_timeout(), Duration::from_millis(1_500));
429        assert_eq!(options.get_idle_timeout(), Some(Duration::from_secs(9)));
430    }
431
432    #[test]
433    fn errors_that_are_not_from_the_database_are_not_unique_violations() {
434        assert!(!is_unique_violation(&sqlx::Error::RowNotFound));
435    }
436}
axa_sqlx/lib.rs

axa_sqlx/
lib.rs
