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

axa_sqlx/
lib.rs
